PyTorch批量训练的正确张量形状？恒星识别模型张量相乘报错求助

你的代码报错核心原因是张量形状不匹配、数据类型错误，以及若干语法/导入问题，以下是具体问题点和修复步骤：

1. 缺失必要导入 #

代码用到Dataset、DataLoader和numpy但未导入，需添加：

import numpy as np
from torch.utils.data import Dataset, DataLoader

2. 输入张量形状错误 #

你强行将输入inputsreshape为(-1,500)，但输入特征数是5，DataLoader输出的batch形状已经是[batch_size, 5]，无需额外修改。错误的形状会导致第一层线性层nn.Linear(5,4)接收到维度不匹配的输入，引发张量相乘报错。

3. 标签数据类型与形状错误 #

CrossEntropyLoss要求：

• 标签是一维张量（形状[batch_size]），你的代码中y_data是[n_samples,1]的二维张量
• 标签必须是整数类型（torch.long），你用np.float32读取后转成了float张量

修改Dataset中的标签处理：

self.y_data = torch.from_numpy(xy[:, 0]).type(torch.long)  # 去掉[:, [0]]转为一维，同时转成long类型

4. 拼写错误 #

opimizer是optimizer的错拼，训练时会找不到变量，修正为：

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

5. 训练步数计算错误 #

n_total_steps应该是数据加载器的长度（即batch数量），而非数据集样本数，修正为：

n_total_steps = len(data_loader)

6. 冗余的testSDSS类 #

测试集和训练集读取同一文件，无需重复定义Dataset类，直接复用SDSS类即可：

test_dataset = SDSS()

完整修复后的代码 #

import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import Dataset, DataLoader

# device config
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

input_size = 5
hidden_size = 4
num_classes = 3
num_epochs = 2
batch_size = 100
learning_rate = 0.001

class SDSS(Dataset):
    def __init__(self):
        xy = np.loadtxt('SDSS.csv', delimiter=',', dtype=np.float32, skiprows=0)
        self.n_samples = xy.shape[0]
        # 特征：[n_samples, 5]
        self.x_data = torch.from_numpy(xy[:, 1:])
        # 标签：一维long张量，[n_samples]
        self.y_data = torch.from_numpy(xy[:, 0]).type(torch.long)

    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        return self.n_samples

# 实例化数据集和数据加载器
dataset = SDSS()
test_dataset = SDSS()  # 若测试集是不同文件，可给SDSS类添加文件路径参数

data_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet,self).__init__()
        self.l1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.LeakyReLU()
        self.l2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.l1(x)
        out = self.relu(out)
        out = self.l2(out)
        return out

model = NeuralNet(input_size, hidden_size, num_classes).to(device)

# 验证batch形状（可选）
dataiter = iter(data_loader)
features, labels = next(dataiter)
print(f"特征形状: {features.shape}, 标签形状: {labels.shape}, 标签类型: {labels.dtype}")

# loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# training loop
n_total_steps = len(data_loader)
for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(data_loader):
        inputs = inputs.to(device)
        labels = labels.to(device)
        
        # forward pass
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # backward pass and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 10 == 0:  # 根据实际batch数量调整打印间隔
            print(f'epoch {epoch + 1}/{num_epochs}, step {i+1}/{n_total_steps}, loss = {loss.item():.4f}')

数据格式说明 #

针对你的恒星识别任务：

• 输入特征：每个样本是5维向量，数据集特征张量形状为[总样本数, 5]，经DataLoader后每个batch形状为[batch_size, 5]，完全匹配模型第一层输入要求。
• 标签：3分类任务下，标签是0/1/2的整数，张量形状为[总样本数]，每个batch形状为[batch_size]，符合CrossEntropyLoss的要求。

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