PyTorch线性神经网络出现ArrayRef无效索引RuntimeError问题求助

错误原因分析 #

报错核心是输入张量维度不匹配：nn.Linear(1,1)要求输入为至少1维的张量（形状如[N,1]或[1]），但原代码中x_train的每个元素是0维标量，传入Linear层时触发维度错误；同时labels也是标量，和模型输出的1维张量形状不匹配，会导致后续损失计算异常。

修正后的可运行代码 #

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 生成数据时直接转为二维张量，避免后续维度调整
samples = torch.linspace(0, 100, 100).unsqueeze(1)  # 形状变为(100, 1)
train_split = int(len(samples)*0.8)
x_train, x_test = samples[:train_split], samples[train_split:]
y_labels = 2*samples - 4  # 形状(100, 1)
y_labels += torch.tensor(np.random.normal(0, 5, len(samples))).unsqueeze(1)  # 匹配维度

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 1)
    def forward(self, x):
        return self.fc1(x)

model = NeuralNetwork()
loss_func = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)  # 调大学习率加速收敛
num_epochs = 200  # 增加训练轮数

for epoch in range(num_epochs):
    # 批量处理（用整个训练集作为一个批次，效率更高）
    y_pred = model(x_train)
    loss = loss_func(y_pred, y_labels[:train_split])
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # 每10轮打印一次损失，避免输出冗余
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1} - loss: {loss.item():.4f}")

# 测试集验证
with torch.no_grad():
    test_pred = model(x_test)
    test_loss = loss_func(test_pred, y_labels[train_split:])
    print(f"\nTest Loss: {test_loss.item():.4f}")

关键修正点 #

• 数据维度调整：通过unsqueeze(1)给每个样本增加维度，确保符合Linear层输入要求
• 优化学习率：从0.001改为0.01，针对简单线性拟合任务加速收敛
• 批量训练替换逐样本训练：效率更高，符合PyTorch常规使用逻辑
• 修复打印格式：去掉原代码中损失打印语句多余的%符号

进阶改进建议 #

• 使用DataLoader管理数据：支持自动批处理、数据打乱、多线程加载，代码更规范：

  from torch.utils.data import TensorDataset, DataLoader
  train_dataset = TensorDataset(x_train, y_labels[:train_split])
  train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
  # 训练循环改为遍历DataLoader
  for epoch in range(num_epochs):
      for batch_x, batch_y in train_loader:
          y_pred = model(batch_x)
          loss = loss_func(y_pred, batch_y)
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  

• 尝试添加隐藏层（可选）：虽然任务是线性拟合，可练习非线性模型构建：

  class NeuralNetwork(nn.Module):
      def __init__(self):
          super().__init__()
          self.fc1 = nn.Linear(1, 16)
          self.fc2 = nn.Linear(16, 1)
          self.relu = nn.ReLU()
      def forward(self, x):
          x = self.relu(self.fc1(x))
          return self.fc2(x)
  

• 加入学习率调度器：训练后期降低学习率，帮助模型收敛到更优值：

  from torch.optim.lr_scheduler import StepLR
  scheduler = StepLR(optimizer, step_size=50, gamma=0.5)
  # 训练循环中每次optimizer.step()后调用
  scheduler.step()
  

• 可视化训练结果：用matplotlib绘制损失曲线、预测值与真实值对比图，直观观察模型效果

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