PyTorch实现多分类任务

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

'''定义模型'''
class SimpleModel(nn.Module):
    '''
    方便理解，这里只定义了一层网络
    input_size: 输入维度（这里表示每个样本的特征数量）
    num_classes: 输出维度（这里表示类别数量）
    '''
    def __init__(self, input_size, num_classes) -> None:
        super().__init__()
        self.linear = nn.Linear(input_size, num_classes)
        
    def forward(self, x):
        x = self.linear(x)
        # 此处没有使用激活函数  
        return x    
    
'''构造模拟数据'''
input_data = torch.randn(5, 10)  # 5个样本，每个样本有10个特征  
labels = torch.tensor([0, 1, 2, 0, 1])  # 5个样本，每个样本对应一个类别标签（三分类：0，1，2）

'''实例化模型'''
model = SimpleModel(input_size=10, num_classes=3)

'''定义损失函数'''
criterion = nn.CrossEntropyLoss()

'''定义优化器'''
optimizer = optim.SGD(model.parameters(), lr=0.1)   

epoch = 1000
for i in range(epoch):
    '''前向传播'''
    y = model(input_data)

    '''计算损失'''
    # 此处，criterion会自动将y进行softmax处理，所以不需要显式地在模型中定义softmax  
    # 同时，labels不需要进行one-hot编码，因为criterion会自动完成这一操作    
    loss = criterion(y, labels)
    print(f'Epoch {i+1}/{epoch}, Loss: {loss}')   
    
    '''反向传播 和 更新参数'''
    # 梯度清零
    optimizer.zero_grad()
    # 计算梯度
    loss.backward()
    # 更新参数
    optimizer.step()