创建一个基本的pytorch框架

# 导入需要的库
import torch
import torch.nn as nn
import  torch.optim as optim
from   torchvision import datasets, transforms

1. 读取数据集

一种是导入python自带的数据

train_db = datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ]))
train_loader = torch.utils.data.DataLoader(train_db,
    batch_size=batch_size, shuffle=True)

test_db = datasets.MNIST('../data', train=False,
                         transform=transforms.Compose([
                             transforms.ToTensor(),
                             transforms.Normalize((0.1307,), (0.3081,))
                         ]))
test_loader = torch.utils.data.DataLoader(test_db,
    batch_size=batch_size, shuffle=True)
print('train:', len(train_db), 'test:', len(test_db))

train_db, val_db = torch.utils.data.random_split(train_db, [50000, 10000])
print('db1:', len(train_db), 'db2:', len(val_db))

train_loader = torch.utils.data.DataLoader(train_db,
    batch_size=batch_size, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_db,
    batch_size=batch_size, shuffle=True)

2. 定义网络

class MLP(nn.Module):

    def __init__(self):
        super(MLP, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(784, 200),
            nn.LeakyReLU(inplace=True),
            nn.Linear(200, 200),
            nn.LeakyReLU(inplace=True),
            nn.Linear(200, 10),
            nn.LeakyReLU(inplace=True),
        )

    def forward(self, x):
        x = self.model(x)

        return x
      
device = torch.device('cuda:0') # 使用GPU训练
net = MLP().to(device)  
print(net)

3. 训练过程

batch_size=200  # 每个batch中训练样本的数量
learning_rate=0.01  # 学习率
epochs=10  # 训练的次数，权重迭代的次数
optimizer = optim.SGD(net.parameters(), lr=learning_rate)  # 优化器
criteon = nn.CrossEntropyLoss().to(device)   #loss

for epoch in range(epochs):

    for batch_idx, (data, target) in enumerate(train_loader):
    # data.size [B, N, H, W]
        data = data.view(-1, 28*28)
        data, target = data.to(device), target.cuda()

        logits = net(data)
        loss = criteon(logits, target)

        optimizer.zero_grad()
        loss.backward()
        # print(w1.grad.norm(), w2.grad.norm())
        optimizer.step()

        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))
     
    test_loss = 0
    correct = 0
    for data, target in val_loader:
        data = data.view(-1, 28 * 28)
        data, target = data.to(device), target.cuda()
        logits = net(data)
        test_loss += criteon(logits, target).item()

        pred = logits.data.max(1)[1]
        correct += pred.eq(target.data).sum()

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

4. 测试过程

test_loss = 0
correct = 0
for data, target in test_loader:
    data = data.view(-1, 28 * 28)
    data, target = data.to(device), target.cuda()
    logits = net(data)
    test_loss += criteon(logits, target).item()

    pred = logits.data.max(1)[1]
    correct += pred.eq(target.data).sum()

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

Train Epoch: 0 [0/50000 (0%)]	Loss: 2.308843
Train Epoch: 0 [20000/50000 (40%)]	Loss: 2.074609
Train Epoch: 0 [40000/50000 (80%)]	Loss: 1.860805

VAL set: Average loss: 0.0083, Accuracy: 5077/10000 (51%)

Train Epoch: 1 [0/50000 (0%)]	Loss: 1.779289
Train Epoch: 1 [20000/50000 (40%)]	Loss: 1.278288
Train Epoch: 1 [40000/50000 (80%)]	Loss: 1.209769

VAL set: Average loss: 0.0052, Accuracy: 7120/10000 (71%)

Train Epoch: 2 [0/50000 (0%)]	Loss: 1.106976
Train Epoch: 2 [20000/50000 (40%)]	Loss: 0.663975
Train Epoch: 2 [40000/50000 (80%)]	Loss: 0.684854

VAL set: Average loss: 0.0034, Accuracy: 8020/10000 (80%)

Train Epoch: 3 [0/50000 (0%)]	Loss: 0.753274
Train Epoch: 3 [20000/50000 (40%)]	Loss: 0.670728
Train Epoch: 3 [40000/50000 (80%)]	Loss: 0.530107

VAL set: Average loss: 0.0022, Accuracy: 8856/10000 (89%)

Train Epoch: 4 [0/50000 (0%)]	Loss: 0.365646
Train Epoch: 4 [20000/50000 (40%)]	Loss: 0.370881
Train Epoch: 4 [40000/50000 (80%)]	Loss: 0.475091

VAL set: Average loss: 0.0018, Accuracy: 8946/10000 (89%)

Train Epoch: 5 [0/50000 (0%)]	Loss: 0.290509
Train Epoch: 5 [20000/50000 (40%)]	Loss: 0.313986
Train Epoch: 5 [40000/50000 (80%)]	Loss: 0.336088

VAL set: Average loss: 0.0017, Accuracy: 9024/10000 (90%)

Train Epoch: 6 [0/50000 (0%)]	Loss: 0.308336
Train Epoch: 6 [20000/50000 (40%)]	Loss: 0.311062
Train Epoch: 6 [40000/50000 (80%)]	Loss: 0.339352

VAL set: Average loss: 0.0016, Accuracy: 9092/10000 (91%)

Train Epoch: 7 [0/50000 (0%)]	Loss: 0.265708
Train Epoch: 7 [20000/50000 (40%)]	Loss: 0.322460
Train Epoch: 7 [40000/50000 (80%)]	Loss: 0.226619

VAL set: Average loss: 0.0015, Accuracy: 9143/10000 (91%)

Train Epoch: 8 [0/50000 (0%)]	Loss: 0.286767
Train Epoch: 8 [20000/50000 (40%)]	Loss: 0.186286
Train Epoch: 8 [40000/50000 (80%)]	Loss: 0.260167

VAL set: Average loss: 0.0014, Accuracy: 9173/10000 (92%)

Train Epoch: 9 [0/50000 (0%)]	Loss: 0.275809
Train Epoch: 9 [20000/50000 (40%)]	Loss: 0.204801
Train Epoch: 9 [40000/50000 (80%)]	Loss: 0.196037

VAL set: Average loss: 0.0013, Accuracy: 9213/10000 (92%)


Test set: Average loss: 0.0012, Accuracy: 9300/10000 (93%)


Process finished with exit code 0