def main(args):
trainloader, testloader = build_dataset(
"cifar10",
dataroot=args.dataroot,
batch_size=args.batch_size,
eval_batch_size=args.eval_batch_size,
num_workers=2,
)
if args.fname:
print("Loading model from %s" % args.fname)
model = torch.load(args.fname, map_location="cpu").cuda()
else:
model = build_model("ResNet18", num_classes=10)
criterion = torch.nn.CrossEntropyLoss()
eigenvals, eigenvecs = compute_hessian_eigenthings(
model,
testloader,
criterion,
args.num_eigenthings,
mode=args.mode,
# power_iter_steps=args.num_steps,
max_samples=args.max_samples,
# momentum=args.momentum,
full_dataset=args.full_dataset,
use_gpu=args.cuda,
)
print("Eigenvecs:")
print(eigenvecs)
print("Eigenvals:")
print(eigenvals)
def do_training(args):
trainloader, testloader = build_dataset(
args.dataset,
dataroot=args.dataroot,
batch_size=args.batch_size,
eval_batch_size=args.eval_batch_size,
num_workers=2)
model = build_model(args.arch, num_classes=num_classes(args.dataset))
if args.cuda:
model = torch.nn.DataParallel(model).cuda()
# Calculate total number of model parameters
num_params = sum(p.numel() for p in model.parameters())
track.metric(iteration=0, num_params=num_params)
if args.optimizer == 'sgd':
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
optimizer = EKFAC(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
eps=args.eps,
update_freq=args.update_freq)
criterion = torch.nn.CrossEntropyLoss()
best_acc = 0.0
for epoch in range(args.epochs):
track.debug("Starting epoch %d" % epoch)
args.lr = adjust_learning_rate(epoch, optimizer, args.lr,
args.schedule, args.gamma)
train_loss, train_acc = train(trainloader, model, criterion, optimizer,
epoch, args.cuda)
test_loss, test_acc = test(testloader, model, criterion, epoch,
args.cuda)
track.debug('Finished epoch %d... | train loss %.3f | train acc %.3f '
'| test loss %.3f | test acc %.3f' %
(epoch, train_loss, train_acc, test_loss, test_acc))
# Save model
model_fname = os.path.join(track.trial_dir(),
"model{}.ckpt".format(epoch))
torch.save(model, model_fname)
if test_acc > best_acc:
best_acc = test_acc
best_fname = os.path.join(track.trial_dir(), "best.ckpt")
track.debug("New best score! Saving model")
torch.save(model, best_fname)
def do_training(args):
trainloader, testloader = build_dataset(
args.dataset,
dataroot=args.dataroot,
batch_size=args.batch_size,
eval_batch_size=args.eval_batch_size,
num_workers=2)
model = build_model(args.arch, num_classes=num_classes(args.dataset))
if args.cuda:
model = torch.nn.DataParallel(model).cuda()
# Calculate total number of model parameters
num_params = sum(p.numel() for p in model.parameters())
track.metric(iteration=0, num_params=num_params)
num_chunks = max(1, args.batch_size // args.max_samples_per_gpu)
optimizer = LARS(params=model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
eta=args.eta,
max_epoch=args.epochs)
criterion = torch.nn.CrossEntropyLoss()
best_acc = 0.0
for epoch in range(args.epochs):
track.debug("Starting epoch %d" % epoch)
train_loss, train_acc = train(trainloader,
model,
criterion,
optimizer,
epoch,
args.cuda,
num_chunks=num_chunks)
test_loss, test_acc = test(testloader, model, criterion, epoch,
args.cuda)
track.debug('Finished epoch %d... | train loss %.3f | train acc %.3f '
'| test loss %.3f | test acc %.3f' %
(epoch, train_loss, train_acc, test_loss, test_acc))
# Save model
model_fname = os.path.join(track.trial_dir(),
"model{}.ckpt".format(epoch))
torch.save(model, model_fname)
if test_acc > best_acc:
best_acc = test_acc
best_fname = os.path.join(track.trial_dir(), "best.ckpt")
track.debug("New best score! Saving model")
torch.save(model, best_fname)
def main(args):
trainloader, testloader = build_dataset(
'cifar10',
dataroot=args.dataroot,
batch_size=args.batch_size,
eval_batch_size=args.eval_batch_size,
num_workers=2)
model = build_model('ResNet18', num_classes=10)
criterion = torch.nn.CrossEntropyLoss()
eigenvals, eigenvecs = compute_hessian_eigenthings(model,
testloader,
criterion,
args.num_eigenthings,
args.num_steps,
momentum=args.momentum,
use_gpu=args.cuda)
print("Eigenvecs:")
print(eigenvecs)
print("Eigenvals:")
print(eigenvals)
track.metric(iteration=0, eigenvals=eigenvals)
def test_optimizers():
print("Testing optimizers")
model = build_model('ResNet18', num_classes=10)
for name in optimizers:
_ = build_optimizer(name, params=model.parameters(), lr=.1)
def test_models():
print("Testing models")
for name in models:
_ = build_model(name, num_classes=10)