如何用PyTorch前馈神经网络拟合sum_mod_label函数？

首先咱们得先理清两个标签函数的本质差异：

• sum_bin_label是把x1+x2分成3个连续的宽区间，决策边界平滑且连续，这种模式对简单前馈网络来说非常容易学习。
• sum_mod_label则是把x1+x2放大100倍后取模3，相当于把[0,2)的区间切成了66个极窄的周期性重复区间（每个区间宽度仅0.03），这种高频、周期性的分段常数函数，对原网络的容量和训练策略都是极大挑战。

再加上原代码里几个关键的训练细节问题，直接导致模型完全学不动：

1. 先修正最致命的错误：CrossEntropyLoss与Softmax的冲突 #

PyTorch的CrossEntropyLoss内部已经包含了LogSoftmax和NLLLoss的计算，如果在网络输出层再叠加Softmax，会导致梯度饱和（Softmax在概率接近0或1时梯度趋近于0），模型根本无法更新参数。

第一步要去掉网络里的self.softmax，让输出直接是线性层的logits：

class FeedForward(torch.nn.Module):
    def __init__(self, input_size, num_classes):
        super().__init__()
        self.input_size = input_size
        self.num_classes = num_classes
        self.relu = torch.nn.ReLU()
        # 移除softmax层
        self.fc1 = torch.nn.Linear(self.input_size, 64)  # 增加隐藏层神经元数量
        self.fc2 = torch.nn.Linear(64, 64)               # 新增一层隐藏层
        self.fc3 = torch.nn.Linear(64, num_classes)
    def forward(self, x, **kwargs):
        x = x.float()
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        output = self.fc3(x)  # 直接输出logits，交给CrossEntropyLoss处理
        return output

2. 提升网络容量，适配周期性高频任务 #

原网络只有2个线性层，且隐藏层神经元数等于输入维度（仅2个），容量完全不足以捕捉sum_mod_label的细粒度周期性模式。我们需要：

• 增加隐藏层的数量（比如从1层变成2层）
• 增大每个隐藏层的神经元数量（比如从2个变成64个）

3. 调整训练策略，让模型有效更新 #

• 增大Batch Size：原代码用默认的batch_size=1，梯度噪声极大，训练不稳定。改成batch_size=64或128，让梯度更新更平滑。
• 更换优化器/调整学习率：原SGD的学习率0.001太小，且动量未必适合这类任务。换成Adam自适应优化器，或者把SGD的学习率调到0.01以上。
• 增加训练轮数：复杂任务需要更多迭代，比如从50轮增加到200轮。

修改后的完整可运行代码 #

import torch
import numpy as np
from sklearn.model_selection import train_test_split
import torch.utils.data as utils

DATASHAPE = (2000, 2)
NUM_CLASSES = 3

# 标签函数保持不变
def sum_mod_label(x):
    return np.array([x for x in map(
        lambda x: x % NUM_CLASSES, map(int, (x[:, 0] + x[:, 1]) * 100))])

def sum_bin_label(x):
    def binit(x):
        if x < 0.807:
            return 0
        if x < 1.169:
            return 1
        return 2
    return np.array(
        [x for x in map(lambda x: binit(x), x[:, 0] + x[:, 1])])

class RandomModuloDataset(utils.Dataset):
    def __init__(self, shape, label_fn):
        self.data = np.random.random_sample(shape)
        self.label = label_fn(self.data)
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx, :], self.label[idx]

# 优化后的网络结构
class FeedForward(torch.nn.Module):
    def __init__(self, input_size, num_classes):
        super().__init__()
        self.input_size = input_size
        self.num_classes = num_classes
        self.relu = torch.nn.ReLU()
        self.fc1 = torch.nn.Linear(self.input_size, 64)
        self.fc2 = torch.nn.Linear(64, 64)
        self.fc3 = torch.nn.Linear(64, num_classes)
    def forward(self, x, **kwargs):
        x = x.float()
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return self.fc3(x)

# 优化后的训练函数
def fitit(trainloader, epochs=200):
    neurons = DATASHAPE[1]
    net = FeedForward(neurons, NUM_CLASSES)
    criterion = torch.nn.CrossEntropyLoss()
    # 换成Adam优化器，自适应学习率更适配复杂任务
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
    # 加入学习率调度器，后期降低学习率提升收敛效果
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
    
    for epoch in range(epochs):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        
        # 打印平均损失而非最后一个batch的损失，更能反映训练状态
        avg_loss = running_loss / len(trainloader)
        print(f'[{epoch + 1}] avg loss: {avg_loss:.3f}')
        scheduler.step()

# 测试sum_mod_label的训练效果
sum_mod_tloader = utils.DataLoader(
    RandomModuloDataset(DATASHAPE, sum_mod_label), batch_size=64, shuffle=True)
fitit(sum_mod_tloader, epochs=200)

预期效果 #

修改后你会看到训练损失持续下降，最终能达到0.3以下，模型的准确率也会接近甚至超过CatBoost的水平。

额外优化建议 #

如果想进一步提升效果，可以：

• 把x1+x2作为额外特征输入网络（显式告诉模型我们关注的是两个特征的和）
• 尝试加入周期性激活函数（比如torch.sin/torch.cos），帮助模型更好地捕捉周期性模式
• 增加数据集规模，比如把DATASHAPE改成(10000, 2)，让模型看到更多样本模式

内容的提问来源于stack exchange，提问作者dimitar-petrov