在这代码里面加入获取准确率和损失值并且绘制函数的代码import os import sys import json import torch import torch.nn as nn from torchvision import transforms, datasets, utils import matplotlib.pyplot as plt import numpy as np import torch.optim as optim from tqdm import tqdm from m

import os
import sys
import json
import torch
import torch.nn as nn
from torchvision import transforms, datasets, utils
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim
from tqdm import tqdm

def train(model, device, train_loader, optimizer, epoch):
model.train()
train_loss = 0
correct = 0
total = 0
progress_bar = tqdm(train_loader)

for batch_idx, (data, target) in enumerate(progress_bar):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = nn.CrossEntropyLoss()(output, target)
loss.backward()
optimizer.step()

train_loss += loss.item()
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()

progress_bar.set_description(
'Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader),
loss.item() / len(data)))

accuracy = 100. * correct / total
train_loss /= len(train_loader.dataset)

return accuracy, train_loss

def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
total = 0

with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += nn.CrossEntropyLoss()(output, target).item()
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()

test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / total

return accuracy, test_loss

def main():
# set up device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# set up hyperparameters
batch_size = 64
epochs = 10
lr = 0.01

# set up data loaders
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(
'../data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = datasets.MNIST(
'../data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=batch_size, shuffle=True)

# set up model
model = nn.Sequential(
nn.Conv2d(1, 10, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(10, 20, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2),
nn.Flatten(),
nn.Linear(320, 50),
nn.ReLU(),
nn.Linear(50, 10),
nn.LogSoftmax(dim=1)
).to(device)

# set up optimizer
optimizer = optim.SGD(model.parameters(), lr=lr)

# train and test the model
train_accuracies = []
train_losses = []
test_accuracies = []
test_losses = []

for epoch in range(1, epochs + 1):
train_accuracy, train_loss = train(
model, device, train_loader, optimizer, epoch)
test_accuracy, test_loss = test(model, device, test_loader)

train_accuracies.append(train_accuracy)
train_losses.append(train_loss)
test_accuracies.append(test_accuracy)
test_losses.append(test_loss)

print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))

# plot accuracy and loss curves
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs + 1), train_accuracies, label='Train')
plt.plot(range(1, epochs + 1), test_accuracies, label='Test')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(range(1, epochs + 1), train_losses, label='Train')
plt.plot(range(1, epochs + 1), test_losses, label='Test')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

plt.show()

if __name__ == '__main__':
main()