Pytorch图像分割转移学习

Question

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

Answer 1

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

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

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

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

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