PyTorch运行时错误排查及石头剪刀布游戏模型代码问题分析

问题场景 #

需要实现一个基于PyTorch的石头剪刀布胜负判断程序：

• 输入：每个样本是两个玩家的独热张量，比如[[1,0,0],[0,1,0]]（[1,0,0]代表石头，[0,1,0]代表剪刀，[0,0,1]代表布）
• 输出：每个样本对应胜负结果，比如[1,0]表示玩家1获胜、玩家2失败，[1,1]表示平局

运行代码时触发错误：

RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 1

错误原因分析 #

1. 输入特征维度错误
   每个样本包含两个玩家的3维独热向量，需要将它们拼接成一个6维的特征向量作为模型输入，但原代码直接将形状为(4,2,3)的张量输入模型，且错误设置num_inputs=3，导致模型无法正确理解输入的样本结构。

2. 张量形状不匹配
   原模型接收(4,2,3)的输入后，Linear层会对每个3维向量单独处理，输出形状为(4,2,2)的张量。后续执行res.squeeze(dim=1)时，由于dim=1的维度大小是2而非1，squeeze操作无效，最终res的形状仍为(4,2,2)，与目标张量result的(4,2)形状不匹配，触发维度错误。

修正方案 #

步骤1：调整输入张量形状 #

将每个样本的两个3维向量拼接为6维向量，把输入形状从(4,2,3)转换为(4,6)。

步骤2：修正模型输入维度 #

将模型的num_inputs设置为6，匹配拼接后的输入特征维度。

修正后的完整代码 #

import torch
from torch import nn
torch.manual_seed(42)

# 原始输入：4个样本，每个样本是两个玩家的3维独热向量
input_data = [[1, 0, 0], [0, 1, 0]], [[1, 0, 0], [0, 0, 1]], [[0, 1, 0], [0, 1, 0]], [[0, 0, 1], [1, 0, 0]]
input_data = torch.tensor(input_data, dtype=torch.float32)
# 将每个样本的两个向量拼接成6维特征
input_data = torch.cat([input_data[:, 0], input_data[:, 1]], dim=1)

# 目标结果
result = [[1, 0], [0, 1], [1, 1], [1, 0]]
result = torch.tensor(result, dtype=torch.float32)

class ResultPredictor(nn.Module):
    def __init__(self, num_inputs, num_outputs):
        super().__init__()
        self.linear = nn.Linear(num_inputs, num_outputs)
        self.act_fn = nn.Sigmoid()
    def forward(self, x):
        x = self.linear(x)
        x = self.act_fn(x)
        return x

# 修正输入维度为6
model = ResultPredictor(num_inputs=6, num_outputs=2)
print(model)

# 查看参数形状
for name, param in model.named_parameters():
    print(f"Parameter {name}, shape {param.shape}")

# 测试单样本输入
for x in input_data:
    print(model(x))

print(model(input_data))
print('####################')

# 训练流程
model.train()
optimizer = torch.optim.SGD(model.parameters(), lr=0.3)
lossfunc = nn.MSELoss()

for _ in range(1000):
    res = model(input_data)
    loss = lossfunc(res, result)
    
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print(model(input_data))

说明 #

修正后，输入张量形状为(4,6)，模型输出形状为(4,2)，与目标张量形状完全匹配，可正常执行训练和推理。Sigmoid激活函数适合输出0-1之间的概率值，符合胜负结果的表示需求。

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