一、现有模型的加载、修改、添加

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import torchvision
from torch import nn

vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True)
print(vgg16_true)

train_data = torchvision.datasets.CIFAR10("dataset", train=True, transform=torchvision.transforms.ToTensor(),download=True)

vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10)) # 添加 层
print(vgg16_true)

print(vgg16_false)
vgg16_false.classifier[6] = nn.Linear(4096, 10) # 修改层的内容
print(vgg16_false)

二、网络模型的保存和读取

保存

1
2
3
4
5
6
7
8
9
import torch
import torchvision

vgg16 = torchvision.models.vgg16()
# 保存方式1 模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth")

# 保存方式2 模型参数 （将模型中的参数保存为字典）[官方推荐，存储量小]
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

读取

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import torch
import torchvision
from torch import nn

# 方式1 对应保存方式1
model = torch.load("vgg16_method1.pth")
print(model)

# 方式2
vgg16 = torchvision.models.vgg16()
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16)


class Test(nn.Module):
    def __init__(self):
        super(Test, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)

    def forward(self, x):
        x = self.conv1(x)
        return x


# 陷阱1 需要引入自己的模型，要么在开头from model_save import *,要么在这个文件中把模型复制过来
model = torch.load('test_method1.pth')
print(model)

三、完整模型的训练

两种类型分类的模型

模型预测得到的outputs=[0.3,0.5] (第一张图)

​ [0.5,0.2] (第二张图)

使用argmax（1）方法可以得到该行最大数的位置，即对应的那个类型 argmax（0）是得到列最大数的位置

这里argmax(1)=[2][1]

将模型预测的结果和真实情况进行比较：

1
2
3
preds =  torch.argmax(1)

print((preds == targets).sum)  # 得到正确率

yourmodel.train()、yourmodel.eval()对某些特定的层有影响，比如dropout层，详见pytorch官网

四、使用gpu训练

法一：

改动的地方：

  1. 网络模型
  2. 数据（输入、检测）
  3. 损失函数

改动方法:data = data.cuda()

法二：

先在开头定义：device = torch.device("cuda")

​ 语法糖写法：device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

接着像法一一样修改：

1
2
test = Test()
test = test.to(device)

五、查看开源项目

修改参数时有些参数有required = True,可以把这个参数删掉，然后改成default='data的地址'

TIPS

ctrl + D 可以复制本行内容