使用 3 通道 (RGB) PyTorch 模型对 4 通道 (RGBY) 图像进行分类

Question

我用 4 通道图像 (RGBY) 标记了数据集。我想使用预训练的分类模型（使用pytorch 和 ResNet50 作为模型）。不过，所有pytorch 型号都适用于 3 通道。 所以，问题是：如何我可以使用 3-channel 预训练模型来处理 4-channels 数据？我正在加载模型：

import torchvision.models as models
resnet50 = models.resnet50(pretrained=True)

Answer 1

您可以修改 CNN 的第一层，使其期望 4 个输入通道而不是 3 个。在您的情况下，第一层是 resnet50.conv1。所以：

import torchvision.models as models
resnet50 = models.resnet50(pretrained=True)

# modify first layer so it expects 4 input channels; all other parameters unchanged
resnet50.conv1 = torch.nn.Conv2d(4,64,kernel_size = (7,7),stride = (2,2), padding = (3,3), bias = False) 

# test
inp = torch.rand([1,4,512,512])
resnet50.eval()
resnet50.training = False
out = resnet50(inp) # should evaluate without error

以下实现细节使这种更改的简单性成为可能：对于 2D 卷积（也适用于其他维度的卷积），pytorch 为每个所需的输出平面（特征图）与每个输入平面卷积一个内核。这导致 n_input_planes x n_output_planes 总特征图（在本例中分别为 4 和 64）。 Pytorch 然后对每个输出平面的所有输入平面求和，无论输入平面的数量如何，都会产生总共 n_output_planes 个平面。

好消息是，这意味着您可以添加额外的输入平面（地图），而无需修改第一层之后的网络。 （可能在某些情况下）不利的部分是所有输入特征图都被同等对待，并且每个特征图的信息在第一次卷积结束时完全合并。在某些情况下，可能需要在开始时对输入特征图进行不同的处理，在这种情况下，您需要定义两个单独的 CNN 分支，以便不会在每一层将特征加在一起。