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,缩小一半终于把程序跑出来了。