如何使用预训练的编码器来自定义Unet

huangapple go评论73阅读模式
英文:

How to use pretrained encoder for customized Unet

问题

如果您有一个标准的Unet编码器,例如resnet50,那么很容易将其添加到Unet模型中。例如:

  1. ENCODER = 'resnet50'
  2. ENCODER_WEIGHTS = 'imagenet'
  3. CLASSES = class_names
  4. ACTIVATION = 'sigmoid' # 可以是None(用于logits)或'softmax2d'(用于多类别分割)
  5. # 使用预训练的编码器创建分割模型
  6. model = smp.Unet(
  7. encoder_name=ENCODER,
  8. encoder_weights=ENCODER_WEIGHTS,
  9. classes=len(CLASSES),
  10. activation=ACTIVATION,
  11. )
  12. preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)

然而,如果您有一个自定义的Unet编码器(不一定使用resnet50),您可以像下面这样做:

  1. class VGGBlock(nn.Module):
  2. # 这里是VGGBlock的定义
  3. class UNet(nn.Module):
  4. # 这里是UNet的定义

对于自定义编码器,通常不会直接进行ImageNet预训练,因为自定义编码器的结构与ImageNet预训练模型的结构不匹配。预训练通常适用于具有相似结构的模型。如果要利用现有的预训练编码器(如resnet50),您可以将其加载到模型中,然后进行微调以适应您的任务。可以使用以下代码加载预训练的resnet50编码器:

  1. import torchvision.models as models
  2. # 加载预训练的resnet50模型
  3. pretrained_resnet50 = models.resnet50(pretrained=True)
  4. # 将预训练模型的权重加载到自定义UNet模型的编码器部分
  5. custom_unet_model.encoder.load_state_dict(pretrained_resnet50.state_dict())

这将加载resnet50的权重到自定义UNet模型的编码器中,然后您可以对整个模型进行微调以适应您的分割任务。

英文:

if you have a standard Unet encoder such as resnet50, then it's easy to add pertaining to it. for example:

  1. ENCODER = 'resnet50'
  2. ENCODER_WEIGHTS = 'imagenet'
  3. CLASSES = class_names
  4. ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multiclass segmentation
  5. # create segmentation model with pretrained encoder
  6. model = smp.Unet(
  7. encoder_name=ENCODER,
  8. encoder_weights=ENCODER_WEIGHTS,
  9. classes=len(CLASSES),
  10. activation=ACTIVATION,
  11. )
  12. preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)

However, suppose you have a custom-made Unet (not necessarily use resent50) encoder such as:

  1. class VGGBlock(nn.Module):
  2. def __init__(self, in_channels, middle_channels, out_channels):
  3. super().__init__()
  4. self.relu = nn.ReLU(inplace=True)
  5. self.conv1 = nn.Conv2d(in_channels, middle_channels, 3, padding=1)
  6. self.bn1 = nn.BatchNorm2d(middle_channels)
  7. self.conv2 = nn.Conv2d(middle_channels, out_channels, 3, padding=1)
  8. self.bn2 = nn.BatchNorm2d(out_channels)
  9. def forward(self, x):
  10. out = self.conv1(x)
  11. out = self.bn1(out)
  12. out = self.relu(out)
  13. out = self.conv2(out)
  14. out = self.bn2(out)
  15. out = self.relu(out)
  16. return out
  17. class UNet(nn.Module):
  18. def __init__(self, num_classes, input_channels=3, **kwargs):
  19. super().__init__()
  20. nb_filter = [32, 64, 128, 256, 512]
  21. self.pool = nn.MaxPool2d(2, 2)
  22. self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
  23. self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
  24. self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
  25. self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
  26. self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
  27. self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])
  28. self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])
  29. self.conv2_2 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
  30. self.conv1_3 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
  31. self.conv0_4 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
  32. self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  33. def forward(self, input):
  34. x0_0 = self.conv0_0(input)
  35. x1_0 = self.conv1_0(self.pool(x0_0))
  36. x2_0 = self.conv2_0(self.pool(x1_0))
  37. x3_0 = self.conv3_0(self.pool(x2_0))
  38. x4_0 = self.conv4_0(self.pool(x3_0))
  39. x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
  40. x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
  41. x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
  42. x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))
  43. output = self.final(x0_4)
  44. return output

How to do Imagenet pretraining for the encoder. I assume doing pretraining for the encoder from scratch will take long time. Is there a way to utilize an existing pre-trained encoder such as the resnet50 for such Unet.

答案1

得分: 0

是的,可以只使用一个预训练的块,而不是使用整个网络,比如来自Torchvisionresnet50。由于你提到了基于VGG类型块的自定义编码器,我是基于这个来回答的。

VGGBlock中,你可以调用预训练的VGG网络,然后选择最多到第二个卷积层,而不是手动定义层:

  1. # 必要的导入
  2. from torchvision.models import vgg16_bn
  3. import torch
  4. import torch.nn as nn
  5. from copy import deepcopy
  6. # 初始化带有BatchNorm的预训练vgg16网络
  7. model = vgg16_bn(pretrained=True)

然后,你可以通过以下方式修改你的VGGBlock

  1. class VGGBlock(nn.Module):
  2. def __init__(self, in_channels, out_channels):
  3. super().__init__()
  4. self.vggblock = deepcopy(model.features[:6])
  5. self.vggblock[0].in_channels = in_channels
  6. self.vggblock[0].out_channels = out_channels
  7. self.vggblock[1].num_features = out_channels
  8. self.vggblock[3].in_channels = out_channels
  9. self.vggblock[3].out_channels = out_channels
  10. self.vggblock[4].num_features = out_channels
  11. def forward(self, x):
  12. out = self.vggblock(x)
  13. return out

我也稍微修改了你的UNet类,以下是修改后的代码:

  1. class UNet(nn.Module):
  2. def __init__(self, num_classes, input_channels):
  3. super().__init__()
  4. nb_filter = [32, 64, 128, 256, 512]
  5. self.pool = nn.MaxPool2d(2, 2)
  6. self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
  7. self.conv0_0 = VGGBlock(input_channels, nb_filter[0])
  8. self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1])
  9. self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2])
  10. self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3])
  11. self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4])
  12. self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3])
  13. self.conv2_2 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2])
  14. self.conv1_3 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1])
  15. self.conv0_4 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0])
  16. self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  17. def forward(self, input):
  18. x0_0 = self.conv0_0(input)
  19. x1_0 = self.conv1_0(self.pool(x0_0))
  20. x2_0 = self.conv2_0(self.pool(x1_0))
  21. x3_0 = self.conv3_0(self.pool(x2_0))
  22. x4_0 = self.conv4_0(self.pool(x3_0))
  23. x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
  24. x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
  25. x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
  26. x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))
  27. output = self.final(x0_4)
  28. return output

请注意,在VGGBlockUNet类中,我跳过了你在代码片段中使用的middle_channels,因为这个输入参数实际上是无关紧要的,因为你的middle_channelsout_channels本质上是相同的。上述代码将使用预训练权重构建你在问题中发布的确切UNet架构。

英文:

Yes, it is possible to use only a pre-trained block instead of using the entire network such as resnet50 from Torchvision. Since you mentioned a custom encoder based on a VGG-type block, I'm answering based on that.
Instead of defining the layers in the VGGBlock manually, you can just call the pre-trained VGG network within that class and then select up to the 2nd conv layer.

First, you would need to get the pre-trained VGG network from Torchvision:

  1. # Necessary imports
  2. from torchvision.models import vgg16_bn
  3. import torch
  4. import torch.nn as nn
  5. from copy import deepcopy
  6. # Initializing the pre-trained vgg16 (with BatchNorm) network from torchvision
  7. model = vgg16_bn(pretrained = True)

Then, you can modify your VGGBlock by the following:

  1. class VGGBlock(nn.Module):
  2. def __init__(self, in_channels, out_channels):
  3. super().__init__()
  4. self.vggblock = deepcopy(model.features[:6])
  5. self.vggblock[0].in_channels = in_channels
  6. self.vggblock[0].out_channels = out_channels
  7. self.vggblock[1].num_features = out_channels
  8. self.vggblock[3].in_channels = out_channels
  9. self.vggblock[3].out_channels = out_channels
  10. self.vggblock[4].num_features = out_channels
  11. def forward(self, x):
  12. out = self.vggblock(x)
  13. return out

I also modified your UNet class a bit and this is the modified code:

  1. class UNet(nn.Module):
  2. def __init__(self, num_classes, input_channels):
  3. super().__init__()
  4. nb_filter = [32, 64, 128, 256, 512]
  5. self.pool = nn.MaxPool2d(2, 2)
  6. self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
  7. self.conv0_0 = VGGBlock(input_channels, nb_filter[0])
  8. self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1])
  9. self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2])
  10. self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3])
  11. self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4])
  12. self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3])
  13. self.conv2_2 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2])
  14. self.conv1_3 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1])
  15. self.conv0_4 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0])
  16. self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  17. def forward(self, input):
  18. x0_0 = self.conv0_0(input)
  19. x1_0 = self.conv1_0(self.pool(x0_0))
  20. x2_0 = self.conv2_0(self.pool(x1_0))
  21. x3_0 = self.conv3_0(self.pool(x2_0))
  22. x4_0 = self.conv4_0(self.pool(x3_0))
  23. x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
  24. x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
  25. x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
  26. x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))
  27. output = self.final(x0_4)
  28. return output

You would notice that, both in the VGGBlock and in the UNet class, I skipped the use of middle_channels as you did in your snippets. That input argument is actually irrelevant since your middle_channels and out_channels are essentially the same. The above code would build you the exact UNet architecture that you posted in the question with pre-trained weights.

huangapple
  • 本文由 发表于 2023年7月14日 04:21:42
  • 转载请务必保留本文链接:https://go.coder-hub.com/76682993.html
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定