在 Pytorch 中使用预训练模型进行语义分割,然后使用我们自己的数据集仅训练完全连接的层



我正在学习Pytorch,并试图了解该库如何用于语义分割。 到目前为止,我所理解的是,我们可以在pytorch中使用预先训练的模型。我找到了一篇在 .eval(( 模式下使用此模型的文章,但我找不到任何关于使用这种模型在我们自己的数据集上进行训练的教程。我有一个非常小的数据集,我需要迁移学习才能获得结果。我的目标是只用我自己的数据训练 FC 层。如何在 Pytorch 中实现这一点,而不会使 OOP 或如此多.py文件的代码复杂化。我一直很难在 github 中找出这样的存储库,因为在 OOP 方面我不是最熟练的人。直到最近,我一直在使用 Keras 进行深度学习,那里的一切都简单明了。我在 Pycharm 中也有相同的选项吗? 我感谢对此的任何指导。我需要运行一段代码来执行语义分割,我对需要采取的许多步骤感到非常困惑。

假设您从一个名为model的预训练模型开始。所有这些都发生在向模型传递任何数据之前。

您希望通过查看所有图层,然后使用model.children()对它们进行索引来查找要训练的图层。运行此命令将显示所有块和图层。

list(model.children())

假设您现在已经找到了要微调的层(如您所描述的 FC 层(。如果要训练的层是最后 5 个层,则可以抓取除最后 5 个层之外的所有层,以便将它们的requires_grad参数设置为False这样它们就不会在运行训练算法时进行训练。

list(model.children())[-5:]

删除这些图层:

layer_list = list(model.children())[-5:]

使用顺序重建模型:

model_small = nn.Sequential(*list(model.children())[:-5])

requires_grad参数设置为False

for param in model_small.parameters():
param.requires_grad = False

现在,您有一个名为model_small的模型,该模型包含除要训练的层之外的所有层。现在,您可以重新附加已删除的图层,它们本质上会将requires_grad参数设置为True。现在,当您训练模型时,它只会更新这些层上的权重。

model_small.avgpool_1 = nn.AdaptiveAvgPool2d()
model_small.lin1 = nn.Linear()
model_small.logits = nn.Linear()
model_small.softmax = nn.Softmax()
model = model_small.to(device)

最新更新