pytorch 多分类的问题

训练的数据是 300x10 的矩阵，一共 14000 个，每个矩阵对应三种情况，表示为 0 1 2

需要特别说明的是，大部分的数据所对应的类别是 0，少数是对应着 1 2 的。可以理解为数据中 50%是 0 类，剩下的 50%分给了 1 2 类

我使用的是逻辑回归

网络如下

class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression, self).__init__()
        self.lr = Linear(300 * 10, 3)
        self.sm = nn.Sigmoid()
       
    def forward(self, x: Tensor):
    	# 展开
    	x = x.view(-1, self.in_features)
        x = self.lr(x)
        x = self.sm(x)
        return x

优化器使用的是 SGD 梯度下降，学习率是 0.001 损失函数用的是 CrossEntropyLoss

训练部分代码如下 其中 x_train 是训练数据，y_train 对应的类别，y_train 中的每个元素均为 0/1/2，没有经过独热编码 x_train 的 shape 为torch.Size([14000, 300, 10]) y_train 的 shape 为torch.Size([14000, 1])

model = LogisticRegression().to(device)

optimizer = optim.SGD(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(1, 10001):
    for i, (data, target) in enumerate(zip(x_train, y_train)):
        x, y = data.to(device), target.to(device)

        y_pred = model(x)
		
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            if y == 2:
                print(f'Epoch: {epoch}, Loss: {loss.item()}')

如果训练所有数据，会出现 loss 前期下降快，后面基本不动 我请教过他人，提出我的网络层数太少，出现了过拟合，让我多加几层网络 加了几层后我的网络是这样的

class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression, self).__init__()
        self.lr = Linear(300 * 10, 200 * 10)
        self.lr2 = Linear(200 * 10, 3)
        self.sm = nn.Sigmoid()
       
    def forward(self, x: Tensor):
        # 展开
    	x = x.view(-1, self.in_features)
        x = self.lr(x)
        x = self.lr(x)
        x = self.sm(x)
        return x

接着我训练，loss 会一直增长

期间我试过取 0 和 2 两个类别的数据训练 10000 次，结果是 97 的准确率

我想知道问题出在哪了，谢谢