def get_base_parses(run_id, trials_per=800, reverse=False, dry_run=False):
print('run_id: %i' % run_id)
print('Loading model...')
type_model = TypeModel().eval()
if torch.cuda.is_available():
type_model = type_model.cuda()
score_fn = lambda parses : model_score_fn(type_model, parses)
print('Loading classification dataset...')
dataset = ClassificationDataset(osc_folder='./one-shot-classification')
run = dataset.runs[run_id]
if reverse:
imgs = run.test_imgs
else:
imgs = run.train_imgs
results_dir = './results'
run_dir = os.path.join(results_dir, 'run%0.2i' % (run_id+1))
save_dir = os.path.join(run_dir, 'base_parses')
if not dry_run:
mkdir(results_dir)
mkdir(run_dir)
mkdir(save_dir)
print('Collecting top-K parses for each train image...')
nimg = len(imgs)
for i in range(nimg):
start_time = time.time()
parses, log_probs = get_topK_parses(
imgs[i], k=5, score_fn=score_fn, configs_per=1,
trials_per=trials_per)
total_time = time.time() - start_time
print('image %i/%i took %s' % (i+1, nimg, time_string(total_time)))
if dry_run:
continue
save_img_results(save_dir, i, parses, log_probs, reverse)
def select_parses(score_fn, images, configs_per=1, trials_per=800):
nimg = len(images)
base_parses = []
for i in range(nimg):
start_time = time.time()
parses, log_probs = get_topK_parses(
images[i], k=5, score_fn=score_fn,
configs_per=configs_per, trials_per=trials_per)
total_time = time.time() - start_time
print('parsing image %i/%i took %s' % (i+1, nimg, time_string(total_time)))
base_parses.append(parses)
return base_parses
def optimize_parses(model, parse_list, target_imgs, iterations=1500):
render_params = [p for parse in parse_list for p in parse.render_params]
stroke_params = [p for parse in parse_list for p in parse.stroke_params]
param_groups = [
{'params': render_params, 'lr': 0.306983},
{'params': stroke_params, 'lr': 0.044114}
]
optimizer = torch.optim.Adam(param_groups)
# optimize
start_time = time.time()
losses, states = opt.optimize_parselist(
parse_list, target_imgs,
loss_fn=model.losses_fn,
iterations=iterations,
optimizer=optimizer,
tune_blur=True,
tune_fn=model.likelihood_losses_fn
)
total_time = time.time() - start_time
time.sleep(0.5)
print('Took %s' % time_string(total_time))
parse_scores = -losses[-1]
return parse_list, parse_scores
def optimize_parses(run_id, iterations=1500, reverse=False, dry_run=False):
run_dir = './results/run%0.2i' % (run_id+1)
load_dir = os.path.join(run_dir, 'base_parses')
save_dir = os.path.join(run_dir, 'tuned_parses')
assert os.path.exists(run_dir)
assert os.path.exists(load_dir)
if not dry_run:
mkdir(save_dir)
print('Loading model and data...')
type_model = TypeModel().eval()
token_model = TokenModel()
renderer = Renderer()
# move to GPU if available
if torch.cuda.is_available():
torch.backends.cudnn.enabled = False
type_model = type_model.cuda()
token_model = token_model.cuda()
renderer = renderer.cuda()
# build full model
model = opt.FullModel(
renderer=renderer, type_model=type_model, token_model=token_model,
denormalize=True)
print('Loading data...')
# load classification dataset and select run
dataset = ClassificationDataset(osc_folder='./one-shot-classification')
run = dataset.runs[run_id]
# load images and base parses for this run
base_parses, images, K_per_img = load_parses(run, load_dir, reverse)
assert len(base_parses) == len(images)
print('total # parses: %i' % len(images))
print('Optimizing parses...')
# initialize Parse modules and optimizer
parse_list = [opt.ParseWithToken(p) for p in base_parses]
render_params = [p for parse in parse_list for p in parse.render_params]
stroke_params = [p for parse in parse_list for p in parse.stroke_params]
param_groups = [
{'params': render_params, 'lr': 0.306983},
{'params': stroke_params, 'lr': 0.044114}
]
optimizer = torch.optim.Adam(param_groups)
# optimize
start_time = time.time()
losses, states = opt.optimize_parselist(
parse_list, images,
loss_fn=model.losses_fn,
iterations=iterations,
optimizer=optimizer,
tune_blur=True,
tune_fn=model.likelihood_losses_fn
)
total_time = time.time() - start_time
time.sleep(0.5)
print('Took %s' % time_string(total_time))
if dry_run:
return
parse_scores = -losses[-1]
save_new_parses(parse_list, parse_scores, save_dir, K_per_img, reverse)
def refit_parses_single(run_id,
test_id,
iterations=1500,
reverse=False,
run=None,
dry_run=False):
run_dir = './results/run%0.2i' % (run_id + 1)
load_dir = os.path.join(run_dir, 'tuned_parses')
save_dir = os.path.join(run_dir, 'refitted_parses')
assert os.path.exists(run_dir)
assert os.path.exists(load_dir)
if not dry_run:
mkdir(save_dir)
print('Loading model...')
token_model = TokenModel()
renderer = Renderer(blur_fsize=21)
if torch.cuda.is_available():
token_model = token_model.cuda()
renderer = renderer.cuda()
model = opt.FullModel(renderer=renderer, token_model=token_model)
print('Loading parses...')
# load classification dataset and select run
if run is None:
dataset = ClassificationDataset(osc_folder='./one-shot-classification')
run = dataset.runs[run_id]
if reverse:
ntrain = len(run.test_imgs)
test_img = torch.from_numpy(run.train_imgs[test_id]).float()
else:
ntrain = len(run.train_imgs)
test_img = torch.from_numpy(run.test_imgs[test_id]).float()
if torch.cuda.is_available():
test_img = test_img.cuda()
# load tuned parses
parse_list, K_per_img = load_tuned_parses(load_dir, ntrain, reverse)
images = test_img.expand(len(parse_list), 105, 105)
print('total # parses: %i' % len(images))
print('Optimizing parses...')
# initialize Parse modules and optimizer
render_params = [p for parse in parse_list for p in parse.render_params]
stroke_params = [
p for parse in parse_list for p in parse.stroke_params
if p.requires_grad
]
param_groups = [{
'params': render_params,
'lr': 0.087992
}, {
'params': stroke_params,
'lr': 0.166810
}]
optimizer = torch.optim.Adam(param_groups)
# optimize
start_time = time.time()
losses, states = opt.optimize_parselist(parse_list,
images,
loss_fn=model.losses_fn,
iterations=iterations,
optimizer=optimizer,
tune_blur=True,
tune_fn=model.likelihood_losses_fn)
total_time = time.time() - start_time
time.sleep(0.5)
print('Took %s' % time_string(total_time))
if dry_run:
return
parse_scores = -losses[-1]
save_new_parses(parse_list, parse_scores, save_dir, K_per_img, test_id,
reverse)