Pytorch 实现Resnet分类Cifar10数据集＞95%

主要参数设置

1.图片预处理 – 关键1

import torchvision.transforms as transforms
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
])

2.优化函数和学习率 – 关键2

import torch
optimizer = torch.optim.SGD(Resnet50.parameters(), lr=0.001,momentum=0.9)

完整代码

1. 加载数据集和预处理

# 导入数据集
from torchvision import datasets
import torch
import torchvision.transforms as transforms
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
])

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'trunk')
cifar_train = datasets.CIFAR10(root="/data",train=True, download=True, transform=transform)
cifar_test = datasets.CIFAR10(root="/data",train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(cifar_train, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(cifar_test, batch_size=32, shuffle=False)

2.导入resnet50预训练模型（修改最后一层输出)

import torchvision
Resnet50 = torchvision.models.resnet50(pretrained=True)
Resnet50.fc.out_features=10
print(Resnet50)

3.训练

# 定义损失函数和优化器
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(Resnet50.parameters(), lr=0.001)
epoch = 10
Resnet50 = Resnet50.to(device)
SE_Resnet50 = SE_Resnet50.to(device)
total_step = len(train_loader)
train_all_loss = []
val_all_loss = []
for i in range(epoch):
    Resnet50.train()
    train_total_loss = 0
    train_total_num = 0
    train_total_correct = 0

    for iter, (images,labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        
        outputs = Resnet50(images)
        loss = criterion(outputs,labels)
        train_total_correct += (outputs.argmax(1) == labels).sum().item()
            
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_total_num += labels.shape[0]
        train_total_loss += loss.item()
        print("Epoch [{}/{}], Iter [{}/{}], train_loss:{:4f}".format(i+1,epoch,iter+1,total_step,loss.item()/labels.shape[0]))
    Resnet50.eval()
    test_total_loss = 0
    test_total_correct = 0
    test_total_num = 0
    for iter,(images,labels) in enumerate(test_loader):
        images = images.to(device)
        labels = labels.to(device)
        
        outputs = Resnet50(images)
        loss = criterion(outputs,labels)
        test_total_correct += (outputs.argmax(1) == labels).sum().item()
        test_total_loss += loss.item()
        test_total_num += labels.shape[0]
    print("Epoch [{}/{}], train_loss:{:.4f}, train_acc:{:.4f}%, test_loss:{:.4f}, test_acc:{:.4f}%".format(
        i+1, epoch, train_total_loss / train_total_num, train_total_correct / train_total_num * 100, test_total_loss / test_total_num, test_total_correct / test_total_num * 100
    
    ))
    train_all_loss.append(np.round(train_total_loss / train_total_num,4))
    test_all_loss.append(np.round(test_total_loss / test_total_num,4))

4.可视化训练过程的损失

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

plt.figure()
plt.title("Train Loss and Test Loss Curve")
plt.xlabel('plot_epoch')
plt.ylabel('loss')
plt.plot(train_all_loss)
plt.plot(test_all_loss)
plt.legend(['train loss', 'test loss'])

5.训练结果和可视化结果

这个训练结果也表明，优化算法的选择对模型的表现也很重要。如果懒人直接选择Adam,后期会因为学习率太低导致性能没有提升，最终test_acc收敛在80%多。这里使用SGD的效果反而很nice。