| import os |
| from importlib import import_module |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.parallel as P |
| import torch.utils.model_zoo |
|
|
| class Model(nn.Module): |
| def __init__(self, args, ckp): |
| super(Model, self).__init__() |
| print('Making model...') |
|
|
| self.scale = args.scale |
| self.idx_scale = 0 |
| self.input_large = (args.model == 'VDSR') |
| self.self_ensemble = args.self_ensemble |
| self.chop = args.chop |
| self.precision = args.precision |
| self.cpu = args.cpu |
| self.device = torch.device('cpu' if args.cpu else 'cuda') |
| self.n_GPUs = args.n_GPUs |
| self.save_models = args.save_models |
|
|
| module = import_module('model.' + args.model.lower()) |
| self.model = module.make_model(args).to(self.device) |
| if args.precision == 'half': |
| self.model.half() |
|
|
| self.load( |
| ckp.get_path('model'), |
| pre_train=args.pre_train, |
| resume=args.resume, |
| cpu=args.cpu |
| ) |
| print(self.model, file=ckp.log_file) |
|
|
| def forward(self, x, idx_scale): |
| self.idx_scale = idx_scale |
| if hasattr(self.model, 'set_scale'): |
| self.model.set_scale(idx_scale) |
|
|
| if self.training: |
| if self.n_GPUs > 1: |
| return P.data_parallel(self.model, x, range(self.n_GPUs)) |
| else: |
| return self.model(x) |
| else: |
| if self.chop: |
| forward_function = self.forward_chop |
| else: |
| forward_function = self.model.forward |
|
|
| if self.self_ensemble: |
| return self.forward_x8(x, forward_function=forward_function) |
| else: |
| return forward_function(x) |
|
|
| def save(self, apath, epoch, is_best=False): |
| save_dirs = [os.path.join(apath, 'model_latest.pt')] |
|
|
| if is_best: |
| save_dirs.append(os.path.join(apath, 'model_best.pt')) |
| if self.save_models: |
| save_dirs.append( |
| os.path.join(apath, 'model_{}.pt'.format(epoch)) |
| ) |
|
|
| for s in save_dirs: |
| torch.save(self.model.state_dict(), s) |
|
|
| def load(self, apath, pre_train='', resume=-1, cpu=False): |
| load_from = None |
| kwargs = {} |
| if cpu: |
| kwargs = {'map_location': lambda storage, loc: storage} |
|
|
| if resume == -1: |
| load_from = torch.load( |
| os.path.join(apath, 'model_latest.pt'), |
| **kwargs |
| ) |
| elif resume == 0: |
| if pre_train == 'download': |
| print('Download the model') |
| dir_model = os.path.join('..', 'models') |
| os.makedirs(dir_model, exist_ok=True) |
| load_from = torch.utils.model_zoo.load_url( |
| self.model.url, |
| model_dir=dir_model, |
| **kwargs |
| ) |
| elif pre_train: |
| print('Load the model from {}'.format(pre_train)) |
| load_from = torch.load(pre_train, **kwargs) |
| else: |
| load_from = torch.load( |
| os.path.join(apath, 'model_{}.pt'.format(resume)), |
| **kwargs |
| ) |
|
|
| if load_from: |
| self.model.load_state_dict(load_from, strict=False) |
|
|
| def forward_chop(self, *args, shave=10, min_size=160000): |
| scale = 1 if self.input_large else self.scale[self.idx_scale] |
| n_GPUs = min(self.n_GPUs, 4) |
| |
| h, w = args[0].size()[-2:] |
|
|
| top = slice(0, h//2 + shave) |
| bottom = slice(h - h//2 - shave, h) |
| left = slice(0, w//2 + shave) |
| right = slice(w - w//2 - shave, w) |
| x_chops = [torch.cat([ |
| a[..., top, left], |
| a[..., top, right], |
| a[..., bottom, left], |
| a[..., bottom, right] |
| ]) for a in args] |
|
|
| y_chops = [] |
| if h * w < 4 * min_size: |
| for i in range(0, 4, n_GPUs): |
| x = [x_chop[i:(i + n_GPUs)] for x_chop in x_chops] |
| y = P.data_parallel(self.model, *x, range(n_GPUs)) |
| if not isinstance(y, list): y = [y] |
| if not y_chops: |
| y_chops = [[c for c in _y.chunk(n_GPUs, dim=0)] for _y in y] |
| else: |
| for y_chop, _y in zip(y_chops, y): |
| y_chop.extend(_y.chunk(n_GPUs, dim=0)) |
| else: |
| for p in zip(*x_chops): |
| y = self.forward_chop(*p, shave=shave, min_size=min_size) |
| if not isinstance(y, list): y = [y] |
| if not y_chops: |
| y_chops = [[_y] for _y in y] |
| else: |
| for y_chop, _y in zip(y_chops, y): y_chop.append(_y) |
|
|
| h *= scale |
| w *= scale |
| top = slice(0, h//2) |
| bottom = slice(h - h//2, h) |
| bottom_r = slice(h//2 - h, None) |
| left = slice(0, w//2) |
| right = slice(w - w//2, w) |
| right_r = slice(w//2 - w, None) |
|
|
| |
| b, c = y_chops[0][0].size()[:-2] |
| y = [y_chop[0].new(b, c, h, w) for y_chop in y_chops] |
| for y_chop, _y in zip(y_chops, y): |
| _y[..., top, left] = y_chop[0][..., top, left] |
| _y[..., top, right] = y_chop[1][..., top, right_r] |
| _y[..., bottom, left] = y_chop[2][..., bottom_r, left] |
| _y[..., bottom, right] = y_chop[3][..., bottom_r, right_r] |
|
|
| if len(y) == 1: y = y[0] |
|
|
| return y |
|
|
| def forward_x8(self, *args, forward_function=None): |
| def _transform(v, op): |
| if self.precision != 'single': v = v.float() |
|
|
| v2np = v.data.cpu().numpy() |
| if op == 'v': |
| tfnp = v2np[:, :, :, ::-1].copy() |
| elif op == 'h': |
| tfnp = v2np[:, :, ::-1, :].copy() |
| elif op == 't': |
| tfnp = v2np.transpose((0, 1, 3, 2)).copy() |
|
|
| ret = torch.Tensor(tfnp).to(self.device) |
| if self.precision == 'half': ret = ret.half() |
|
|
| return ret |
|
|
| list_x = [] |
| for a in args: |
| x = [a] |
| for tf in 'v', 'h', 't': x.extend([_transform(_x, tf) for _x in x]) |
|
|
| list_x.append(x) |
|
|
| list_y = [] |
| for x in zip(*list_x): |
| y = forward_function(*x) |
| if not isinstance(y, list): y = [y] |
| if not list_y: |
| list_y = [[_y] for _y in y] |
| else: |
| for _list_y, _y in zip(list_y, y): _list_y.append(_y) |
|
|
| for _list_y in list_y: |
| for i in range(len(_list_y)): |
| if i > 3: |
| _list_y[i] = _transform(_list_y[i], 't') |
| if i % 4 > 1: |
| _list_y[i] = _transform(_list_y[i], 'h') |
| if (i % 4) % 2 == 1: |
| _list_y[i] = _transform(_list_y[i], 'v') |
|
|
| y = [torch.cat(_y, dim=0).mean(dim=0, keepdim=True) for _y in list_y] |
| if len(y) == 1: y = y[0] |
|
|
| return y |
|
|
