深度学习模型加载张量计算RMSE时遇AttributeError问题排查

核心错误原因 #

• enumerate使用错误：for x, rul in enumerate(self.train_data)里的x是迭代的整数索引，并非数据张量，这直接导致x.to(self.device)触发int对象无to方法的报错。
• DataLoader返回值处理错误：你的CustomDataset.__getitem__返回的是字典{'input': input_tensor, 'labels': label_tensor}，但train_epoch里错误地将rul当成了标签张量，实际rul是包含输入和标签的batch字典。

修复步骤 #

1. 修正train_epoch循环逻辑 #

去掉enumerate，直接迭代DataLoader，从返回的字典中提取输入和标签：

def train_epoch(self, epoch):
    train_loss = 0.0
    self.model.train()
    # 直接迭代train_data，每个batch是包含input和labels的字典
    for batch in self.train_data:
        self.model.zero_grad()
        # 从字典中取出输入与标签，并转移到指定设备
        x = batch['input'].to(self.device).float()
        rul = batch['labels'].to(self.device).float()
        out = self.model(x)
        # 计算RMSE：基于MSE损失开根号
        loss = torch.sqrt(self.criterion(out.float(), rul))
        loss.backward()
        self.optimizer.step()
        train_loss += loss.item()  # 提取标量值避免张量累积
    # 计算平均训练损失
    avg_train_loss = train_loss / len(self.train_data)
    print(f'Epoch {epoch}, Train RMSE: {avg_train_loss:.4f}')

2. 修复CustomDataset的维度错误（潜在隐患） #

在load_dataset中，self.data_ = np.expand_dims(self.data_, axis=0)会让数据集维度变为(1, N, ...)，导致__len__返回1，整个数据集只有1个样本，DataLoader的batch_size设置完全失效。需删除这行代码：

def load_dataset(self):
    # ... 其他代码保持不变
    self.data_ = x_data['data'].transpose(1,0)
    # 移除错误的维度扩展
    # self.data_ = np.expand_dims(self.data_, axis=0)
    self.labels = y_data['label']
    self.config['data_shape'] = self.data_.shape[1:]

3. 验证数据集维度 #

确保self.data_.shape为(样本数, 特征维度...)，self.labels.shape为(样本数,)或(样本数, 1)，保证DataLoader能按batch_size正确打包数据。

额外优化建议 #

• 计算训练损失时使用loss.item()提取标量值，避免累积大量张量占用内存。
• 若criterion已设置为torch.nn.MSELoss()，torch.sqrt(criterion(out, rul))即可得到RMSE，无需额外重复转换float类型（除非模型输出与标签类型不匹配）。

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