Pytorch预训练模型

Pytorch预训练模型

Pytorch支持预训练的模型如下

From torchvision package:

• ResNet (resnet18, resnet34, resnet50, resnet101, resnet152)
• DenseNet (densenet121, densenet169, densenet201, densenet161)
• Inception v3 (inception_v3)
• VGG (vgg11, vgg11_bn, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn)
• SqueezeNet (squeezenet1_0, squeezenet1_1)
• AlexNet (alexnet)

From Pretrained models for PyTorch package:

• ResNeXt (resnext101_32x4d, resnext101_64x4d)
• NASNet-A Large (nasnetalarge)
• NASNet-A Mobile (nasnetamobile)
• Inception-ResNet v2 (inceptionresnetv2)
• Dual Path Networks (dpn68, dpn68b, dpn92, dpn98, dpn131, dpn107)
• Inception v4 (inception_v4)
• Xception (xception)
• Squeeze-and-Excitation Networks (senet154, se_resnet50, se_resnet101, se_resnet152, se_resnext50_32x4d, se_resnext101_32x4d)
• PNASNet-5-Large (pnasnet5large)
• PolyNet (polynet)

Pytorch预训练运行环境

• Python 3.5+
• PyTorch 0.3+

安装方法

pip install cnn_finetune

用例

from cnn_finetune import make_model
from torch import optim
model = make_model(model_name='resnext101_64x4d', num_classes=10, pretrained=True)
# optimizer
optimizer = optim.SGD(model.parameters(), lr=0.0001)

导入包，将make_model中的num_classes改成你项目中分类类别数。如果你选择的model是VGG系列或者其他需要设置固定输入格式大小的模型还会有input_size需要设置。需要将输入通过cv.resize完成或者特征提取成相应大小。

简单的完整示例

"""
!/usr/bin/env python
-*- coding:utf-8 -*-
Author: eric.lai
Created on 2019/7/22 9:48
"""
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
from cnn_finetune import make_model
from torch import optim
import torch.nn as nn
import torch

mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
model = make_model(model_name='resnext101_64x4d', num_classes=10, pretrained=True)
# optimizer
optimizer = optim.SGD(model.parameters(), lr=0.0001)

def load_data():
    model.train()
    criterion = nn.MSELoss() # CrossEntropyLoss does not expect a one-hot encoded vector as the target
    total_loss = 0
    total_size = 0

    for i in range(100):
        optimizer.zero_grad()
        batch_xs, batch_ys = mnist.train.next_batch(10)
        batch_xs = list(batch_xs)*3
        batch_xs = np.array(batch_xs)
        batch_xs = np.reshape(batch_xs, [-1, 3, 28, 28])
        batch_xs = torch.from_numpy(batch_xs)
        batch_ys = torch.from_numpy(batch_ys)
        output = model(batch_xs)
        loss = criterion(output,target=batch_ys.float())
        total_loss += loss.item()
        total_size += batch_xs.shape[0]
        loss.backward()
        optimizer.step()
        print(loss)
        test()

def test():
    model.eval()
    with torch.no_grad():
        batch_xs, batch_ys = mnist.test.next_batch(100)
        batch_xs = list(batch_xs) * 3
        batch_xs = np.array(batch_xs)
        batch_xs = np.reshape(batch_xs, [-1, 3, 28, 28])
        batch_xs = torch.from_numpy(batch_xs)
        output = model(batch_xs)
        correct = list(np.equal(np.argmax(output, axis=1),np.argmax(batch_ys, axis=1)))
        nums = correct.count(1)
    print('test accuracy: ', nums / len(batch_xs))
    return output

if __name__ == '__main__':
    load_data()

需要注意的几点：

1.Pytorch是channel_first，在reshape的时候要把通道数放在第一个位置

2.输入的数据要通过torch_from_numpy转换成Tensor

参考来源

https://github.com/Laicheng0830/pytorch-cnn-finetune#from-torchvision-package