shuffle=False,
num_workers=args.workers
)
if args.model in ['VGG11_bn', 'ResNet18', 'DenseNet3_40', 'LeNet', 'VGG11']:
model = Network().construct(args.model, config)
else:
raise Exception('Unknown model: {}'.format())
model = model.to(device)
if args.dp:
model = nn.DataParallel(model, args.dp)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.mom, weight_decay=args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=args.lr_gamma)
acc_list = train(model,
optimizer,
scheduler,
dataloaders,
criterion,
device,
config.num_classes,
num_epochs=args.num_epochs,
args=args,
ckpt_dir=ckpt_dir,
dataset_sizes=dataset_sizes,
images_dir=images_dir
)
train_test_stats.append(acc_list)
# if args.model in ['VGG11', 'VGG11_bn', 'VGG13', 'VGG13_bn', 'VGG16', 'VGG16_bn', 'VGG19', 'VGG19_bn', 'ResNet18', 'DenseNet3_40', 'LeNet', 'MobileNet', 'FullyConnected']:
if args.model in model_choices:
model = Network().construct(args.model, config)
else:
raise Exception('Unknown model: {}'.format())
model = model.to(device)
if args.loss == 'ce':
criterion = nn.CrossEntropyLoss()
else:
raise Exception('Only cross entropy is allowed: {}'.format(args.loss))
if args.optimizer == 'sgd':
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay,
momentum=0.9)
elif args.optimizer == 'adam':
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
else:
raise Exception('Optimizer not allowed: {}'.format(args.optimizer))
scheduler = lr_scheduler.MultiStepLR(optimizer,
args.milestones,
gamma=args.step_gamma)
if args.resume is not None:
# if args.model in ['VGG11', 'VGG11_bn', 'VGG13', 'VGG13_bn', 'VGG16', 'VGG16_bn', 'VGG19', 'VGG19_bn', 'ResNet18', 'DenseNet3_40', 'LeNet', 'MobileNet', 'FullyConnected']:
if args.model in model_choices:
model = Network().construct(args.model, config)
else:
raise Exception('Unknown model: {}'.format())
model = model.to(device)
if args.loss == 'ce':
criterion = nn.CrossEntropyLoss()
else:
raise Exception('Only cross entropy is allowed: {}'.format(args.loss))
if args.optimizer == 'sgd':
optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, momentum=0.9)
elif args.optimizer == 'adam':
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
else:
raise Exception('Optimizer not allowed: {}'.format(args.optimizer))
scheduler = lr_scheduler.MultiStepLR(optimizer, args.milestones, gamma=args.step_gamma)
if args.resume is not None:
assert osp.exists(args.resume)
assert osp.isfile(args.resume)
ckpt = torch.load(args.resume)
assert 'model' in ckpt.keys()
model.load_state_dict(ckpt['model'], strict=True)
system = train(model,