Keras model.fit与PyTorch训练流程差异及不平衡数据集适配问题

问题背景 #

将Keras实现的ECG心跳分类模型转换为PyTorch代码后，发现无论是否使用类权重，训练损失均无法正常下降。Keras仅需调用model.fit(X, Y, epochs=1000, verbose=2, validation_split=0.1)即可完成训练，而PyTorch需要手动编写训练流程，现需明确两者的核心差异及数据集采样处理的不同。

用户提供的PyTorch模型代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.utils.class_weight import compute_class_weight
import numpy as np
import torch.optim as optim

class Model(nn.Module):
    def __init__(self, nclass):
        super(ConvModel, self).__init__()
        
        self.conv1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=5, padding=0, dtype=float)
        self.conv2 = nn.Conv1d(in_channels=16, out_channels=16, kernel_size=5, padding=0, dtype=float)
        self.pool1 = nn.MaxPool1d(kernel_size=2)
        self.dropout1 = nn.Dropout(0.1)
        
        self.conv3 = nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3, padding=0, dtype=float)
        self.conv4 = nn.Conv1d(in_channels=32, out_channels=32, kernel_size=3, padding=0, dtype=float)
        self.pool2 = nn.MaxPool1d(kernel_size=2)
        self.dropout2 = nn.Dropout(0.1)
        
        self.conv5 = nn.Conv1d(in_channels=32, out_channels=32, kernel_size=3, padding=0, dtype=float)
        self.conv6 = nn.Conv1d(in_channels=32, out_channels=32, kernel_size=3, padding=0, dtype=float)
        self.pool3 = nn.MaxPool1d(kernel_size=2)
        self.dropout3 = nn.Dropout(0.1)
        
        self.conv7 = nn.Conv1d(in_channels=32, out_channels=256, kernel_size=3, padding=0, dtype=float)
        self.conv8 = nn.Conv1d(in_channels=256, out_channels=256, kernel_size=3, padding=0, dtype=float)
        self.global_max_pool = nn.AdaptiveMaxPool1d(1)
        self.dropout4 = nn.Dropout(0.2)
        
        self.fc1 = nn.Linear(256, 64, dtype=float)
        self.fc2 = nn.Linear(64, 64, dtype=float)
        self.fc3 = nn.Linear(64, nclass, dtype=float)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.pool1(x)
        x = self.dropout1(x)
        
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.pool2(x)
        x = self.dropout2(x)
        
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = self.pool3(x)
        x = self.dropout3(x)
        
        x = F.relu(self.conv7(x))
        x = F.relu(self.conv8(x))
        
        x = self.global_max_pool(x)
        x = torch.flatten(x, 1)  # Flatten the output for fully connected layers
        x = self.dropout4(x)
        
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        
        return torch.softmax(x, dim=1)

用户提供的PyTorch训练代码：

# Hyperparameters
input_dim = 187  # Original time series length
num_classes = 5  # Number of output classes

# Instantiate the model
model = Model(num_classes)
# print(model)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# Define loss function and optimizer
# 假设Y是训练标签张量
class_weights=compute_class_weight('balanced',classes=np.unique(Y),y=Y.numpy())
class_weights=torch.tensor(class_weights,dtype=torch.float64).to(device)
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Train the model
n_epochs = 1000
for epoch in range(n_epochs):
    model.train()
    last_loss = 0
    running_loss = 0
    
    for i, (X_batch, Y_batch) in enumerate(train_loader):
        X_batch = X_batch.to(device)
        Y_batch = Y_batch.to(device)
        output = model(X_batch)
        loss = criterion(output, Y_batch)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
     
        running_loss += loss.item()
        if i % 200 == 199:
            last_loss = running_loss / 200 # loss per batch
            print('  batch {} loss: {}'.format(i + 1, last_loss))
            running_loss = 0.

对比的Keras训练代码：

model.fit(X, Y, epochs=1000, verbose=2, validation_split=0.1)

核心差异与问题排查 #

1. 模型输出层的本质差异 #

PyTorch的nn.CrossEntropyLoss内部已集成LogSoftmax和NLLLoss，而你在模型forward方法中额外添加了torch.softmax，相当于做了两次概率归一化，会导致梯度消失，这是损失无法下降的核心原因。

修复方案：删除输出层的softmax，直接返回logits：

# 替换原return语句
return x

2. 训练流程的梯度处理错误 #

你的PyTorch训练代码中，optimizer.zero_grad()的位置错误——应该在每个批次计算损失之前清零梯度，否则会累积上一批次的梯度，导致参数更新混乱。

修复后的训练循环顺序：

for i, (X_batch, Y_batch) in enumerate(train_loader):
    optimizer.zero_grad()  # 移到最前面，先清零梯度
    X_batch = X_batch.to(device)
    Y_batch = Y_batch.to(device)
    output = model(X_batch)
    loss = criterion(output, Y_batch)
    loss.backward()
    optimizer.step()  # 最后更新参数

3. 数据集采样与类权重的处理差异 #

• Keras默认行为：model.fit默认开启shuffle=True，每个epoch会随机打乱数据集。如果原Keras代码无需手动传class_weight也能训练，大概率是原代码在数据集阶段做了过采样/欠采样处理（比如对少数类数据进行复制），而非依赖损失权重。
• PyTorch的处理：若你的train_loader未设置shuffle=True，模型会按固定顺序训练，容易陷入局部最优。另外，需确保compute_class_weight计算的权重顺序与标签的类别顺序完全匹配，否则权重会失效。

4. 数据维度与模型初始化错误 #

• 输入维度匹配：PyTorch的Conv1d要求输入格式为(batch_size, in_channels, seq_len)，而Keras的Conv1D默认是(batch_size, seq_len, in_channels)。需确保你的X_batch已转换为(batch, 1, 187)的格式，否则in_channels=1会不匹配。
• 模型继承错误：你的模型类名为Model，但__init__中写了super(ConvModel, self).__init__()，这会导致类初始化异常，需修正为：

  super(Model, self).__init__()
  

5. 张量类型一致性问题 #

代码中使用dtype=float属于Python原生类型，PyTorch中应明确指定torch.float32或torch.float64，同时确保class_weights的 dtype 与模型参数 dtype 一致，否则会出现梯度计算不兼容的问题。

总结修复步骤 #

1. 修正模型继承语句，确保类初始化正确
2. 删除模型输出层的torch.softmax，直接返回logits
3. 调整梯度清零语句的位置，放在每个批次的最开始
4. 确保输入数据维度为(batch_size, 1, 187)，匹配Conv1d的输入要求
5. 给train_loader添加shuffle=True参数，保证每个epoch数据随机打乱
6. 统一张量类型，将dtype=float替换为torch.float32，并同步类权重的类型

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