pytorch-使用预训练模型,分类器训练

时间:2020-05-05 18:05:23

标签: python deep-learning pytorch

我想使用预先训练的模型并仅训练分类器部分,而使网络主要部分的权重保持不变。

我已经导入了模型,更改了输出层以匹配数据中的类数,并实例化了优化器:

model = torch.hub.load('pytorch/vision:v0.6.0', 'mobilenet_v2', pretrained=True, force_reload=True)

model.classifier[1] = nn.Linear(1280, 50, True)

optimizer = torch.optim.Adam(model.classifier.parameters(), lr=0.05)

然后在培训部分中,我照常做:

optimizer.zero_grad()
loss.backward()
optimizer.step()

以这种方式,我是只训练网络中的分类器部分还是还是以某种方式调整了整个网络的权重,而我却缺少了一些东西?

0 个答案:

没有答案