Use o PyTorch para criar um modelo de classificação CNN simples para classificar MNIST

Importar dados do archivision

import torchvision 
from torchvision import transforms 

transform = transforms.Compose([
    transforms.ToTensor()
])

mnist_train_dataset = torchvision.datasets.MNIST(
    root='./', train=True, transform=transform,download=True)
minist_test_dataset = torchvision.datasets.MNIST(
   root='./', train=False, transform=transform, download=True)

Construir DataLoader

from torch.utils.data import DataLoader

batch_size = 64 
train_dl = DataLoader(mnist_train_dataset, batch_size, shuffle=True)
test_dl = DataLoader(mnist_test_dataset, batch_size, shuffle=True)

Construir um modelo CNN

import torch.nn as nn 

model = nn.Sequential()
model.add_module('conv1', nn.Conv2d(in_channels=1,out_channels=32, kernel_size=5, padding=2))
model.add_module('relu1', nn.ReLU())
model.add_module('pool1', nn.MaxPool2d(kernel_size=2))
model.add_module('conv2', nn.Conv2d(in_channels=32,out_channels=64, kernel_size=5, padding=2))
model.add_module('relu2', nn.ReLU())
model.add_module('pool2', nn.MaxPool2d(kernel_size=2))
model.add_module('flatten', nn.Flatten())

x = torch.ones((4, 1, 28, 28))
print(model(x).shape)

output:
torch.Size([4, 3136])

Pool=2 representa o mesmo preenchimento, ou seja, a altura e a largura do mapa de recursos antes e depois de rolar permanecem inalteradas.

A entrada do pytorch é o modo NCHW, N: número do lote. C: número do canal. H: altura do mapa de recursos. W: largura do mapa de recursos.

(O modo de entrada do tensorflow é NHWC)

Em seguida, a camada totalmente conectada: 

model.add_module('fc1', nn.Linear(3136, 1024))
model.add_module('relu3', nn.ReLU())
model.add_module('dropout', nn.Dropout(p=0.5))
model.add_module('fc2', nn.Linear(1024, 10))

Definir perda e otimizador

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

 nn.CrossEntropyLoss() contém a função softmax, então a última camada não adiciona a camada softmax ao definir o modelo

modelo de treinamento

def train(model, num_epochs, trin_dl, valid_dl): 
    loss_hist_train = [0] * num_epochs 
    accuracy_hist_train = [0] * num_epochs
    loss_hist_test = [0] * num_epochs 
    accuracy_hist_test = [0] * num_epochs 
    for epoch in range(num_epochs):
        model.train()
        for x_batch, y_batch in train_dl:
            
            pred = model(x_batch)
            loss = loss_fn(pred, y_batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            
            loss_hist_train[epoch] += loss.item()*y_batch.size(0)
            is_crrect = (torch.argmax(pred, dim=1) == y_batch).float()
            accuracy_hist_train[epoch] += is_crrect.sum()
        loss_hist_train[epoch] /= len(train_dl.dataset)
        accuracy_hist_train[epoch] /= len(train_dl.dataset)
        
        
        model.eval()
        with torch.no_grad():
            for x_batch, y_batch in test_dl:
                pred = model(x_batch)
                loss = loss_fn(pred, y_batch)
                loss_hist_test[epoch] += loss.item()*y_batch.size(0)
                is_correct = (torch.argmax(pred, dim=1) == y_batch).float()
                accuracy_hist_test[epoch] += is_correct.sum()
                
        loss_hist_test[epoch] /= len(test_dl.dataset)
        accuracy_hist_test[epoch] /= len(test_dl.dataset)
            
        print(f'Epoch {epoch+1} accuracy: {accuracy_hist_train[epoch]:.4f}'
                  f'test_accuracy: {accuracy_hist_test[epoch]:.4f}')
    return loss_hist_train, loss_hist_test, accuracy_hist_train, accuracy_hist_test

num_epochs= 20 
hist = train(model, num_epochs, train_dl, test_dl)


Epoch 1 accuracy: 0.9868test_accuracy: 0.9925
Epoch 2 accuracy: 0.9902test_accuracy: 0.9914

参考自:   Machine Learning with PyTorch and Scikit-Learn  Book by Sebastian Raschka