def load_stage1_model():
net_stage1 = DnCNN(channels=12, out_ch=6, num_of_layers=opt.num_of_layers)
if use_irregular:
PATH = '{}/irregular_{}_net.pth'.format(opt.stage1_outf, opt.dataset)
print('Loading from ', PATH)
ckpt = torch.load(PATH, map_location='cpu')
net_stage1.load_state_dict(ckpt)
return net_stage1
def load_stage1_model():
# net_stage1 = DnCNN(channels=12, out_ch=6, num_of_layers=opt.num_of_layers)
device_ids = [0]
net_stage1 = DnCNN(channels=4, out_ch=2, num_of_layers=opt.num_of_layers)
net_stage1 = nn.DataParallel(net_stage1, device_ids=device_ids).cuda()
if use_irregular:
PATH = '{}/irregular_{}_net.pth'.format(opt.stage1_outf, opt.dataset)
print('Loading from ', PATH)
ckpt = torch.load(PATH, map_location='cpu')
net_stage1.load_state_dict(ckpt)
return net_stage1
def load_stage1_model():
net_stage1 = DnCNN(channels=6, num_of_layers=opt.num_of_layers)
if opt.dataset != 'olivetti':
net_stage1 = net_stage1.cuda()
if use_regular:
PATH = '{}/regular_{}_net.pth'.format(opt.stage1_outf, opt.dataset)
else:
PATH = '{}/{}_net.pth'.format(opt.stage1_outf, opt.dataset)
print('Loading from ', PATH)
ckpt = torch.load(PATH, map_location='cpu')
net_stage1.load_state_dict(ckpt)
return net_stage1