pytorch 多分类

PyTorch中的多分类任务通常使用交叉熵损失函数进行训练。以下是一个简单的多分类示例，假设你有一个包含10个类别的数据集：

import torch
import torch.nn as nn
import torch.optim as optim

# 假设你的数据集有10个类别
num_classes = 10

# 定义一个简单的神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 64)
        self.fc3 = nn.Linear(64, num_classes)

    def forward(self, x):
        x = x.view(-1, 784)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 加载数据集
train_loader = ...

# 定义模型、损失函数、优化器
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):
    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        if i % 100 == 0:
            print(f"Epoch {epoch}, Iteration {i}, Loss: {loss.item()}")

在这个示例中，我们定义了一个简单的神经网络，包含3个全连接层，并使用ReLU作为激活函数。我们使用交叉熵损失函数进行训练，并使用随机梯度下降优化器进行参数更新。在训练循环中，我们使用`nn.CrossEntropyLoss`计算损失，然后通过反向传播更新模型参数。