图片二分类在VGG16基础上finetune

# -*- coding: utf-8 -*- """ Created on ${DATE} ${TIME} @author: """ import torch import torchvision #prepare model1_vgg16 = torchvision.models.vgg16(pretrained=True) #设置网络参数 #读取输入特征的维度 num_fc = model1_vgg16.classifier[6].in_features #修改最后一层的输出维度，即分类数 model1_vgg16.classifier[6] = torch.nn.Linear(num_fc,2) #固定权值参数 #对于模型的每个权重，使其不进行反向传播，即固定参数 for param in model1_vgg16.parameters(): param.requires_grad = False #但是参数全部固定了，也没法进行学习，所以我们不固定最后一层，即全连接层 for param in model1_vgg16.classifier[6].parameters(): param.requires_grad = True #训练模型 import numpy as np from torchvision import transforms,utils from torch.utils.data import DataLoader from torch.autograd import Variable #加载图片 这种读取图片的方式用的是Torch自带的ImageFloder,读取的文件夹必须在一个大的子文件下，按类别归好类 #train data train_data = torchvision.datasets.ImageFolder("C:/datafolder/hymenoptera_data/train",transform=transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor() ])) train_loader = DataLoader(train_data,batch_size=20,shuffle=True) #test data test_data = torchvision.datasets.ImageFolder("C:/datafolder/hymenoptera_data/val",transform=transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor() ])) test_loader = DataLoader(test_data,batch_size=20,shuffle=True) #训练 定义自己的优化器 criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model1_vgg16.parameters(),lr=0.001) #训练模型 EPOCH = 20 for epoch in range(EPOCH): train_loss = 0. train_acc = 0. for step, data in enumerate(train_loader): batch_x, batch_y = data batch_x, batch_y = Variable(batch_x), Variable(batch_y) # batch_y not one hot # out is the probability of eatch class # such as one sample[-1.1009 0.1411 0.0320],need to calculate the max index # out shape is batch_size * class out = model1_vgg16(batch_x) loss = criterion(out, batch_y) train_loss += loss.item() # pred is the expect class # batch_y is the true label pred = torch.max(out, 1)[1] train_correct = (pred == batch_y).sum() train_acc += train_correct.item() optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 == 0: print('Epoch: ', epoch, 'Step', step, 'Train_loss: ', train_loss / ((step + 1) * 20), 'Train acc: ', train_acc / ((step + 1) * 20)) print('Epoch: ', epoch, 'Train_loss: ', train_loss / len(train_data), 'Train acc: ', train_acc / len(train_data)) #测试 model1_vgg16.eval() eval_loss = 0 eval_acc = 0 for step,data in enumerate(test_loader): batch_x,batch_y = data batch_x,batch_y = Variable(batch_x),Variable(batch_y) out = model1_vgg16(batch_x) loss = criterion(out,batch_y) eval_loss+=loss.item() #pred is the expect class #batch_y is the true label pred = torch.max(out,1)[1] test_correct = (pred==batch_y).sum() eval_acc += test_correct.item() print('Test_loss: ', eval_loss / len(test_data), 'Test acc: ', eval_acc / len(test_data)) 

参考：

https://blog.csdn.net/zhonglongshen/article/details/