用Pytorch定义并训练一个简单的全连接网络

用Pytorch定义并训练一个简单的全连接网络，完整步骤如下：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# Step1: Define Fully connected network
class NN(nn.Module):
    def __init__(self, num_features, num_classes=10):
        super().__init__()
        self.fc1 =nn.Linear(num_features, 50)
        self.fc2 = nn.Linear(50, num_classes)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Set device & Hyperparameters
device = "cuda" if torch.cuda.is_available() else "cpu"

num_features = 784
num_classes = 10
learning_rate = 1e-3
batch_size = 64
num_epochs = 3

# Step2: Load data
train_dataset = datasets.MNIST(root="dataset/", train=True, transform=transforms.ToTensor(), download=True)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = datasets.MNIST(root="dataset/", train=False, transform=transforms.ToTensor(), download=True)
test_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)

# Step3: Initialize network
model = NN(num_features, num_classes).to(device)

# Step4: define Loss and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Step5: Train Network
for epoch in range(num_epochs):
    losses=[]
    for batch_idx, (data, targets) in enumerate(train_dataloader):
        data = data.to(device=device)
        targets = targets.to(device=device)
        data = data.reshape(data.shape[0], -1)

        # forward
        preds = model(data)
        loss = loss_fn(preds, targets)
        losses.append(loss)
        
        # backward
        optimizer.zero_grad()
        loss.backward()

        # GSD
        optimizer.step()
    print(f"Epoch:{epoch}, loss is {sum(losses)/len(losses)}.")

# Step6: Chekc accuracy on test dataset
num_correct = 0
num_samples = 0
model.eval()

with torch.no_grad():
    for data, targets in test_dataloader:
        data = data.to(device)
        targets = targets.to(device)
        data = data.reshape(data.shape[0], -1)

        preds = model(data)
        _, results = preds.max(1)
        # print(preds.shape, results.shape, targets.shape)
        num_correct += (results == targets).sum()
        num_samples += results.size(0)
    print(f"The accuracy on test dataset is : {float(num_correct)/float(num_samples)*100:.2f}%")

运行结果：

Epoch:0, loss is 0.42042621970176697.
Epoch:1, loss is 0.20832164585590363.
Epoch:2, loss is 0.15770433843135834.
The accuracy is : 96.20%