def test(args, model, epoch=None, grad=False): model.eval() if args.dataset == 'mnist': _, test_loader = datagen.load_mnist(args) elif args.dataset == 'fashion_mnist': _, test_loader = datagen.load_fashion_mnist(args) test_loss = 0 correct = 0. criterion = nn.CrossEntropyLoss() for data, target in test_loader: data, target = data.cuda(), target.cuda() output = model(data) if grad is False: test_loss += criterion(output, target).item() # sum up batch loss else: test_loss += criterion(output, target) pred = output.data.max( 1, keepdim=True)[1] # get the index of the max log-probability correct += pred.eq(target.data.view_as(pred)).long().cpu().sum() test_loss /= len(test_loader.dataset) acc = (correct.float() / len(test_loader.dataset)).item() print(acc) if epoch: print('Average loss: {}, Accuracy: {}/{} ({}%)'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset))) return acc, test_loss
def train(args, model, grad=False): if args.dataset == 'mnist': train_loader, _ = datagen.load_mnist(args) elif args.dataset == 'fashion_mnist': train_loader, _ = datagen.load_fashion_mnist(args) train_loss, train_acc = 0., 0. criterion = nn.CrossEntropyLoss() if args.ft: for child in list(model.children())[:2]: print('removing {}'.format(child)) for param in child.parameters(): param.requires_grad = False optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3) for epoch in range(args.epochs): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.cuda(), target.cuda() optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() acc, loss = test(args, model, epoch) return acc, loss
def load_data(args): if args.dataset == 'mnist': return datagen.load_mnist(args) if args.dataset == 'cifar': return datagen.load_cifar(args) if args.dataset == 'fmnist': return datagen.load_fashion_mnist(args) if args.dataset == 'cifar_hidden': class_list = [0] ## just load class 0 return datagen.load_cifar_hidden(args, class_list) else: print ('Dataset not specified correctly') print ('choose --dataset <mnist, fmnist, cifar, cifar_hidden>')