1

pip install torchvision

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

28

from torchvision import models

from torch import nn

from torch import optim

resnet_model = models.resnet18(pretrained=True)

# pretrained 设置为 True，会自动下载模型 所对应权重，并加载到模型中

# 也可以自己下载 权重，然后 load 到 模型中，源码中有 权重的地址。

# 假设 我们的 分类任务只需要 分 100 类，那么我们应该做的是

# 1. 查看 resnet 的源码

# 2. 看最后一层的 名字是啥 （在 resnet 里是 self.fc = nn.Linear(512 * block.expansion, num_classes)）

# 3. 在外面替换掉这个层

resnet_model.fc= nn.Linear(in_features=..., out_features=100)

# 这样就 哦了，修改后的模型除了输出层的参数是 随机初始化的，其他层都是用预训练的参数初始化的。

# 如果只想训练 最后一层的话，应该做的是：

# 1. 将其它层的参数 requires_grad 设置为 False

# 2. 构建一个 optimizer， optimizer 管理的参数只有最后一层的参数

# 3. 然后 backward， step 就可以了

# 这一步可以节省大量的时间，因为多数的参数不需要计算梯度

for para in list(resnet_model.parameters())[:-2]:

para.requires_grad=False

optimizer = optim.SGD(params=[resnet_model.fc.weight, resnet_model.fc.bias], lr=1e-3)

...

1

2

3

4

5

def __setattr__(self, name, value):

def remove_from(*dicts):

for d in dicts:

if name in d:

del d[name]

1

2

3

4

5

6

7

8

9

def weight_init(m):

# 使用isinstance来判断m属于什么类型

if isinstance(m, nn.Conv2d):

n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels

m.weight.data.normal_(0, math.sqrt(2. / n))

elif isinstance(m, nn.BatchNorm2d):

# m中的weight，bias其实都是Variable，为了能学习参数以及后向传播

m.weight.data.fill_(1)

m.bias.data.zero_()

1

2

3

4

5

6

7

8

9

model = torchvision.models.resnet18(pretrained=True)

for param in model.parameters():

param.requires_grad = False

# 替换最后的全连接层， 改为训练100类

# 新构造的模块的参数默认requires_grad为True

model.fc = nn.Linear(512, 100)

# 只优化最后的分类层

optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)

1

2

3

4

5

6

7

8

ignored_params = list(map(id, model.fc.parameters()))

base_params = filter(lambda p: id(p) not in ignored_params,

model.parameters())

optimizer = torch.optim.SGD([

{'params': base_params},

{'params': model.fc.parameters(), 'lr': 1e-3}

], lr=1e-2, momentum=0.9)

1

torch.save(the_model.state_dict(), PATH)

1

2

the_model = TheModelClass(*args, **kwargs)

the_model.load_state_dict(torch.load(PATH))

1

torch.save(the_model, PATH)

1

the_model = torch.load(PATH)

1

2

3

model_dict = the_model.state_dict(),

for k,v in model_dict.items():

print(k)

1

model_dict[k] = new_value

1

the_model.load_state_dict(model_dict)

1

2

3

4

5

6

7

8

9

10

11

model_dict = model.state_dict()

pretrained_dict = torch.load(model_path)

# 1. filter out unnecessary keys

diff = {k: v for k, v in model_dict.items() if \

k in pretrained_dict and pretrained_dict[k].size() == v.size()}

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].size() == v.size()}

pretrained_dict.update(diff)

# 2. overwrite entries in the existing state dict

model_dict.update(pretrained_dict)

# 3. load the new state dict

model.load_state_dict(model_dict)

1

2

3

4

5

6

7

8

9

10

11

12

model_dict = model.state_dict()

pretrained_dict = torch.load(model_path)

keys = []

for k,v in pretrained_dict.items():

keys.append(k)

i = 0

for k,v in model_dict.items():

if v.size() == pretrained_dict[keys[i]].size():

print(k, ',', keys[i])

model_dict[k]=pretrained_dict[keys[i]]

i = i + 1

model.load_state_dict(model_dict)