Pytorch图像分割迁移学习

Question

我是Pytorch的新手。我的问题是：如何将迁移学习应用于自定义数据集？我正在做脑肿瘤的图像分割。我可以找到使用U-net结构的例子，但我找不到使用预训练模型的权重进行U-net图像分割的例子？

Answer 1

您可以通过两种方式获得预先训练的模型：

1. 以.pt或.pth等格式共享的模型权重或完整模型：在这种情况下，保存和加载模型是一个很好的起点。从那里的教程中复制，可以将模型加载为

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

2. 另一种方法是从torchvision加载模型。Torchvision models提供了可用型号列表。U-Net尚不可用。然而，可以加载一个预训练的模型作为编码器，并编写一个单独的解码器，以使用预训练的编码器形成U-Net。在这种情况下，从API中显示的函数调用返回的模型对象在pretrained=True时已经加载了预训练的权重

对于编写自定义数据加载器，PyTorch数据加载器可能是一个有用的指南。