week2

目标：

尝试完成一个多分类任务的训练:一个随机向量，哪一维数字最大就属于第几类。

内容：

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

# 1. 固定随机种子
torch.manual_seed(42)

# 2. 定义参数
input_dim = 5
num_classes = 5
num_samples = 10000
epochs = 100
batch_size = 128

# 3. 生成训练数据
X_train = torch.rand(num_samples, input_dim)

# 找出每个向量中最大值的位置
y_train = torch.argmax(X_train, dim=1)

# 生成独立测试数据
X_test = torch.rand(2000, input_dim)
y_test = torch.argmax(X_test, dim=1)

# 4. 定义神经网络
class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.fc = nn.Linear(5, 5)

    def forward(self, x):
        return self.fc(x)

model = Model()

# 5. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 6. 训练模型
for epoch in range(epochs):
    model.train()

    indices = torch.randperm(num_samples)
    total_loss = 0
    total_correct = 0

    for start in range(0, num_samples, batch_size):
        idx = indices[start:start + batch_size]

        X_batch = X_train[idx]
        y_batch = y_train[idx]

        # 前向传播
        outputs = model(X_batch)

        # 计算损失
        loss = criterion(outputs, y_batch)

        # 梯度清零
        optimizer.zero_grad()

        # 反向传播
        loss.backward()

        # 更新参数
        optimizer.step()

        total_loss += loss.item() * len(idx)

        preds = torch.argmax(outputs, dim=1)
        total_correct += (preds == y_batch).sum().item()

    if (epoch + 1) % 10 == 0:
        print(
            f"Epoch {epoch + 1}, "
            f"Loss: {total_loss / num_samples:.4f}, "
            f"Accuracy: {total_correct / num_samples:.2%}"
        )

# 7. 测试模型
model.eval()

with torch.no_grad():
    outputs = model(X_test)
    predictions = torch.argmax(outputs, dim=1)

    accuracy = (predictions == y_test).float().mean()

    print(f"\n测试准确率: {accuracy.item():.2%}")

# 8. 预测新的随机向量
x = torch.tensor([
    [0.12, 0.35, 0.91, 0.43, 0.28]
])

with torch.no_grad():
    output = model(x)
    prediction = torch.argmax(output, dim=1)

print(f"预测类别：第 {prediction.item() + 1} 类")
print(f"真实类别：第 {torch.argmax(x).item() + 1} 类")