def __init__(
self,
args_flow,
model_name="FlowNet2C",
restore_path="flow_inference/models/FlowNet2-C_checkpoint.pth.tar",
train_flow=False,
resume=None):
super(Inference_flow, self).__init__()
args_flow.rgb_max = 255.0
args_flow.fp16 = False
args_flow.grads = {}
if model_name == 'FlowNet2':
self.model = models.FlowNet2(args_flow).cuda()
elif model_name == 'FlowNet2C':
self.model = models.FlowNet2C(args_flow).cuda()
elif model_name == 'FlowNet2S':
self.model = models.FlowNet2S(args_flow).cuda()
elif model_name == 'FlowNet2SD':
self.model = models.FlowNet2SD(args_flow).cuda()
elif model_name == 'FlowNet2CS':
self.model = models.FlowNet2CS(args_flow).cuda()
elif model_name == 'FlowNet2CSS':
self.model = models.FlowNet2CSS(args_flow).cuda()
else:
assert False, "No such model %s" % (model_name)
print("loading %s pretrained model..." % (model_name))
if train_flow:
self.model.train()
else:
self.model.eval()
if resume is not None:
self.model.load_state_dict(torch.load(resume)['flow'])
else:
self.model.load_state_dict(torch.load(restore_path)['state_dict'])
weights,
biases,
param_prefix='net2_')
parse_flownets(model.flownets_2.modules(),
weights,
biases,
param_prefix='net3_')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(
state,
os.path.join(args.flownet2_pytorch,
'FlowNet2-CSS-ft-sd_checkpoint.pth.tar'))
elif 'FlowNet2-S/' in args.caffe_model:
model = models.FlowNet2S(args)
parse_flownetsonly(model.modules(), weights, biases, param_prefix='')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(
state,
os.path.join(args.flownet2_pytorch, 'FlowNet2-S_checkpoint.pth.tar'))
elif 'FlowNet2-SD/' in args.caffe_model:
model = models.FlowNet2SD(args)
parse_flownetsd(model.modules(), weights, biases, param_prefix='')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(
state,
weights,
biases,
param_prefix='net2_')
parse_flownets(model.flownets_2.modules(),
weights,
biases,
param_prefix='net3_')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(
state,
os.path.join(opt.weights_path,
'FlowNet2-CSS-ft-sd_checkpoint.pth.tar'))
elif 'FlowNet2-S/' in opt.caffe_model:
model = models.FlowNet2S(opt)
parse_flownetsonly(model.modules(), weights, biases, param_prefix='')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(state,
os.path.join(opt.weights_path, 'FlowNet2-S_checkpoint.pth.tar'))
elif 'FlowNet2-SD/' in opt.caffe_model:
model = models.FlowNet2SD(opt)
parse_flownetsd(model.modules(), weights, biases, param_prefix='')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(
state, os.path.join(opt.weights_path,
'FlowNet2-SD_checkpoint.pth.tar'))