trainloader, testloader = utils.mnist_dataloaders(args.data_path,
args.batch_size,
args.use_gpu)
model = CapsNet(n_conv_channel=256,
n_primary_caps=8,
primary_cap_size=1152,
output_unit_size=16,
n_routing_iter=3)
# load state from past runs
if args.load_checkpoint != '':
model.load_state_dict(torch.load(args.load_checkpoint))
# move to GPU
model = model.cuda() if args.use_gpu else model
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# setup decoder for training
decoder = Decoder()
decoder = decoder.cuda() if args.use_gpu else decoder
decoder_optim = torch.optim.Adam(decoder.parameters(), lr=0.001)
# use decaying learning rate
scheduler = lr_scheduler.ExponentialLR(decoder_optim, 0.5)
for epoch in range(1, args.epochs + 1):
train(epoch, model, trainloader, optimizer, decoder, decoder_optim)
test(epoch, model, testloader)
scheduler.step()
loss = capsule_net.loss(data, output, target, reconstructions)
test_loss += loss.data[0]
correct += sum(
np.argmax(masked.data.cpu().numpy(), 1) == np.argmax(
target.data.cpu().numpy(), 1))
tqdm.write("Epoch: [{}/{}], test accuracy: {:.6f}, loss: {:.6f}".format(
epoch, N_EPOCHS, correct / len(test_loader.dataset),
test_loss / len(test_loader)))
if __name__ == '__main__':
torch.manual_seed(1)
dataset = 'cifar10'
# dataset = 'mnist'
config = Config(dataset)
mnist = Dataset(dataset, BATCH_SIZE)
capsule_net = CapsNet(config)
capsule_net = torch.nn.DataParallel(capsule_net)
if USE_CUDA:
capsule_net = capsule_net.cuda()
capsule_net = capsule_net.module
optimizer = torch.optim.Adam(capsule_net.parameters())
for e in range(1, N_EPOCHS + 1):
train(capsule_net, optimizer, mnist.train_loader, e)
test(capsule_net, mnist.test_loader, e)
drop_last=False)
sup_iterator = train_loader.__iter__()
test_iterator = test_loader.__iter__()
imgs, labels = sup_iterator.next()
sup_iterator = train_loader.__iter__()
"""
Setup model, load it to CUDA and make JIT compilation
"""
#imgs = imgs[:2]
stat = []
model = CapsNet(args,
len(train_dataset) // (2 * args.batch_size) + 3, stat)
use_cuda = not args.disable_cuda and torch.cuda.is_available()
if use_cuda:
model.cuda()
imgs = imgs.cuda()
if args.jit:
model = torch.jit.trace(model, (imgs), check_inputs=[(imgs)])
else:
model(imgs)
print("# model parameters:",
sum(param.numel() for param in model.parameters()))
"""
Construct optimizer, scheduler, and loss functions
"""
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr if args.lr else 1e-3,
betas=(0.9, 0.999))
#optimizer = torch.optim.SGD(model.parameters(), lr=args.lr if args.lr else 1e-3, momentum=0.9, weight_decay=0.0005)
