def load_branch(self):
opt = self.opt
if opt.rgb_model != '':
print('create rgb model')
self.rgb_model_backbone, self.rgb_model_branch = create_inference_model(
opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test)
self.rgb_model_backbone, self.rgb_model_branch = load_inference_model(
self.rgb_model_backbone, self.rgb_model_branch, opt.rgb_model)
self.rgb_model_branch = DataParallel(
self.rgb_model_branch,
device_ids=opt.gpus,
chunk_sizes=opt.chunk_sizes).to(opt.device)
self.rgb_model_branch.eval()
if opt.flow_model != '':
print('create flow model')
self.flow_model_backbone, self.flow_model_branch = create_inference_model(
opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test)
self.flow_model_backbone = convert2flow(opt.ninput,
self.flow_model_backbone)
self.flow_model_backbone, self.flow_model_branch = load_inference_model(
self.flow_model_backbone, self.flow_model_branch,
opt.flow_model)
self.flow_model_branch = DataParallel(
self.flow_model_branch,
device_ids=opt.gpus,
chunk_sizes=opt.chunk_sizes).to(opt.device)
self.flow_model_branch.eval()
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
self.rgb_model_backbone, self.rgb_model_branch = None, None
self.flow_model_backbone, self.flow_model_branch = None, None
if opt.rgb_model != '':
print('create rgb model')
self.rgb_model_backbone, self.rgb_model_branch = create_inference_model(
opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test)
self.rgb_model_backbone, self.rgb_model_branch = load_inference_model(
self.rgb_model_backbone, self.rgb_model_branch, opt.rgb_model)
self.rgb_model_backbone = DataParallel(self.rgb_model_backbone,
device_ids=[opt.gpus[0]],
chunk_sizes=[1]).to(
opt.device)
self.rgb_model_branch = DataParallel(self.rgb_model_branch,
device_ids=[opt.gpus[0]],
chunk_sizes=[1
]).to(opt.device)
self.rgb_model_backbone.eval()
self.rgb_model_branch.eval()
if opt.flow_model != '':
print('create flow model')
self.flow_model_backbone, self.flow_model_branch = create_inference_model(
opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test)
self.flow_model_backbone = convert2flow(opt.ninput,
self.flow_model_backbone)
self.flow_model_backbone, self.flow_model_branch = load_inference_model(
self.flow_model_backbone, self.flow_model_branch,
opt.flow_model)
self.flow_model_backbone = DataParallel(self.flow_model_backbone,
device_ids=[opt.gpus[0]],
chunk_sizes=[1]).to(
opt.device)
self.flow_model_branch = DataParallel(self.flow_model_branch,
device_ids=[opt.gpus[0]],
chunk_sizes=[1]).to(
opt.device)
self.flow_model_backbone.eval()
self.flow_model_branch.eval()
self.num_classes = opt.num_classes
self.opt = opt
self.rgb_buffer = []
self.flow_buffer = []
self.rgb_buffer_flip = []
self.flow_buffer_flip = []
def set_device(self, gpus, chunk_sizes, device):
if len(gpus) > 1:
self.model_with_loss = DataParallel(
self.model_with_loss, device_ids=gpus,
chunk_sizes=chunk_sizes).to(device)
else:
self.model_with_loss = self.model_with_loss.to(device)
for state in self.optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
# MODIFY for pytorch 0.4.0
state[k] = v.to(device=device, non_blocking=True)
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
self.rgb_model = None
self.flow_model = None
if opt.rgb_model != '':
print('create rgb model')
self.rgb_model = create_model(opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test)
self.rgb_model = load_model(self.rgb_model, opt.rgb_model)
self.rgb_model = DataParallel(self.rgb_model,
device_ids=opt.gpus,
chunk_sizes=opt.chunk_sizes).to(
opt.device)
self.rgb_model.eval()
if opt.flow_model != '':
print('create flow model')
self.flow_model = create_model(opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test)
self.flow_model = convert2flow(opt.ninput, self.flow_model)
self.flow_model = load_model(self.flow_model, opt.flow_model)
self.flow_model = DataParallel(self.flow_model,
device_ids=opt.gpus,
chunk_sizes=opt.chunk_sizes).to(
opt.device)
self.flow_model.eval()
self.num_classes = opt.num_classes
self.opt = opt
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
self.rgb_model = None
self.flow_model = None
self.pa_model = None
if opt.rgb_model != '':
print('create rgb model')
self.rgb_model = create_model(opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test,
is_pa=False)
self.rgb_model = load_model(self.rgb_model,
opt.save_root + opt.rgb_model)
'''
# ADDED: debug param weights
for i, child in enumerate(self.rgb_model.children()):
if i < 2 :
continue
for l, param in enumerate(child.parameters()):
if l == 0:
vistensor(param)
#param.requires_grad = False
#print(param.size())
'''
# ORIG
self.rgb_model = DataParallel(self.rgb_model,
device_ids=opt.gpus,
chunk_sizes=opt.chunk_sizes).to(
opt.device)
self.rgb_model.eval()
if opt.flow_model != '':
print('create flow model')
self.flow_model = create_model(opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test)
self.flow_model = convert2flow(opt.ninput, self.flow_model)
self.flow_model = load_model(self.flow_model, opt.flow_model)
self.flow_model = DataParallel(self.flow_model,
device_ids=opt.gpus,
chunk_sizes=opt.chunk_sizes).to(
opt.device)
self.flow_model.eval()
if opt.pa_model != '':
print('create PA model')
self.pa_model = create_model(opt.arch,
opt.branch_info,
opt.head_conv,
opt.K,
flip_test=opt.flip_test,
is_pa=True,
pa_fuse_mode=opt.pa_fuse_mode,
rgb_w3=opt.rgb_w3)
if opt.pa_fuse_mode == 'PAN':
self.pa_model = convert2PAN(opt.ninput,
self.pa_model,
conv_idx=1)
elif opt.pa_fuse_mode == 'TDN':
pass
#self.pa_model = convert2TDN(opt.ninput, self.pa_model, conv_idx=2) # idx 1 or 2? does not matter here as trained weight would be loaded here?
# Single PAN stream
else:
self.pa_model = convert2PAN(opt.ninput,
self.pa_model,
conv_idx=1)
self.pa_model = load_model(self.pa_model,
opt.save_root + opt.pa_model)
self.pa_model = DataParallel(
self.pa_model,
device_ids=opt.gpus, #[0]
chunk_sizes=opt.chunk_sizes).to(opt.device)
self.pa_model.eval()
self.num_classes = opt.num_classes
self.opt = opt
# added: for speed measurement
self.total_time = 0