def load(args):
add_args(args, sub_args)
net = VGG(10)
net.build_gate(VIB)
train_loader, test_loader = get_CIFAR10(args.batch_size)
base_params = []
gate_params = []
for name, param in net.named_parameters():
if 'gate' in name:
gate_params.append(param)
else:
base_params.append(param)
optimizer = optim.Adam([{
'params': gate_params,
'lr': 1e-2
}, {
'params': base_params,
'lr': 1e-3,
'weight_decay': 1e-4
}])
scheduler = optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[int(r * args.num_epochs) for r in [.5, .8]],
gamma=0.1)
return net, train_loader, test_loader, optimizer, scheduler
def load(args):
add_args(args, sub_args)
net = VGG(10)
train_loader, test_loader = get_CIFAR10(args.batch_size)
optimizer = optim.Adam(net.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[int(r * args.num_epochs) for r in [0.5, 0.8]],
gamma=0.1)
return net, train_loader, test_loader, optimizer, scheduler
def load(args):
add_args(args, sub_args)
net = VGG(10)
net.build_gate(L0Reg, {
'weight_decay': 1e-4,
'lamb': args.lamb,
'droprate_init': 0.2
})
train_loader, test_loader = get_CIFAR10(args.batch_size)
optimizer = optim.Adam(net.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[int(r * args.num_epochs) for r in [.5, .8]],
gamma=0.1)
return net, train_loader, test_loader, optimizer, scheduler