期望的输入批次大小(24)与目标批次大小(18)不匹配

Question

class GCNModel(nn.Module):
    def __init__(self, in_channels, hidden_dim, out_channels, edge_dim): # 5, 64, 6, 3
        super(GCNModel, self).__init__() 
        self.conv1 = Edge_GCNConv(in_channels=in_channels, out_channels=hidden_dim, edge_dim=edge_dim)
        self.conv2 = Edge_GCNConv(in_channels=hidden_dim, out_channels=out_channels, edge_dim=edge_dim)
        self.batch_norm1 = nn.BatchNorm1d(hidden_dim)  
        self.batch_norm2 = nn.BatchNorm1d(out_channels)  
        self.linear = nn.Linear(out_channels, out_channels)  

    def forward(self, x, edge_index, edge_attr):
        x, edge_index, edge_attr = x, edge_index, edge_attr
        
        x1 = self.conv1(x, edge_index, edge_attr)
        
        x1 = self.batch_norm1(x1)
        x1 = F.relu(x1)
        x1 = F.dropout(x1, p=0.1, training=self.training)
        # print("GCNModel x1:", x1) 
        # print("GCNModel x1.shape:", x1.shape) # (24, 64)
        

        x2 = self.conv2(x1, edge_index, edge_attr)
        x2 = self.batch_norm2(x2)
        x2 = F.relu(x2)
        
        # print("GCNModel x2:", x2)
        # print("GCNModel x2.shape:", x2.shape) # (24, 6)
        
        x2 = F.dropout(x2, p=0.1, training=self.training)
        
        out = self.linear(x2)
        print("GCNModel out:", out)
        print("GCNModel out.shape:", out.shape) # (24, 6)
        return out

def train_model(train_loader, val_loader, model, optimizer, output_dim, threshold_value, num_epochs=100, early_stopping_rounds=50, batch_size=4):
    """
        训练GNN模型，使用 k 折交叉验证

        Args:
            train_loader: 训练数据加载器
            val_loader: 验证数据加载器
            model: GNN模型
            optimizer: 优化器
            num_epochs: 训练轮数 (default: 100)
            early_stopping_rounds: 早停轮数 (default: 50)
    """        
    best_val_loss = float('inf')
    best_accuracy = 0  # Track best validation accuracy
    rounds_without_improvement = 0

    # 创建损失函数
    criterion = nn.CrossEntropyLoss() 
    # criterion = nn.BCEWithLogitsLoss() # 二分类
    
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        correct = 0
        total = 0

        for data in train_loader:
            optimizer.zero_grad()
            
            #################### error #################
            out = model(data.x, data.edge_index, data.edge_attr)

            # 转换 data.y 为多热编码
            one_hot_labels = convert_to_one_hot(data.y, output_dim)

            print("train_model out.shape:", out.shape) # (24, 6)
            print("train_model data.y.shape:", data.y.shape) # (18, 2)
            print("train_model data.edge_attr.shape:", data.edge_attr.shape) # (18, 3)
            print("train_model data.edge_attr:", data.edge_attr)
            print("train_model one_hot_labels.shape:", one_hot_labels.shape) # (18, 6)
            
            loss = criterion(out, one_hot_labels)
            
            #################################################
            # print("train_model loss:", loss)
            total_loss += loss.item()

            # print("torch.sigmoid(out):", torch.sigmoid(out))
            predicted = (torch.sigmoid(out) >= threshold_value).long()
            # print("predicted:", predicted)

            correct += (predicted == one_hot_labels).all(dim=1).sum().item()
            # print("correct:", correct)
            total += len(data.y)
            # print("total：", total )
            
            loss.backward()
            optimizer.step()

        avg_train_loss = total_loss / len(train_loader)
        train_accuracy = correct / total

这里提到的 out.shape 是通过 data.x.shape 得到的输出。

而我的 data.y.shape 是从 data.edge_attr 中获取的。

我该怎么做才能解决 out 和 one_hot_labels 之间的数量不匹配的问题呢？我应该修改模型，还是调整输出的维度呢？

out.shape 这个张量是通过 data.y.shape 得到的输出。