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'])
elif 'FlowNet2-CS/' in args.caffe_model:
model = models.FlowNet2CS(args)
parse_flownetc(model.flownetc.modules(), weights, biases)
parse_flownets(model.flownets_1.modules(),
weights,
biases,
param_prefix='net2_')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(
state,
os.path.join(args.flownet2_pytorch, 'FlowNet2-CS_checkpoint.pth.tar'))
elif 'FlowNet2-CSS/' in args.caffe_model:
model = models.FlowNet2CSS(args)
parse_flownetc(model.flownetc.modules(), weights, biases)
parse_flownets(model.flownets_1.modules(),
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_checkpoint.pth.tar'))
elif 'FlowNet2-CS/' in opt.caffe_model:
model = models.FlowNet2CS(opt)
parse_flownetc(model.flownetc.modules(), weights, biases)
parse_flownets(model.flownets_1.modules(),
weights,
biases,
param_prefix='net2_')
state = {'epoch': 0, 'state_dict': model.state_dict(), 'best_EPE': 1e10}
torch.save(
state, os.path.join(opt.weights_path,
'FlowNet2-CS_checkpoint.pth.tar'))
elif 'FlowNet2-CSS/' in opt.caffe_model:
model = models.FlowNet2CSS(opt)
parse_flownetc(model.flownetc.modules(), weights, biases)
parse_flownets(model.flownets_1.modules(),
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_checkpoint.pth.tar'))