PyTorch实现Alpha Zero CNN时矩阵维度不匹配错误排查

使用PyTorch实现Alpha Zero游戏玩家的CNN时，遇到矩阵乘法维度不匹配错误。输入是3通道的10×10矩阵，执行代码：

model = Net(10, 10**2+1)
print(summary(model,(3,10,10)))

出现错误：

RuntimeError                              Traceback (most recent call last)
<ipython-input-32-0a101a882eb1> in <cell line: 3>()
      1 model = Net(10, 10**2+1)
      2 
----> 3 print(summary(model,(3,10,10)))

9 frames
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py in forward(self, input)
    112 
    113     def forward(self, input: Tensor) -> Tensor:
--> 114         return F.linear(input, self.weight, self.bias)
    115 
    116     def extra_repr(self) -> str:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (40x10 and 2x101)

当前模型架构：

def conv3x3(in_planes, out_planes):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, padding=1)

def conv1x1(in_planes, out_planes):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, padding=0)


class Net(nn.Module):
    def __init__(self, board_size, action_size, num_resBlocks=20, num_hidden=128):
        super().__init__()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Initial convolution 
        self.startBlock = nn.Sequential(
            conv3x3(3, num_hidden),
            nn.BatchNorm2d(num_hidden),
            nn.ReLU()
        )
        
        # Loop of all 20 Residual Layers
        self.backBone = nn.ModuleList(
            [ResBlock(num_hidden) for i in range(num_resBlocks)]
        )
        
        
        # Outputs expected value of the state 
        self.valueHead = nn.Sequential(
            conv1x1(num_hidden, 1),
            nn.BatchNorm2d(1),
            nn.ReLU(),

            nn.Linear(in_features=1, out_features=num_hidden),
            nn.ReLU(),

            nn.Linear(in_features=num_hidden, out_features=1),
            nn.Tanh()
        )

    
        # Outputs the probabilities of each possible action 
        self.policyHead = nn.Sequential(
            conv1x1(num_hidden, 2),
            nn.BatchNorm2d(2),
            nn.ReLU(),
            nn.Linear(2, out_features=(action_size)),
            nn.Softmax(dim=1)
        )


        self.to(self.device)

    def forward(self, x):
        x = self.startBlock(x)

        for resBlock in self.backBone:
            x = resBlock(x)

        policy = self.policyHead(x)
        value = self.valueHead(x)

        return policy, value
    

class ResBlock(nn.Module):
    def __init__(self, num_hidden):
        super().__init__()
        self.conv1 = conv3x3(num_hidden, num_hidden)
        self.bn1 = nn.BatchNorm2d(num_hidden)
        self.conv2 = conv3x3(num_hidden, num_hidden)
        self.bn2 = nn.BatchNorm2d(num_hidden)
        self.relu = nn.ReLU()

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        # Skip connections
        out = self.relu(out + identity)

        return out

错误核心是卷积层输出的4D张量直接传入Linear层，未展平为2D张量，导致维度不匹配：

• 残差网络输出张量形状为(batch_size, 128, 10, 10)
• 进入policyHead的conv1x1后变为(batch_size, 2, 10, 10)，直接传入Linear层时，PyTorch会自动展平后三个维度为2*10*10=200，但你定义的Linear层输入维度是2，因此出现维度冲突。

需要在policyHead和valueHead的Linear层前添加展平操作，同时修正Linear层的输入维度：

1. 修正valueHead #

conv1x1输出为(batch_size,1,10,10)，展平后维度是1*10*10=100，因此第一个Linear层的in_features需改为1*board_size*board_size：

self.valueHead = nn.Sequential(
    conv1x1(num_hidden, 1),
    nn.BatchNorm2d(1),
    nn.ReLU(),
    nn.Flatten(),  # 添加张量展平操作
    nn.Linear(in_features=1*board_size*board_size, out_features=num_hidden),
    nn.ReLU(),
    nn.Linear(in_features=num_hidden, out_features=1),
    nn.Tanh()
)

2. 修正policyHead #

conv1x1输出为(batch_size,2,10,10)，展平后维度是2*10*10=200，因此Linear层的in_features需改为2*board_size*board_size：

self.policyHead = nn.Sequential(
    conv1x1(num_hidden, 2),
    nn.BatchNorm2d(2),
    nn.ReLU(),
    nn.Flatten(),  # 添加张量展平操作
    nn.Linear(in_features=2*board_size*board_size, out_features=action_size),
    nn.Softmax(dim=1)
)

完整修正后的Net类 #

class Net(nn.Module):
    def __init__(self, board_size, action_size, num_resBlocks=20, num_hidden=128):
        super().__init__()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Initial convolution 
        self.startBlock = nn.Sequential(
            conv3x3(3, num_hidden),
            nn.BatchNorm2d(num_hidden),
            nn.ReLU()
        )
        
        # Loop of all 20 Residual Layers
        self.backBone = nn.ModuleList(
            [ResBlock(num_hidden) for i in range(num_resBlocks)]
        )
        
        
        # Outputs expected value of the state 
        self.valueHead = nn.Sequential(
            conv1x1(num_hidden, 1),
            nn.BatchNorm2d(1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(in_features=1*board_size*board_size, out_features=num_hidden),
            nn.ReLU(),
            nn.Linear(in_features=num_hidden, out_features=1),
            nn.Tanh()
        )

    
        # Outputs the probabilities of each possible action 
        self.policyHead = nn.Sequential(
            conv1x1(num_hidden, 2),
            nn.BatchNorm2d(2),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(in_features=2*board_size*board_size, out_features=action_size),
            nn.Softmax(dim=1)
        )


        self.to(self.device)

    def forward(self, x):
        x = self.startBlock(x)

        for resBlock in self.backBone:
            x = resBlock(x)

        policy = self.policyHead(x)
        value = self.valueHead(x)

        return policy, value

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