def main(stage,cuda=True):
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found")
if cuda:
torch.cuda.manual_seed(123)
else:
torch.manual_seed(123)
fg_path = '/scratch/matting/dataset/Training_set/adobe/fg/'
bg_path = '/scratch/matting/dataset/train2014/'
alpha_path = '/scratch/matting/dataset/Training_set/adobe/alpha/'
vgg_path = '/scratch/matting/dataset/vgg_state_dict.pth'
stage1_path = '/scratch/matting/dataset/full_trained.pth'
fg_file = '/scratch/matting/dataset/Training_set/training_fg_names.txt'
bg_file = '/scratch/matting/dataset/Training_set/training_bg_names.txt'
epochs = 10
files_list = get_files_list(fg_file, bg_file)
dataset = load_dataset(fg_path, alpha_path, bg_path, files_list)
model = get_model(stage)
print(model)
# vgg_dict = torch.load(vgg_path)
# model = copy_pretrain_vals(model, vgg_dict)
stage1_model = torch.load(stage1_path, map_location='cpu')['state_dict']
model = copy_pretrain_vals(model, stage1_model)
if cuda:
model = model.cuda()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-5)
for i in range(epochs):
train_model(stage ,model, optimizer, dataset, i, cuda)
save_model('/scratch/matting/model_' + str(i) + '_stage2.pth', model)
def validate(args, device_id, pt, step):
if pt != '':
test_from = pt
else:
test_from = args.test_from
print('Loading checkpoint from %s' % test_from)
logger.info('Loading checkpoint from %s' % test_from)
checkpoint = torch.load(test_from,
map_location=lambda storage, loc: storage)
opt = vars(checkpoint['opt'])
for k in opt.keys():
if k in model_flags:
setattr(args, k, opt[k])
print(args)
model = Summarizer(args, device, checkpoint)
valid_iter = DataLoader.Dataloader(args,
load_dataset(args,
'valid',
shuffle=False),
args.batch_size,
device,
shuffle=False,
is_test=False)
trainer = build_trainer(args, device_id, model, None)
stats = trainer.validate(valid_iter, step)
return stats.xent()
def train_iter_fct():
return DataLoader.Dataloader(args,
load_dataset(args, 'train', shuffle=True),
args.batch_size,
device,
shuffle=True,
is_test=False)
def main(cuda=False):
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found")
if cuda:
torch.cuda.manual_seed(123)
else:
torch.manual_seed(123)
fg_path = '/scratch/matting/dataset/Test_set/adobe/fg/'
bg_path = '/scratch/matting/dataset/VOCO_dataset/'
alpha_path = '/scratch/matting/dataset/Test_set/adobe/alpha/'
trimap_path = '/scratch/matting/dataset/Test_set/adobe/trimaps/'
# vgg_path = '/scratch/matting/dataset/full_trained.pth'
vgg_path = '/scratch/matting/model_1.pth'
fg_file = '/scratch/matting/dataset/Test_set/test_fg_names.txt'
bg_file = '/scratch/matting/dataset/Test_set/test_bg_names.txt'
files_list = get_files_list(fg_file, bg_file)
# print(files_list)
dataset = load_dataset(fg_path, alpha_path, bg_path, files_list)
model = get_model()
vgg_dict = torch.load(vgg_path, map_location='cpu')
# vgg_dict = torch.load_state_dict(vgg_dict)
# model = copy_pretrain_vals(model, vgg_dict["state_dict"], True)
model = copy_pretrain_vals(model, vgg_dict.state_dict(), True)
print(model)
if cuda:
model = model.cuda()
test_model(model, dataset, trimap_path, cuda)
def test_ext(args, device_id, pt, step):
if pt != '':
test_from = pt
else:
test_from = args.test_from
print('Loading checkpoint from %s' % test_from)
checkpoint = torch.load(test_from,
map_location=lambda storage, loc: storage)
opt = vars(checkpoint['opt'])
for k in opt.keys():
if k in model_flags:
setattr(args, k, opt[k])
print(args)
model = Summarizer(args, device, checkpoint)
test_iter = DataLoader.Dataloader(args,
load_dataset(args, 'test',
shuffle=False),
args.test_batch_size,
device,
shuffle=False,
is_test=True)
trainer = build_trainer(args, device_id, model, None)
trainer.test(test_iter, step)