| import torch |
| import torch.nn as nn |
| import numpy as np |
| from torch.optim import AdamW |
| import torch.optim as optim |
| import itertools |
| from model.warplayer import warp |
| from torch.nn.parallel import DistributedDataParallel as DDP |
| from train_log.IFNet_HDv3 import * |
| import torch.nn.functional as F |
| from model.loss import * |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| |
| class Model: |
| def __init__(self, local_rank=-1): |
| self.flownet = IFNet() |
| self.device() |
| self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4) |
| self.epe = EPE() |
| self.version = 4.25 |
| |
| self.sobel = SOBEL() |
| if local_rank != -1: |
| self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) |
|
|
| def train(self): |
| self.flownet.train() |
|
|
| def eval(self): |
| self.flownet.eval() |
|
|
| def device(self): |
| self.flownet.to(device) |
|
|
| def load_model(self, path, rank=0): |
| def convert(param): |
| if rank == -1: |
| return { |
| k.replace("module.", ""): v |
| for k, v in param.items() |
| if "module." in k |
| } |
| else: |
| return param |
| if rank <= 0: |
| if torch.cuda.is_available(): |
| self.flownet.load_state_dict(convert(torch.load('{}/flownet.pkl'.format(path))), False) |
| else: |
| self.flownet.load_state_dict(convert(torch.load('{}/flownet.pkl'.format(path), map_location ='cpu')), False) |
| |
| def save_model(self, path, rank=0): |
| if rank == 0: |
| torch.save(self.flownet.state_dict(),'{}/flownet.pkl'.format(path)) |
|
|
| def inference(self, img0, img1, timestep=0.5, scale=1.0): |
| imgs = torch.cat((img0, img1), 1) |
| scale_list = [16/scale, 8/scale, 4/scale, 2/scale, 1/scale] |
| flow, mask, merged = self.flownet(imgs, timestep, scale_list) |
| return merged[-1] |
| |
| def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): |
| for param_group in self.optimG.param_groups: |
| param_group['lr'] = learning_rate |
| img0 = imgs[:, :3] |
| img1 = imgs[:, 3:] |
| if training: |
| self.train() |
| else: |
| self.eval() |
| scale = [16, 8, 4, 2, 1] |
| flow, mask, merged = self.flownet(torch.cat((imgs, gt), 1), scale=scale, training=training) |
| loss_l1 = (merged[-1] - gt).abs().mean() |
| loss_smooth = self.sobel(flow[-1], flow[-1]*0).mean() |
| |
| if training: |
| self.optimG.zero_grad() |
| loss_G = loss_l1 + loss_cons + loss_smooth * 0.1 |
| loss_G.backward() |
| self.optimG.step() |
| else: |
| flow_teacher = flow[2] |
| return merged[-1], { |
| 'mask': mask, |
| 'flow': flow[-1][:, :2], |
| 'loss_l1': loss_l1, |
| 'loss_cons': loss_cons, |
| 'loss_smooth': loss_smooth, |
| } |
|
|
