npu-smi set -t pwm-mode -d value
# value 0为手动模式 1为自动模式
npu-smi set -t pwm-duty-ratio -d value
# value 直接调100就可以了，虽然吵了一点但比过热死机强。

# 在开发板或 ModelArts 终端里
git clone https://github.com/pytorch/examples.git
cd examples/mnist
# 安装依赖（如果你环境里没有）
pip install torchvision matplotlib tensorboard -i https://pypi.tuna.tsinghua.edu.cn/simple

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# 可视化几个样本
examples = iter(train_loader)
samples, labels = next(examples)
plt.figure(figsize=(10, 4))
for i in range(6):
    plt.subplot(2, 3, i+1)
    plt.imshow(samples[i][0], cmap='gray')
    plt.title(f"Label: {labels[i]}")
    plt.axis('off')
plt.show()

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

model = Net().to(device)
print(model)

optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.NLLLoss()

epochs = 3
train_losses = []

for epoch in range(epochs):
    model.train()
    running_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    
    avg_loss = running_loss / len(train_loader)
    train_losses.append(avg_loss)
    
    # 验证
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    acc = 100. * correct / len(test_dataset)
    print(f'Epoch {epoch+1}: Loss={avg_loss:.4f}, Test Acc={acc:.2f}%')

# 画 Loss 曲线
plt.plot(range(1, epochs+1), train_losses, marker='o')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss Curve')
plt.grid()
plt.show()

model.eval()
test_iter = iter(test_loader)
data, target = next(test_iter)
data, target = data.to(device), target.to(device)
output = model(data)
pred = output.argmax(dim=1)

# 找出前 6 个预测错的样本
wrong_idx = (pred != target).nonzero(as_tuple=True)[0][:6]
plt.figure(figsize=(12, 8))
for i, idx in enumerate(wrong_idx):
    plt.subplot(2, 3, i+1)
    plt.imshow(data[idx].cpu().squeeze(), cmap='gray')
    plt.title(f'True: {target[idx]}, Pred: {pred[idx]}')
    plt.axis('off')
plt.show()