import os
from importlib import import_module
import torch
import torch.nn as nn
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()
if not args.cpu and args.n_GPUs > 1:
self.model = nn.DataParallel(self.model, range(args.n_GPUs))
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
target = self.get_model()
if hasattr(target, 'set_scale'): target.set_scale(idx_scale)
if self.self_ensemble and not self.training:
if self.chop:
forward_function = self.forward_chop
else:
forward_function = self.model.forward
return self.forward_x8(x, forward_function=forward_function)
elif self.chop and not self.training:
return self.forward_chop(x)
else:
return self.model(x)
def get_model(self):
if self.n_GPUs == 1:
return self.model
else:
return self.model.module
def state_dict(self, **kwargs):
target = self.get_model()
return target.state_dict(**kwargs)
def save(self, apath, epoch, is_best=False):
target = self.get_model()
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(target.state_dict(), s)
def load(self, apath, pre_train='', resume=-1, cpu=False):
if cpu:
kwargs = {'map_location': lambda storage, loc: storage}
else:
kwargs = {}
load_from = None
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.get_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.get_model().load_state_dict(load_from, strict=False)
def forward_chop(self, *args, shave=10, min_size=160000):
if self.input_large:
scale = 1
else:
scale = self.scale[self.idx_scale]
n_GPUs = min(self.n_GPUs, 4)
_, _, h, w = args[0].size()
h_half, w_half = h // 2, w // 2
h_size, w_size = h_half + shave, w_half + shave
list_x = [[
a[:, :, 0:h_size, 0:w_size],
a[:, :, 0:h_size, (w - w_size):w],
a[:, :, (h - h_size):h, 0:w_size],
a[:, :, (h - h_size):h, (w - w_size):w]
] for a in args]
list_y = []
if w_size * h_size < min_size:
for i in range(0, 4, n_GPUs):
x = [torch.cat(_x[i:(i + n_GPUs)], dim=0) for _x in list_x]
y = self.model(*x)
if not isinstance(y, list): y = [y]
if not list_y:
list_y = [[c for c in _y.chunk(n_GPUs, dim=0)] for _y in y]
else:
for _list_y, _y in zip(list_y, y):
_list_y.extend(_y.chunk(n_GPUs, dim=0))
else:
for p in zip(*list_x):
y = self.forward_chop(*p, shave=shave, min_size=min_size)
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)
h, w = scale * h, scale * w
h_half, w_half = scale * h_half, scale * w_half
h_size, w_size = scale * h_size, scale * w_size
shave *= scale
b, c, _, _ = list_y[0][0].size()
y = [_y[0].new(b, c, h, w) for _y in list_y]
for _list_y, _y in zip(list_y, y):
_y[:, :, :h_half, :w_half] \
= _list_y[0][:, :, :h_half, :w_half]
_y[:, :, :h_half, w_half:] \
= _list_y[1][:, :, :h_half, (w_size - w + w_half):]
_y[:, :, h_half:, :w_half] \
= _list_y[2][:, :, (h_size - h + h_half):, :w_half]
_y[:, :, h_half:, w_half:] \
= _list_y[3][:, :, (h_size - h + h_half):, (w_size - w + w_half):]
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
