PyTorch构造TensorDataset报AssertionError张量尺寸不匹配问题求助

问题根因 #

你遇到的AssertionError: Size mismatch between tensors报错，本质是构造TensorDataset时传入的训练特征张量和训练标签张量的样本数量不匹配，触发了PyTorch源码里的长度校验逻辑。

代码错误点梳理 #

• 标签生成逻辑完全错误：train_data.status_group == "functional needs repair"仅做了条件判断，没有赋值给任何变量，后续生成train_labels时用到的train_data_new是未定义变量，直接导致标签张量的长度和特征张量不匹配
• 神经网络前向传播逻辑错误：定义了fc1全连接层和relu激活函数，但forward方法里直接将输入传入输出层，未使用已定义的隐藏层
• 训练循环变量未定义：循环中用到的lrmse变量没有提前定义，运行到这里会直接报错
• 损失函数缩进错误：loss = nn.BCELoss()写在了Net类的内部作用域，属于逻辑错误

修复后的可运行代码 #

import numpy as np
import pandas as pd
import torch
from torch import nn

train_link = "/content/pump-data/train_pumps.csv"
test_link = "/content/pump-data/test_pumps.csv"

train_data = pd.read_csv(train_link)
test_data = pd.read_csv(test_link)

# 修正标签生成逻辑：将status_group转为二分类标签
train_data['label'] = (train_data.status_group == "functional needs repair").astype(int)

all_features = pd.concat((train_data.iloc[:, 1:-2], test_data.iloc[:, 1:]))

numeric_features = all_features.dtypes[all_features.dtypes != 'object'].index

all_features[numeric_features] = all_features[numeric_features].apply(
    lambda x: (x - x.mean()) / (x.std()))

all_features[numeric_features] = all_features[numeric_features].fillna(0)

all_features = pd.get_dummies(all_features, dummy_na=True)

n_train = train_data.shape[0]

train_features = torch.tensor(all_features[:n_train].values,
                              dtype=torch.float32)
test_features = torch.tensor(all_features[n_train:].values,
                             dtype=torch.float32)

# 修正标签来源，直接用train_data的label列
train_labels = torch.tensor(train_data['label'].values.reshape(-1, 1),
                            dtype=torch.float32)


class Net(nn.Module):
  def __init__(self, in_features, out_features):
    super().__init__()
    self.fc1 = nn.Linear(in_features, 20)
    self.output = nn.Linear(20, out_features)
    self.relu = nn.ReLU()

  def forward(self, x):
    # 修正前向传播逻辑，先走隐藏层+激活，再走输出层
    x = self.relu(self.fc1(x))
    x = self.output(x)
    return torch.sigmoid(x)

# 修正损失函数缩进，放到全局作用域
loss = nn.BCELoss()
in_features = train_features.shape[1]

def train(net, train_features, train_labels,
          num_epochs, learning_rate, weight_decay, batch_size):
  train_ls = []
  train_tensor = torch.utils.data.TensorDataset(train_features, 
                                                train_labels)
  train_loader = torch.utils.data.DataLoader(train_tensor,
                                             batch_size=batch_size, shuffle=False)

  optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate,
                               weight_decay=weight_decay)

  for epoch in range(num_epochs):
    epoch_loss = 0
    for X, y in train_loader:
      optimizer.zero_grad()
      output = net(X)
      l = loss(output, y)
      l.backward()
      optimizer.step()
      epoch_loss += l.item()
    # 修正未定义的lrmse变量，改用计算得到的平均损失
    avg_loss = epoch_loss / len(train_loader)
    train_ls.append(avg_loss)
    if(epoch%10==0):
      print('epoch {}, loss {}'.format(epoch, avg_loss))
  return train_ls


net = Net(in_features, 1)

num_epochs = 2
learning_rate = 0.80
weight_decay = 24
batch_size = 24

train_ls = train(net, train_features, train_labels,
                 num_epochs, learning_rate,
                 weight_decay, batch_size)

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