温馨提示×

Torch中如何处理多类分类任务

小樊
85
2024-03-25 11:11:47
栏目: 深度学习

在Torch中处理多类分类任务通常使用交叉熵损失函数和softmax函数。首先,定义一个包含所有可能类别的输出层,并使用softmax函数将输出映射为概率分布。然后,计算预测概率分布与真实标签之间的交叉熵损失,通过反向传播算法更新模型参数来最小化损失。以下是一个处理多类分类任务的示例代码:

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

# 定义模型
class Net(nn.Module):
    def __init__(self, num_classes):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, num_classes)
    
    def forward(self, x):
        x = self.fc(x)
        return x

# 初始化模型和损失函数
model = Net(num_classes=5)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 准备数据
inputs = torch.randn(3, 10)  # batch size为3,特征维度为10
targets = torch.LongTensor([1, 3, 0])  # 真实标签

# 训练模型
for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    
    print('Epoch {}, Loss: {}'.format(epoch, loss.item()))

在上面的示例中,我们定义了一个简单的神经网络模型Net,包含一个全连接层和一个softmax操作。然后,我们使用交叉熵损失函数和随机梯度下降优化器来训练模型,最小化预测概率分布与真实标签之间的交叉熵损失。在每个epoch中,我们计算损失并反向传播更新模型参数,直到达到指定的训练次数为止。

0