用VGG19模型跑CIFAR10的过程与总结

CIFAR10的源代码

遇到的问题

结果

1.源代码

废话不多说，直接上源代码。源代码是网上一位大神写的，原文源代码链接。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.datasets as dset
import torchvision.models as models
import torch.optim as optim
import torchvision.transforms as transforms

image_size = (224,224)
transform=transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
train_data=dset.CIFAR10(root='./data',train=True,transform=transform,download=False)
test_data=dset.CIFAR10(root='./data',train=False,transform=transform,download=False)
train_loader=torch.utils.data.DataLoader(train_data,batch_size=20,shuffle=True)
test_loader=torch.utils.data.DataLoader(test_data,batch_size=20,shuffle=True)
train_len=len(train_data)
test_len=len(test_data)
print(train_len,test_len)

# model=models.resnet18(pretrained=True) vgg19
model=models.vgg19(pretrained=True) #这里采用VGG19模型
n_classes=len(train_data.classes)
print('labels:',n_classes)
model.fc=nn.Linear(512, n_classes)

def get_num_correct(out, labels):  #求准确率
    return out.argmax(dim=1).eq(labels).sum().item()

model=model.cuda()
optimizer=torch.optim.SGD(model.parameters(),lr=0.01)
epoch_num = 10
for epoch in range(epoch_num):
    total_loss=0
    train_correct=0
    test_correct=0
    for batch in train_loader:
        images,labels=batch
        outs=model(images.cuda())
        loss=F.cross_entropy(outs,labels.cuda())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
        train_correct+=get_num_correct(outs,labels.cuda())
    for batch in test_loader:
        images,labels=batch
        outs=model(images.cuda())
        test_correct+=get_num_correct(outs,labels.cuda())
    print('process:',epoch," loss:",total_loss," train_correct:",train_correct/train_len, " test_correct:",test_correct/test_len)

# torch.save(model, './data/resent.pkl')
# print('The model has been saved!!')

2.遇到的问题

第一次下载完VGG19模型，开始跑程序，遇到的问题如下：
图片有些不清晰，错误提示如下：

RuntimeError: CUDA out of memory. Tried to allocate 32.00 MiB (GPU 0;
6.00 GiB total capacity; 3.88 GiB already allocated; 12.63 MiB free; 3.99 GiB reserved in total by PyTorch)

百度之后建议缩小batch_size的值试试。把batch_size=20,缩小一半终于把程序跑出来了。

3.结果