1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81
| import torch from torch import nn import torchvision.datasets as datasets import torchvision.transforms as transforms
torch.manual_seed(2019)
EPOCH = 1 BATCH_SIZE = 32 INIT_LR = 1e-3 DOWNLOAD_MNIST = True
train_data = datasets.MNIST(root='mnist', train=True, transform=transforms.ToTensor(), download=DOWNLOAD_MNIST) train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True) test_data = datasets.MNIST(root='mnist', train=False) test_x = test_data.test_data.type(torch.FloatTensor)[:2000] / 255. test_x.unsqueeze_(1) test_y = test_data.test_labels.numpy()[:2000]
class ConvNet(nn.Module): def __init__(self): super(ConvNet, self).__init__() self.conv = nn.Sequential( nn.Conv2d(1, 32, 5), nn.MaxPool2d(2), nn.ReLU(True), nn.Conv2d(32, 64, 5), nn.Dropout2d(), nn.MaxPool2d(2), nn.ReLU(True) )
self.linear = nn.Sequential( nn.Linear(4 * 4 * 64, 128), nn.ReLU(True), nn.Dropout2d(), nn.Linear(128, 10), nn.Softmax(1) )
def forward(self, x): x = self.conv(x) x = x.view(-1, 4 * 4 * 64) out = self.linear(x) return out
model = ConvNet() print(model)
optimizer = torch.optim.Adam(model.parameters(), lr=INIT_LR) loss_func = nn.CrossEntropyLoss()
for epoch in range(EPOCH): for index, (b_x, b_y) in enumerate(train_loader): model.train() output = model(b_x) loss = loss_func(output, b_y) optimizer.zero_grad() loss.backward() optimizer.step()
if index % 50 == 0: model.eval() prediction = model(test_x) pred_y = torch.max(prediction, 1)[1].data.numpy() accuracy = (pred_y == test_y).sum() / float(test_y.size) print(f'Epoch: [{index}/{epoch}]', f'| train loss: {loss.item()}', f'| test accuracy: {accuracy}')
model.eval() prediction = model(test_x[:20]) pred_y = torch.max(prediction, 1)[1].data.numpy() print(pred_y, 'prediction number') print(test_y[:20], 'real number')
|