PyTorch分割模型通道不匹配RuntimeError问题求助

错误信息 #

RuntimeError: Given groups=1, weight of size [64, 2, 3, 3], expected input[1, 1, 256, 256] to have 2 channels, but got 1 channels instead

用户声明已设置in_channels=2但仍触发该错误，以下是相关代码、输入输出信息及错误回溯：

分割模型代码 #

class SegmentationModel(nn.Module):
    def __init__(self, in_channels):
        super(SegmentationModel, self).__init__()
        self.conv1 = nn.Conv2d(2, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.conv5 = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
        self.upconv1 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.upconv2 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.upconv3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.upconv4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.final_conv = nn.Conv2d(64, 2, kernel_size=1)

    def forward(self, x):
        # Assuming x is of shape [batch_size, channels, height, width]
        x1 = F.relu(self.conv1(x))
        x2 = F.relu(self.conv2(x1))
        x3 = F.relu(self.conv3(x2))
        x4 = F.relu(self.conv4(x3))
        x5 = F.relu(self.conv5(x4))
        
        x6 = F.relu(self.upconv1(x5))
        x7 = F.relu(self.upconv2(x6))
        x8 = F.relu(self.upconv3(x7))
        x9 = F.relu(self.upconv4(x8))
        
        output = self.final_conv(x9)
  
        return output

输入输出形状 #

Image shape: torch.Size([2, 256, 256])

Mask shape: torch.Size([2, 256, 256])
Output shape: torch.Size([1, 2, 4096, 4096])

Resized output shape: torch.Size([1, 2, 256, 256])
Resized_Image shape: torch.Size([1, 2, 256, 256])
Resized_Mask shape: torch.Size([1, 2, 256, 256])

错误回溯 #

---------
RuntimeError                              Traceback (most recent call last)
Cell In[18], line 51
     47 optimizer.zero_grad()
     50  # Forward pass
---> 51 outputs = model(resized_images)
     53 # Resize the output tensor to match the spatial dimensions of the target tensor
     54 resized_outputs = F.interpolate(outputs, size=(256, 256), mode='bilinear', align_corners=False)

File /usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[16], line 17, in SegmentationModel.forward(self, x)
     15 def forward(self, x):
     16     # Assuming x is of shape [batch_size, channels, height, width]
---> 17     x1 = F.relu(self.conv1(x))
     18     x2 = F.relu(self.conv2(x1))
     19     x3 = F.relu(self.conv3(x2))

File /usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.8/site-packages/torch/nn/modules/conv.py:463, in Conv2d.forward(self, input)
    462 def forward(self, input: Tensor) -> Tensor:
--> 463     return self._conv_forward(input, self.weight, self.bias)

File /usr/local/lib/python3.8/site-packages/torch/nn/modules/conv.py:459, in Conv2d._conv_forward(self, input, weight, bias)
    455 if self.padding_mode != 'zeros':
    456     return F.conv2d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode),
    457                     weight, bias, self.stride,
    458                     _pair(0), self.dilation, self.groups)
--> 459 return F.conv2d(input, weight, bias, self.stride,
    460                 self.padding, self.dilation, self.groups)

RuntimeError: Given groups=1, weight of size [64, 2, 3, 3], expected input[1, 1, 256, 256] to have 2 channels, but got 1 channels instead

错误原因及修复方案 #

核心问题 #

错误信息明确显示实际传入模型的输入张量是1通道，但用户打印的Resized_Image shape是2通道，说明两者不是同一个变量，或者数据预处理/传输过程中通道维度被意外修改。另外模型定义存在参数冗余问题，但不是直接报错原因。

具体排查点及修复步骤 #

1. 确认实际输入模型的张量形状
   在模型的forward函数开头添加打印语句，验证传入的张量维度：

   def forward(self, x):
       print("Actual input shape:", x.shape)  # 新增打印
       x1 = F.relu(self.conv1(x))
       # ... 其余代码不变
   

2. 修复模型参数冗余问题
   模型__init__方法接收in_channels参数，但第一层卷积硬编码为2，改为使用传入的参数，保证模型灵活性：

   def __init__(self, in_channels):
       super(SegmentationModel, self).__init__()
       self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)  # 替换硬编码的2为in_channels
       # ... 其余代码不变
   

3. 检查数据预处理流程

   • 确保输入张量的维度顺序是[batch_size, channels, height, width]，如果原始图像是[height, width, channels]格式，需要用x = x.permute(2,0,1)转置，再增加batch维度
   • 如果确实需要2通道输入，但原始数据是单通道，通过repeat扩展通道：

     # 单张图像：[1,256,256] -> [1,2,256,256]
     x = x.unsqueeze(0).repeat(1,2,1,1)
     # 批量图像：[batch,1,256,256] -> [batch,2,256,256]
     x = x.repeat(1,2,1,1)
     

   • 检查resized_images的生成代码，确认使用F.interpolate时没有修改通道数，比如：

     resized_images = F.interpolate(images, size=(256,256), mode='bilinear', align_corners=False)
     

4. 验证批量数据一致性
   确保每个批次的输入通道数都为2，避免部分批次因数据损坏导致通道数异常。

内容的提问来源于stack exchange，提问作者Muhammad Ismail