from __future__ import print_function, division import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms from torch.autograd import Variable import numpy as np from torchvision import models import matplotlib.pyplot as plt batch_size = 4 learning_rate = 0.0002 epoch = 10 train_transforms = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((.5, .5, .5), (.5, .5, .5)) ]) val_transforms = transforms.Compose([ transforms.Resize(256), transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize((.5, .5, .5), (.5, .5, .5)) ]) train_dir = './data/train' train_datasets = datasets.ImageFolder(train_dir, transform=train_transforms) train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=batch_size, shuffle=True) val_dir = './data/val' val_datasets = datasets.ImageFolder(val_dir, transform=val_transforms) val_dataloader = torch.utils.data.DataLoader(val_datasets, batch_size=batch_size, shuffle=True) class VGGNet(nn.Module): def __init__(self, num_classes=2): super(VGGNet, self).__init__() #自己下载预训练模型 net = models.vgg16(pretrained=False) net.load_state_dict(torch.load('./vgg16-397923af.pth')) net.classifier = nn.Sequential() self.features = net self.classifier = nn.Sequential( nn.Linear(512 * 7 * 7, 512), nn.ReLU(True), nn.Dropout(), nn.Linear(512, 128), nn.ReLU(True), nn.Dropout(), nn.Linear(128, num_classes), ) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) x = self.classifier(x) return x if __name__=='__main__': # --------------------训练过程--------------------------------- model = VGGNet() if torch.cuda.is_available(): model.cuda() params = [{'params': md.parameters()} for md in model.children() if md in [model.classifier]] optimizer = optim.Adam(model.parameters(), lr=learning_rate) loss_func = nn.CrossEntropyLoss() Loss_list = [] Accuracy_list = [] for epoch in range(5): print('epoch {}'.format(epoch + 1)) # training----------------------------- train_loss = 0. train_acc = 0. batchIndex = 0 for batch_x, batch_y in train_dataloader: batchIndex += 1 batch_x, batch_y = Variable(batch_x).cuda(), Variable(batch_y).cuda() out = model(batch_x) loss = loss_func(out, batch_y) train_loss += loss.item() pred = torch.max(out, 1)[1] #torch.max(out, 1) 0-每一列最大的 1-每一行最大的 生成两个数组 一个数值 一个序号 train_correct = (pred == batch_y).sum() train_acc += train_correct.item() optimizer.zero_grad() loss.backward() optimizer.step() print('Train Loss: {:.6f}, Acc: {:.6f}'.format(train_loss / (len( train_datasets)), train_acc / (len(train_datasets)))) # evaluation-------------------------------- model.eval() eval_loss = 0. eval_acc = 0. for batch_x, batch_y in val_dataloader: batch_x, batch_y = Variable(batch_x, volatile=True).cuda(), Variable(batch_y, volatile=True).cuda() out = model(batch_x) loss = loss_func(out, batch_y) eval_loss += loss.item() pred = torch.max(out, 1)[1] num_correct = (pred == batch_y).sum() eval_acc += num_correct.item() print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len( val_datasets)), eval_acc / (len(val_datasets)))) Loss_list.append(eval_loss / (len(val_datasets))) Accuracy_list.append(100 * eval_acc / (len(val_datasets))) x1 = range(0, 5) x2 = range(0, 5) y1 = Accuracy_list y2 = Loss_list plt.subplot(2, 1, 1) plt.plot(x1, y1, 'o-') plt.title('Test accuracy vs. epoches') plt.ylabel('Test accuracy') plt.subplot(2, 1, 2) plt.plot(x2, y2, '.-') plt.xlabel('Test loss vs. epoches') plt.ylabel('Test loss') plt.show() # plt.savefig("accuracy_loss.jpg")