def modelDeploy(args, model, optimizer, scheduler, logger):
if args.num_gpus >= 1:
from torch.nn.parallel import DataParallel
model = DataParallel(model)
model = model.cuda()
if torch.backends.cudnn.is_available():
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
cudnn.deterministic = True
trainData = {'epoch': 0, 'loss': [], 'miou': [], 'val': [], 'bestMiou': 0}
if args.resume:
if os.path.isfile(args.resume):
logger.info("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume,
map_location=torch.device('cpu'))
# model&optimizer
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
# stop point
trainData = checkpoint['trainData']
for i in range(trainData['epoch']):
scheduler.step()
# print(trainData)
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, trainData['epoch']))
else:
logger.error("=> no checkpoint found at '{}'".format(args.resume))
assert False, "=> no checkpoint found at '{}'".format(args.resume)
if args.finetune:
if os.path.isfile(args.finetune):
logger.info("=> finetuning checkpoint '{}'".format(args.finetune))
state_all = torch.load(args.finetune, map_location='cpu')['model']
state_clip = {} # only use backbone parameters
# print(model.state_dict().keys())
for k, v in state_all.items():
state_clip[k] = v
# print(state_clip.keys())
model.load_state_dict(state_clip, strict=False)
else:
logger.warning("finetune is not a file.")
pass
if args.freeze_bn:
logger.warning('Freezing batch normalization layers')
for m in model.modules():
if isinstance(m, nn.BatchNorm2d):
m.eval()
m.weight.requires_grad = False
m.bias.requires_grad = False
return model, trainData
