pytorch 入门 迁移学习？

知识点1、对图像标准化
知识点2、调用官方预训练模型
知识点3、预测一张图，返回类型

import torch
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import models
from torchvision import transforms as tfs
from torchvision.datasets import ImageFolder
from PIL import Image
import matplotlib.pyplot as plt

知识点1
对数据预处理
对三通道标准化
这里标准化的方式有点奇怪，可以看这里

train_tf = tfs.Compose([
    tfs.RandomResizedCrop(224),
    tfs.RandomHorizontalFlip(),
    tfs.ToTensor(),
    tfs.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

valid_tf = tfs.Compose([
    tfs.Resize(256),
    tfs.CenterCrop(224),
    tfs.ToTensor(),
    tfs.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

知识点2
net = models.resnet34(pretrained=True) # 官方预训练
net.fc = nn.Linear(512, 2)
这里用的是官方与训练的ResNet，这种操作记住

train_set = ImageFolder('./data/hymenoptera_data/train', train_tf)
valid_set = ImageFolder('./data/hymenoptera_data/val', valid_tf)
train_data = DataLoader(train_set, 64, True, num_workers=2)
valid_data = DataLoader(valid_set, 128, False, num_workers=2)

net = models.resnet34(pretrained=True)   # 官方预训练
net.fc = nn.Linear(512, 2)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr=1e-2, weight_decay=1e-4)
from jc_utils import train
train(net, train_data, valid_data, 20, optimizer, criterion)

知识点3
预测一张图，并返回预测类型，这一系列的操作也许有到的时候

net = net.eval()
im1 = Image.open('./data/hymenoptera_data/train/ants/0013035.jpg')
plt.imshow(im1)
plt.show()
im = valid_tf(im1)
out = net(Variable(im.unsqueeze(0)).cuda())
pred_label = out.max(1)[1].data[0]
print('predict label: {}'.format(train_set.classes[pred_label]))

未与训练的对比

net = models.resnet34()
net.fc = nn.Linear(512, 2)
optimizer = torch.optim.SGD(net.parameters(), lr=1e-2, weight_decay=1e-4)
train(net, train_data, valid_data, 20, optimizer, criterion)