def construct_model(model_path, device, nframes):
from rbpn import Net as RBPN
model = RBPN(num_channels=3,
base_filter=256,
feat=64,
num_stages=3,
n_resblock=5,
nFrames=nframes,
scale_factor=4)
ckpt = torch.load(model_path, map_location='cuda:0')
new_ckpt = {}
for key in ckpt:
if key.startswith('module'):
new_key = key[7:]
else:
new_key = key
new_ckpt[new_key] = ckpt[key]
model = model.to(device)
model.load_state_dict(new_ckpt)
model.eval()
return model
opt.data_augmentation, opt.file_list,
opt.other_dataset, opt.patch_size,
opt.future_frame)
#test_set = get_eval_set(opt.test_dir, opt.nFrames, opt.upscale_factor, opt.data_augmentation)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
#testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=False)
print('===> Building model ', opt.model_type)
if opt.model_type == 'RBPN':
model = RBPN(num_channels=3,
base_filter=256,
feat=64,
num_stages=3,
n_resblock=5,
nFrames=opt.nFrames,
scale_factor=opt.upscale_factor)
model = torch.nn.DataParallel(model, device_ids=gpus_list)
criterion = nn.L1Loss()
print('---------- Networks architecture -------------')
print_network(model)
print('----------------------------------------------')
if opt.pretrained:
model_name = os.path.join(opt.save_folder + opt.pretrained_sr)
if os.path.exists(model_name):
#model= torch.load(model_name, map_location=lambda storage, loc: storage)
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
test_set = get_test_set(opt.data_dir, opt.nFrames, opt.upscale_factor,
opt.file_list, opt.other_dataset, opt.future_frame)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
print('===> Building model ', opt.model_type)
if opt.model_type == 'RBPN':
model = RBPN(num_channels=3,
base_filter=256,
feat=64,
num_stages=3,
n_resblock=5,
nFrames=opt.nFrames,
scale_factor=opt.upscale_factor)
if cuda:
model = torch.nn.DataParallel(model, device_ids=gpus_list)
model.load_state_dict(
torch.load(opt.model, map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
if cuda:
model = model.cuda(gpus_list[0])
def __init__(self, base_filter, feat, num_stages, n_resblock, scale_factor, pretrained=True, freeze=False):
super(Net, self).__init__()
if scale_factor == 2:
kernel = 6
stride = 2
padding = 2
elif scale_factor == 4:
kernel = 8
stride = 4
padding = 2
elif scale_factor == 8:
kernel = 12
stride = 8
padding = 2
#Initial Feature Extraction
self.motion_feat = ConvBlock(4, base_filter, 3, 1, 1, activation='lrelu', norm=None)
###INTERPOLATION
#Interp_block
motion_net = [
ResnetBlock(base_filter, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(2)]
motion_net.append(ConvBlock(base_filter, feat, 3, 1, 1, activation='lrelu', norm=None))
self.motion = nn.Sequential(*motion_net)
t_net2 = [ConvBlock(feat*3, feat, 1, 1, 0, bias=True, activation='lrelu', norm=None)]
t_net2.append(PyramidModule(feat,activation='lrelu'))
t_net2.append(ConvBlock(feat, feat, 3, 1, 1, activation='lrelu', norm=None))
self.t_net_hr = nn.Sequential(*t_net2)
self.upsample_layer = nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=True)
interp_b = [
ResnetBlock(feat*5, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(n_resblock)]
interp_b.append(DeconvBlock(feat*5, feat, kernel, stride, padding, activation='lrelu', norm=None))
self.interp_block = nn.Sequential(*interp_b)
###ITERATIVE REFINEMENT
#Motion Up FORWARD
modules_up_f = [
ResnetBlock(feat*5, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(n_resblock)]
modules_up_f.append(DeconvBlock(feat*5, feat, kernel, stride, padding, activation='lrelu', norm=None))
self.motion_up_f = nn.Sequential(*modules_up_f)
#Motion Up BACKWARD
modules_up_b = [
ResnetBlock(feat*5, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(n_resblock)]
modules_up_b.append(DeconvBlock(feat*5, feat, kernel, stride, padding, activation='lrelu', norm=None))
self.motion_up_b = nn.Sequential(*modules_up_b)
#Motion Down
modules_down = [
ResnetBlock(feat, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(2)]
modules_down.append(ConvBlock(feat, feat*2, kernel, stride, padding, activation='lrelu', norm=None))
self.motion_down = nn.Sequential(*modules_down)
self.relu_bp = torch.nn.LeakyReLU(negative_slope=0.1, inplace=True)#torch.nn.PReLU()
#Reconstruction
self.reconstruction_l = ConvBlock(feat*2, 3, 3, 1, 1, activation=None, norm=None)
self.reconstruction_h = ConvBlock(feat, 3, 3, 1, 1, activation=None, norm=None)
####ALIGNMENT
###RBPN
self.RBPN = RBPN(num_channels=3, base_filter=base_filter, feat = feat, num_stages=num_stages, n_resblock=5, nFrames=2, scale_factor=scale_factor)
for m in self.modules():
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('ConvTranspose2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
if pretrained:
if scale_factor == 4:
self.RBPN.load_state_dict(torch.load("weights/pretrained/rbpn_pretrained_F2_4x.pth", map_location=lambda storage, loc: storage))
elif scale_factor == 2:
self.RBPN.load_state_dict(torch.load("weights/pretrained/rbpn_pretrained_F2_2x.pth", map_location=lambda storage, loc: storage))
if freeze:
self.freeze_model(self.RBPN)
torch.cuda.manual_seed(args.seed)
print('==> Loading datasets')
test_set = get_test_set(args.data_dir, args.nFrames, args.upscale_factor,
args.file_list, args.other_dataset, args.future_frame)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=args.threads,
batch_size=args.testBatchSize,
shuffle=False)
print('==> Building model ', args.model_type)
if args.model_type == 'RBPN':
model = RBPN(num_channels=3,
base_filter=256,
feat=64,
num_stages=3,
n_resblock=5,
nFrames=args.nFrames,
scale_factor=args.upscale_factor)
if cuda:
model = torch.nn.DataParallel(model, device_ids=gpus_list)
device = torch.device(
"cuda:0" if cuda and torch.cuda.is_available() else "cpu")
if cuda:
model = model.cuda(gpus_list[0])
def eval():
class Net(nn.Module):
def __init__(self, base_filter, feat, num_stages, n_resblock, scale_factor, pretrained=True, freeze=False):
super(Net, self).__init__()
if scale_factor == 2:
kernel = 6
stride = 2
padding = 2
elif scale_factor == 4:
kernel = 8
stride = 4
padding = 2
elif scale_factor == 8:
kernel = 12
stride = 8
padding = 2
#Initial Feature Extraction
self.motion_feat = ConvBlock(4, base_filter, 3, 1, 1, activation='lrelu', norm=None)
###INTERPOLATION
#Interp_block
motion_net = [
ResnetBlock(base_filter, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(2)]
motion_net.append(ConvBlock(base_filter, feat, 3, 1, 1, activation='lrelu', norm=None))
self.motion = nn.Sequential(*motion_net)
t_net2 = [ConvBlock(feat*3, feat, 1, 1, 0, bias=True, activation='lrelu', norm=None)]
t_net2.append(PyramidModule(feat,activation='lrelu'))
t_net2.append(ConvBlock(feat, feat, 3, 1, 1, activation='lrelu', norm=None))
self.t_net_hr = nn.Sequential(*t_net2)
self.upsample_layer = nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=True)
interp_b = [
ResnetBlock(feat*5, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(n_resblock)]
interp_b.append(DeconvBlock(feat*5, feat, kernel, stride, padding, activation='lrelu', norm=None))
self.interp_block = nn.Sequential(*interp_b)
###ITERATIVE REFINEMENT
#Motion Up FORWARD
modules_up_f = [
ResnetBlock(feat*5, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(n_resblock)]
modules_up_f.append(DeconvBlock(feat*5, feat, kernel, stride, padding, activation='lrelu', norm=None))
self.motion_up_f = nn.Sequential(*modules_up_f)
#Motion Up BACKWARD
modules_up_b = [
ResnetBlock(feat*5, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(n_resblock)]
modules_up_b.append(DeconvBlock(feat*5, feat, kernel, stride, padding, activation='lrelu', norm=None))
self.motion_up_b = nn.Sequential(*modules_up_b)
#Motion Down
modules_down = [
ResnetBlock(feat, kernel_size=3, stride=1, padding=1, bias=True, activation='lrelu', norm=None) \
for _ in range(2)]
modules_down.append(ConvBlock(feat, feat*2, kernel, stride, padding, activation='lrelu', norm=None))
self.motion_down = nn.Sequential(*modules_down)
self.relu_bp = torch.nn.LeakyReLU(negative_slope=0.1, inplace=True)#torch.nn.PReLU()
#Reconstruction
self.reconstruction_l = ConvBlock(feat*2, 3, 3, 1, 1, activation=None, norm=None)
self.reconstruction_h = ConvBlock(feat, 3, 3, 1, 1, activation=None, norm=None)
####ALIGNMENT
###RBPN
self.RBPN = RBPN(num_channels=3, base_filter=base_filter, feat = feat, num_stages=num_stages, n_resblock=5, nFrames=2, scale_factor=scale_factor)
for m in self.modules():
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('ConvTranspose2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
if pretrained:
if scale_factor == 4:
self.RBPN.load_state_dict(torch.load("weights/pretrained/rbpn_pretrained_F2_4x.pth", map_location=lambda storage, loc: storage))
elif scale_factor == 2:
self.RBPN.load_state_dict(torch.load("weights/pretrained/rbpn_pretrained_F2_2x.pth", map_location=lambda storage, loc: storage))
if freeze:
self.freeze_model(self.RBPN)
def freeze_model(self, model):
for child in model.children():
for param in child.parameters():
param.requires_grad = False
def forward(self, t_im1, t_im2, t_flow_f, t_flow_b, train=True):
result_l = []
result_h1 = []
result_ht = []
result_h2 = []
###ALIGNMENT
aux_H1, H1 = self.RBPN(t_im1,[t_im2],[t_flow_f])
aux_H2, H2 = self.RBPN(t_im2,[t_im1],[t_flow_b])
L1 = self.motion_down(H1)
L2 = self.motion_down(H2)
###MOTION & DEPTH
motion_feat0 = self.motion_feat(torch.cat((t_flow_f, t_flow_b),1))
M = self.motion(motion_feat0)
motion_feat1 = self.motion_feat(torch.cat((t_flow_f/2.0, t_flow_b/2.0),1))
M_half = self.motion(motion_feat1)
###INTERPOLATION
Ht = self.interp_block(torch.cat((L1,L2,M),1))
Ht = Ht + self.relu_bp(Ht - self.t_net_hr(torch.cat((H1,H2,self.upsample_layer(M)),1)))
L = self.motion_down(Ht)
aux_Ht = self.reconstruction_h(Ht)
aux_L = self.reconstruction_l(L)
result_l.append(aux_L)
result_h1.append(aux_H1)
result_ht.append(aux_Ht)
result_h2.append(aux_H2)
####Projection
backward1 = torch.cat((L1, L, M_half),1)
H_b = self.motion_up_b(backward1)
H1 = H1 + self.relu_bp(H1 - H_b)
L1 = L1 + self.relu_bp(L1 - self.motion_down(H_b))
forwardd2 = torch.cat((L, L2, M_half),1)
H_f = self.motion_up_f(forwardd2)
H2 = H2 + self.relu_bp(H2 - H_f)
L2 = L2 + self.relu_bp(L2 - self.motion_down(H_f))
forwardd = torch.cat((L1, L, M_half),1)
H_t_f = self.motion_up_f(forwardd)
Ht = Ht + self.relu_bp(Ht - H_t_f)
L = L + self.relu_bp(L - self.motion_down(H_t_f))
backward = torch.cat((L, L2, M_half),1)
H_t_b = self.motion_up_b(backward)
Ht = Ht + self.relu_bp(Ht - H_t_b)
L = L + self.relu_bp(L - self.motion_down(H_t_b))
output_ht = self.reconstruction_h(Ht)
output_h1 = self.reconstruction_h(H1)
output_h2 = self.reconstruction_h(H2)
output_l = self.reconstruction_l(L)
result_l.append(output_l)
result_h1.append(output_h1)
result_ht.append(output_ht)
result_h2.append(output_h2)
if train:
return result_ht, result_h1, result_h2, result_l
else:
return output_ht, output_h1, output_h2, output_l
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
train_set = get_training_set(opt.data_dir, opt.nFrames, opt.upscale_factor, opt.data_augmentation, opt.file_list, opt.other_dataset, opt.patch_size, opt.future_frame)
#test_set = get_eval_set(opt.test_dir, opt.nFrames, opt.upscale_factor, opt.data_augmentation)
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
#testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=False)
print('===> Building model ', opt.model_type)
if opt.model_type == 'RBPN':
model = RBPN(num_channels=3, base_filter=256, feat = 64, num_stages=3, n_resblock=5, nFrames=opt.nFrames, scale_factor=opt.upscale_factor)
model = torch.nn.DataParallel(model, device_ids=gpus_list)
criterion = nn.L1Loss()
print('---------- Networks architecture -------------')
print_network(model)
print('----------------------------------------------')
if opt.pretrained:
model_name = os.path.join(opt.save_folder + opt.pretrained_sr)
if os.path.exists(model_name):
#model= torch.load(model_name, map_location=lambda storage, loc: storage)
model.load_state_dict(torch.load(model_name, map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
def main():
""" Lets begin the training process! """
args = parser.parse_args()
# Initialize Logger
logger.initLogger(args.debug)
# Load dataset
logger.info('==> Loading datasets')
# print(args.file_list)
# sys.exit()
train_set = get_training_set(args.data_dir, args.nFrames,
args.upscale_factor, args.data_augmentation,
args.file_list, args.other_dataset,
args.patch_size, args.future_frame)
training_data_loader = DataLoader(dataset=train_set,
num_workers=args.threads,
batch_size=args.batchSize,
shuffle=True)
# Use generator as RBPN
netG = RBPN(num_channels=3,
base_filter=256,
feat=64,
num_stages=3,
n_resblock=5,
nFrames=args.nFrames,
scale_factor=args.upscale_factor)
logger.info('# of Generator parameters: %s',
sum(param.numel() for param in netG.parameters()))
# Use DataParallel?
if args.useDataParallel:
gpus_list = range(args.gpus)
netG = torch.nn.DataParallel(netG, device_ids=gpus_list)
# Use discriminator from SRGAN
netD = Discriminator()
logger.info('# of Discriminator parameters: %s',
sum(param.numel() for param in netD.parameters()))
# Generator loss
generatorCriterion = nn.L1Loss() if not args.APITLoss else GeneratorLoss()
# Specify device
device = torch.device(
"cuda:0" if torch.cuda.is_available() and args.gpu_mode else "cpu")
if args.gpu_mode and torch.cuda.is_available():
utils.printCUDAStats()
netG.cuda()
netD.cuda()
netG.to(device)
netD.to(device)
generatorCriterion.cuda()
# Use Adam optimizer
optimizerG = optim.Adam(netG.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-8)
optimizerD = optim.Adam(netD.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-8)
if args.APITLoss:
logger.info(
"Generator Loss: Adversarial Loss + Perception Loss + Image Loss + TV Loss"
)
else:
logger.info("Generator Loss: L1 Loss")
# print iSeeBetter architecture
utils.printNetworkArch(netG, netD)
if args.pretrained:
modelPath = os.path.join(args.save_folder + args.pretrained_sr)
utils.loadPreTrainedModel(gpuMode=args.gpu_mode,
model=netG,
modelPath=modelPath)
# sys.exit()
for epoch in range(args.start_epoch, args.nEpochs + 1):
runningResults = trainModel(epoch, training_data_loader, netG, netD,
optimizerD, optimizerG, generatorCriterion,
device, args)
if (epoch + 1) % (args.snapshots) == 0:
saveModelParams(epoch, runningResults, netG, netD)
#training code
if args.phase == 'train':
dataloaders = data.DataLoader(DataLoader(args),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
device = torch.device("cuda:0")
print("constructing model ....")
model = RBPN(num_channels=3,
base_filter=256,
feat=64,
num_stages=3,
n_resblock=5,
nFrames=args.nframes,
scale_factor=4)
model = nn.DataParallel(model.to(device), gpuids)
if args.resume:
ckpt = torch.load(args.model_path)
new_ckpt = {}
for key in ckpt:
if not key.startswith('module'):
new_key = 'module.' + key
else:
new_key = key
new_ckpt[new_key] = ckpt[key]
if sys.version_info.minor == 5:
print('===> Disabling cpu threads because of a bug in Python 3.5')
opt.threads = 0
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
test_set = get_test_set(opt.data_dir, opt.nFrames, opt.upscale_factor, opt.file_list, opt.other_dataset, opt.future_frame)
testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False)
print('===> Building model', opt.model_type)
if opt.model_type == 'RBPN':
model = RBPN(num_channels=3, base_filter=256, feat = 64, num_stages=3, n_resblock=5, nFrames=opt.nFrames, scale_factor=opt.upscale_factor)
if cuda:
model = torch.nn.DataParallel(model, device_ids=gpus_list)
model_name = opt.model
if not model_name:
# Loading the NTIRE2019 model doesn't actually work (size missmatch)
if False and opt.upscale_factor == 4:
model_name = 'weights/RBPN_4x_F11_NTIRE2019.pth'
else:
model_name = 'weights/RBPN_' + str(opt.upscale_factor) + 'x.pth'
print('===> Using pretrained model', model_name)
model_state_dict = torch.load(model_name, map_location=lambda storage, loc: storage)
