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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145
|
import torch from torch import nn, optim
from torch.autograd import Variable from torch.utils.data import DataLoader from torchvision import transforms from torchvision import datasets
batch_size = 128 learning_rate = 1e-2 num_epoches = 20
train_dataset = datasets.MNIST( root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST( root='./data', train=False, transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
class Cnn(nn.Module): def __init__(self, in_dim, n_class): super(Cnn, self).__init__() self.conv = nn.Sequential( nn.Conv2d(in_dim, 6, 5, stride=1, padding=2), nn.ReLU(True), nn.MaxPool2d(2, 2), nn.Conv2d(6, 16, 5, stride=1, padding=0), nn.ReLU(True), nn.MaxPool2d(2, 2))
self.fc = nn.Sequential( nn.Linear(400, 120), nn.Linear(120, 84), nn.Linear(84, n_class))
def forward(self, x): out = self.conv(x) out = out.view(out.size(0), -1) out = self.fc(out) return out
model = Cnn(1, 10) use_gpu = torch.cuda.is_available() if use_gpu: model = model.cuda()
criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=learning_rate)
for epoch in range(num_epoches): print('epoch {}'.format(epoch + 1)) print('*' * 10) running_loss = 0.0 running_acc = 0.0 for i, data in enumerate(train_loader, 1): img, label = data if use_gpu: img = img.cuda() label = label.cuda() img = Variable(img) label = Variable(label) out = model(img) loss = criterion(out, label) running_loss += loss.item() * label.size(0) _, pred = torch.max(out, 1) num_correct = (pred == label).sum() accuracy = (pred == label).float().mean() running_acc += num_correct.item() optimizer.zero_grad() loss.backward() optimizer.step() """ # ========================= Log ====================== step = epoch * len(train_loader) + i # (1) Log the scalar values info = {'loss': loss.data[0], 'accuracy': accuracy.data[0]}
for tag, value in info.items(): logger.scalar_summary(tag, value, step)
# (2) Log values and gradients of the parameters (histogram) for tag, value in model.named_parameters(): tag = tag.replace('.', '/') logger.histo_summary(tag, to_np(value), step) logger.histo_summary(tag + '/grad', to_np(value.grad), step)
# (3) Log the images info = {'images': to_np(img.view(-1, 28, 28)[:10])}
for tag, images in info.items(): logger.image_summary(tag, images, step) if i % 300 == 0: print('[{}/{}] Loss: {:.6f}, Acc: {:.6f}'.format( epoch + 1, num_epoches, running_loss / (batch_size * i), running_acc / (batch_size * i))) """ print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format( epoch + 1, running_loss / (len(train_dataset)), running_acc / (len(train_dataset)))) model.eval() eval_loss = 0 eval_acc = 0 for data in test_loader: img, label = data if use_gpu: img = Variable(img, volatile=True).cuda() label = Variable(label, volatile=True).cuda() else: img = Variable(img, volatile=True) label = Variable(label, volatile=True) out = model(img) loss = criterion(out, label) eval_loss += loss.item() * label.size(0) _, pred = torch.max(out, 1) num_correct = (pred == label).sum() eval_acc += num_correct.item() print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len( test_dataset)), eval_acc / (len(test_dataset)))) print()
torch.save(model.state_dict(), './cnn.pth')
|