pytorch VGG11识别cifar10数据集(训练+预测单张输入图片操作)

def vgg_stack(num_convs, channels): # vgg_net = vgg_stack((1, 1, 2, 2, 2), ((3, 64), (64, 128), (128, 256), (256, 512), (512, 512)))


 net = []
 for n, c in zip(num_convs, channels):
  in_c = c[0]
  out_c = c[1]
  net.append(vgg_block(n, in_c, out_c))
 return nn.Sequential(*net)

#vgg类
class vgg(nn.Module):
 def __init__(self):
  super(vgg, self).__init__()
  self.feature = vgg_net
  self.fc = nn.Sequential(
   nn.Linear(512, 100),
   nn.ReLU(True),
   nn.Linear(100, 10)
  )
 
 def forward(self, x):
  x = self.feature(x)
  x = x.view(x.shape[0], -1)
  x = self.fc(x)
  return x

def vgg_block(num_convs, in_channels, out_channels):
 net = [nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.ReLU(True)]
 
 for i in range(num_convs - 1): # 定义后面的许多层
  net.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))
  net.append(nn.ReLU(True))
 
 net.append(nn.MaxPool2d(2, 2)) # 定义池化层
 return nn.Sequential(*net)
 
# 下面我们定义一个函数对这个 vgg block 进行堆叠
def vgg_stack(num_convs, channels): # vgg_net = vgg_stack((1, 1, 2, 2, 2), ((3, 64), (64, 128), (128, 256), (256, 512), (512, 512)))
 net = []
 for n, c in zip(num_convs, channels):
  in_c = c[0]
  out_c = c[1]
  net.append(vgg_block(n, in_c, out_c))
 return nn.Sequential(*net)
 
#确定vgg的类型，是vgg11 还是vgg16还是vgg19
vgg_net = vgg_stack((1, 1, 2, 2, 2), ((3, 64), (64, 128), (128, 256), (256, 512), (512, 512)))
#vgg类
class vgg(nn.Module):
 def __init__(self):
  super(vgg, self).__init__()
  self.feature = vgg_net
  self.fc = nn.Sequential(
   nn.Linear(512, 100),
   nn.ReLU(True),
   nn.Linear(100, 10)
  )
 def forward(self, x):
  x = self.feature(x)
  x = x.view(x.shape[0], -1)
  x = self.fc(x)
  return x
 
#获取vgg网络
net = vgg() 

import sys
import numpy as np
import torch
from torch import nn
from torch.autograd import Variable
from torchvision.datasets import CIFAR10
import torchvision.transforms as transforms
 
def vgg_block(num_convs, in_channels, out_channels):
 net = [nn.Conv<mark style="color:transparent">本文来源gaodaimacom搞#^代%!码&网*</mark>2d(in_channels, out_channels, kernel_size=3, padding=1), nn.ReLU(True)]
 
 for i in range(num_convs - 1): # 定义后面的许多层
  net.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))
  net.append(nn.ReLU(True))
 
 net.append(nn.MaxPool2d(2, 2)) # 定义池化层
 return nn.Sequential(*net)
 
# 下面我们定义一个函数对这个 vgg block 进行堆叠
def vgg_stack(num_convs, channels): # vgg_net = vgg_stack((1, 1, 2, 2, 2), ((3, 64), (64, 128), (128, 256), (256, 512), (512, 512)))
 net = []
 for n, c in zip(num_convs, channels):
  in_c = c[0]
  out_c = c[1]
  net.append(vgg_block(n, in_c, out_c))
 return nn.Sequential(*net)
 
#vgg类
class vgg(nn.Module):
 def __init__(self):
  super(vgg, self).__init__()
  self.feature = vgg_net
  self.fc = nn.Sequential(
   nn.Linear(512, 100),
   nn.ReLU(True),
   nn.Linear(100, 10)
  )
 def forward(self, x):
  x = self.feature(x)
  x = x.view(x.shape[0], -1)
  x = self.fc(x)
  return x
 
# 然后我们可以训练我们的模型看看在 cifar10 上的效果
def data_tf(x):
 x = np.array(x, dtype='float32') / 255
 x = (x - 0.5) / 0.5
 x = x.transpose((2, 0, 1)) ## 将 channel 放到第一维，只是 pytorch 要求的输入方式
 x = torch.from_numpy(x)
 return x
 
transform = transforms.Compose([transforms.ToTensor(),
         transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
         ])
def get_acc(output, label):
 total = output.shape[0]
 _, pred_label = output.max(1)
 num_correct = (pred_label == label).sum().item()
 return num_correct / total
 
def train(net, train_data, valid_data, num_epochs, optimizer, criterion):
 if torch.cuda.is_available():
  net = net.cuda()
 for epoch in range(num_epochs):
  train_loss = 0
  train_acc = 0
  net = net.train()
  for im, label in train_data:
   if torch.cuda.is_available():
    im = Variable(im.cuda())
    label = Variable(label.cuda())
   else:
    im = Variable(im)
    label = Variable(label)
   # forward
   output = net(im)
   loss = criterion(output, label)
   # forward
   optimizer.zero_grad()
   loss.backward()
   optimizer.step()
 
   train_loss += loss.item()
   train_acc += get_acc(output, label)
 
  if valid_data is not None:
   valid_loss = 0
   valid_acc = 0
   net = net.eval()
   for im, label in valid_data:
    if torch.cuda.is_available():
     with torch.no_grad():
      im = Variable(im.cuda())
      label = Variable(label.cuda())
    else:
     with torch.no_grad():
      im = Variable(im)
      label = Variable(label)
    output = net(im)
    loss = criterion(output, label)
    valid_loss += loss.item()
    valid_acc += get_acc(output, label)
   epoch_str = (
     "Epoch %d. Train Loss: %f, Train Acc: %f, Valid Loss: %f, Valid Acc: %f, "
     % (epoch, train_loss / len(train_data),
      train_acc / len(train_data), valid_loss / len(valid_data),
      valid_acc / len(valid_data)))
  else:
   epoch_str = ("Epoch %d. Train Loss: %f, Train Acc: %f, " %
       (epoch, train_loss / len(train_data),
       train_acc / len(train_data)))
 
  # prev_time = cur_time
  print(epoch_str)
 
if __name__ == '__main__':
 # 作为实例，我们定义一个稍微简单一点的 vgg11 结构，其中有 8 个卷积层
 vgg_net = vgg_stack((1, 1, 2, 2, 2), ((3, 64), (64, 128), (128, 256), (256, 512), (512, 512)))
 print(vgg_net)
 
 train_set = CIFAR10('./data', train=True, transform=transform, download=True)
 train_data = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)
 test_set = CIFAR10('./data', train=False, transform=transform, download=True)
 test_data = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False)
 
 net = vgg()
 optimizer = torch.optim.SGD(net.parameters(), lr=1e-1)
 criterion = nn.CrossEntropyLoss() #损失函数为交叉熵
 
 train(net, train_data, test_data, 50, optimizer, criterion)
 torch.save(net, 'vgg_model.pth')