def forward(self, lr_curr, lr_prev, hr_prev):
"""
Parameters:
:param lr_curr: the current lr data in shape nchw
:param lr_prev: the previous lr data in shape nchw
:param hr_prev: the previous hr data in shape nc(4h)(4w)
"""
# estimate lr flow (lr_curr -> lr_prev)
lr_flow = self.fnet(lr_curr, lr_prev)
# pad if size is not a multiple of 8
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')
# upsample lr flow
hr_flow = self.scale * self.upsample_func(lr_flow_pad)
# warp hr_prev
hr_prev_warp = backward_warp(hr_prev, hr_flow)
# compute hr_curr
hr_curr = self.srnet(lr_curr, space_to_depth(hr_prev_warp, self.scale))
return hr_curr
def forward_sequence(self, lr_data):
"""
Parameters:
:param lr_data: lr data in shape ntchw
"""
n, t, c, lr_h, lr_w = lr_data.size()
hr_h, hr_w = lr_h * self.scale, lr_w * self.scale
# calculate optical flows
lr_prev = lr_data[:, :-1, ...].reshape(n * (t - 1), c, lr_h, lr_w)
lr_curr = lr_data[:, 1:, ...].reshape(n * (t - 1), c, lr_h, lr_w)
lr_flow = self.fnet(lr_curr, lr_prev) # n*(t-1),2,h,w
# upsample lr flows
hr_flow = self.scale * self.upsample_func(lr_flow)
hr_flow = hr_flow.view(n, (t - 1), 2, hr_h, hr_w)
# compute the first hr data
hr_data = []
hr_prev = self.srnet(
lr_data[:, 0, ...],
torch.zeros(n,
self.out_nc,
lr_h,
lr_w,
dtype=torch.float32,
device=lr_data.device))
hr_data.append(hr_prev)
# compute the remaining hr data
for i in range(1, t):
# warp hr_prev
hr_prev_warp = backward_warp(hr_prev, hr_flow[:, i - 1, ...])
# compute hr_curr
hr_curr = self.srnet(lr_data[:, i, ...],
space_to_depth(hr_prev_warp, self.scale))
# save and update
hr_data.append(hr_curr)
hr_prev = hr_curr
hr_data = torch.stack(hr_data, dim=1) # n,t,c,hr_h,hr_w
# construct output dict
ret_dict = {
'hr_data': hr_data, # n,t,c,hr_h,hr_w
'hr_flow': hr_flow, # n,t,2,hr_h,hr_w
'lr_prev': lr_prev, # n(t-1),c,lr_h,lr_w
'lr_curr': lr_curr, # n(t-1),c,lr_h,lr_w
'lr_flow': lr_flow, # n(t-1),2,lr_h,lr_w
}
return ret_dict
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