optimizer, milestones=[60, 120, 160], gamma=0.2, last_epoch=args.begin_epoch - 1
)
else:
raise ValueError('Unknown optimizer {}'.format(args.optimizer))
best_test_bpd = math.inf
if (args.resume is not None):
logger.info('Resuming model from {}'.format(args.resume))
with torch.no_grad():
x = torch.rand(1, *input_size[1:]).to(device)
model(x)
checkpt = torch.load(args.resume)
sd = {k: v for k, v in checkpt['state_dict'].items() if 'last_n_samples' not in k}
state = model.state_dict()
state.update(sd)
model.load_state_dict(state, strict=True)
ema.set(checkpt['ema'])
if 'optimizer_state_dict' in checkpt:
optimizer.load_state_dict(checkpt['optimizer_state_dict'])
# Manually move optimizer state to GPU
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.to(device)
del checkpt
del state
logger.info(optimizer)
fixed_z = standard_normal_sample([min(32, args.batchsize),
(im_dim + args.padding) * args.imagesize * args.imagesize]).to(device)
first_resblock=args.first_resblock,
learn_p=args.learn_p,
block_type=args.block,
)
model.to(device)
print('Initializing model.', flush=True)
with torch.no_grad():
x = torch.rand(1, *input_size[1:]).to(device)
model(x)
print('Restoring from checkpoint.', flush=True)
checkpt = torch.load(args.resume)
state = model.state_dict()
model.load_state_dict(checkpt['state_dict'], strict=True)
ema = utils.ExponentialMovingAverage(model)
ema.set(checkpt['ema'])
ema.swap()
print(model, flush=True)
model.eval()
print('Updating lipschitz.', flush=True)
update_lipschitz(model)
def visualize(model):
utils.makedirs('{}_imgs_t{}'.format(args.data, args.temp))