def profile(opt, lr_size, test_speed=False):
# logging
logger = base_utils.get_logger('base')
logger.info('{} Model Information {}'.format('='*20, '='*20))
base_utils.print_options(opt['model']['generator'], logger)
# basic configs
scale = opt['scale']
device = torch.device(opt['device'])
# create model
net_G = define_generator(opt).to(device)
# get dummy input
dummy_input_dict = net_G.generate_dummy_input(lr_size)
for key in dummy_input_dict.keys():
dummy_input_dict[key] = dummy_input_dict[key].to(device)
# profile
register(net_G, dummy_input_dict)
gflops, params = profile_model(net_G)
logger.info('-' * 40)
logger.info('Super-resolute data from {}x{}x{} to {}x{}x{}'.format(
*lr_size, lr_size[0], lr_size[1]*scale, lr_size[2]*scale))
logger.info('Parameters (x10^6): {:.3f}'.format(params/1e6))
logger.info('FLOPs (x10^9): {:.3f}'.format(gflops))
logger.info('-' * 40)
# test running speed
if test_speed:
n_test = 3
tot_time = 0
for i in range(n_test):
start_time = time.time()
with torch.no_grad():
_ = net_G(**dummy_input_dict)
end_time = time.time()
tot_time += end_time - start_time
logger.info('Speed (FPS): {:.3f} (averaged for {} runs)'.format(
n_test / tot_time, n_test))
logger.info('-' * 40)
def profile(self, lr_size, device):
gflops_dict, params_dict = OrderedDict(), OrderedDict()
# generate dummy input data
lr_curr, lr_prev, hr_prev = self.generate_dummy_data(lr_size, device)
# profile module 1: flow estimation module
lr_flow = register(self.fnet, [lr_curr, lr_prev])
gflops_dict['FNet'], params_dict['FNet'] = parse_model_info(self.fnet)
# profile module 2: sr module
pad_h = lr_curr.size(2) - lr_curr.size(2) // 8 * 8
pad_w = lr_curr.size(3) - lr_curr.size(3) // 8 * 8
lr_flow_pad = F.pad(lr_flow, (0, pad_w, 0, pad_h), 'reflect')
hr_flow = self.scale * self.upsample_func(lr_flow_pad)
hr_prev_warp = backward_warp(hr_prev, hr_flow)
_ = register(
self.srnet,
[lr_curr, space_to_depth(hr_prev_warp, self.scale)])
gflops_dict['SRNet'], params_dict['SRNet'] = parse_model_info(
self.srnet)
return gflops_dict, params_dict