activation_fn=args.act,
fc_end=args.fc_end,
fc_idim=args.fc_idim,
n_exact_terms=args.n_exact_terms,
preact=args.preact,
neumann_grad=args.neumann_grad,
grad_in_forward=args.mem_eff,
first_resblock=args.first_resblock,
learn_p=args.learn_p,
classification=args.task in ['classification', 'hybrid'],
classification_hdim=args.cdim,
n_classes=n_classes,
block_type=args.block,
)
model.to(device)
ema = utils.ExponentialMovingAverage(model)
def parallelize(model):
return torch.nn.DataParallel(model)
logger.info(model)
logger.info('EMA: {}'.format(ema))
# Optimization
def tensor_in(t, a):
for a_ in a:
if t is a_:
activation_fn=args.act,
fc_end=args.fc_end,
fc_idim=args.fc_idim,
n_exact_terms=args.n_exact_terms,
preact=args.preact,
neumann_grad=args.neumann_grad,
grad_in_forward=args.mem_eff,
first_resblock=args.first_resblock,
learn_p=args.learn_p,
classification=args.task in ['classification', 'hybrid'],
classification_hdim=args.cdim,
n_classes=n_classes,
block_type=args.block,
)
model = model.to(device)
# load model
with torch.no_grad():
x = torch.rand(1, *input_size[1:]).to(device)
model(x)
model.load_state_dict(checkpoint['state_dict'], strict=True)
model.eval()
def postprocess(x):
x = torch.clamp(x, 0.0, 1.0)
x = x * 255 #2**args.nbits
x += 0.5
return torch.clamp(x, 0, 255).byte()