def __init__(self, nframes, num_stages):
super(stacked_EDVR, self).__init__()
self.nframes = nframes
self.num_stages = num_stages
self.edvr = EDVR(128, nframes, 8, 5, 40)
self.edges = FSRCNN()
def __init__(self, nframes, num_stages):
super(stacked_EDVR, self).__init__()
self.nframes = nframes
self.num_stages = num_stages
self.edvr = EDVR(128, nframes, 8, 5, 40)
self.restore = nn.ModuleList()
for i in range(num_stages):
self.restore.append(restore(3))
def construct_model(model_path, device):
from modules.EDVR_arch import EDVR
model = EDVR(128, args.nframes, 8, 5, 40)
ckpt = torch.load(model_path)
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
import torch
import torch.backends.cudnn as cudnn
#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 = EDVR(128, args.nframes, 8, 5, 40)
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]
model.load_state_dict(new_ckpt, strict=False)
print("model constructed")
model.load_state_dict(new_ckpt, strict=True)
return model
#training code
if args.phase == 'train':
dataloader = 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 ....")
model1 = EDVR(128, args.nframes, 8, 5, 40)
model1 = nn.DataParallel(model1.to(device), gpuids)
model2 = D()
model2 = nn.DataParallel(model2.to(device), gpuids)
if args.resume:
model1 = load_state_dict(model1, args.model_path)
model2 = load_state_dict(model2, args.dis_model_path)
print("model constructed")
optimizer1 = torch.optim.Adam(model1.parameters(), lr=args.lr)
scheduler1 = ExponentialLR(optimizer1, gamma=args.gamma)
optimizer2 = torch.optim.Adam(model2.parameters(), lr=args.lr)
scheduler2 = ExponentialLR(optimizer2, gamma=args.gamma)