def upernetmod(root, tree, instructor):
encoder_path, decoder_path = utils.upernet_ckpt(root)
seg_model = upernet.get_upernet(tree, weights_encoder=encoder_path, weights_decoder=decoder_path)
fpn = upernet.extract_fpn(seg_model)
fpn = add_taskmod(fpn)
model = ModModel(instructor, fpn, seg_model.decoder.object_head)
return model
def taskmod(root, tree, obj_classes=None):
encoder_path, decoder_path = utils.upernet_ckpt(root)
instructor = Instructor(tree, obj_classes=obj_classes)
fpn = upernet.get_fpn(tree, weights_encoder=encoder_path, weights_decoder=decoder_path)
fpn = add_taskmod(fpn)
model = TaskMod(instructor, fpn)
return model, instructor
def main(args):
broden_root = Path(args.root).resolve()
tree = parts.ObjectTree.from_meta_folder(broden_root / 'meta')
encoder_path, decoder_path = utils.upernet_ckpt(broden_root)
model = Head3(tree, encoder_path, decoder_path)
db = upernet_data_pipeline(broden_root)
metrics = partial(parts.BrodenMetricsClbk,
obj_tree=tree,
split_func=lambda o: o)
learn = Learner(
db,
model,
loss_func=Loss(tree),
callback_fns=metrics,
)
learn.split((learn.model.obj_branch, ))
learn.freeze()
utils.fit_and_log(learn,
'object-P.A',
save=args.save,
epochs=20,
lr=1e-2,
pct_start=args.pct_start)
def main(args):
broden_root = Path(args.root).resolve()
tree = parts.ObjectTree.from_meta_folder(broden_root / 'meta')
encoder_path, decoder_path = utils.upernet_ckpt(broden_root)
model = cs_head.CSHead2(tree,
encoder_path,
decoder_path,
hidden=args.hidden)
clbk = cs_head.Head2Clbk(tree)
db = upernet_data_pipeline(broden_root)
metrics = partial(parts.BrodenMetricsClbk,
obj_tree=tree,
split_func=lambda o: (o[0], o[1]))
clbks = [clbk]
if not args.train_bn:
clbks.append(utils.BnFreeze(model.fpn))
learn = Learner(db,
model,
loss_func=cs_head.Head2Loss(tree),
callbacks=clbks,
callback_fns=metrics,
train_bn=args.train_bn)
learn.split((learn.model.embedding, ))
learn.freeze()
utils.fit_and_log(learn,
'object-P.A',
save=args.save,
epochs=20,
lr=1e-2,
pct_start=args.pct_start)
def fixup_upernet(root, tree, fix_type='add', n_iter=1):
encoder_path, decoder_path = utils.upernet_ckpt(root)
seg_model = upernet.get_upernet(tree, weights_encoder=encoder_path, weights_decoder=decoder_path)
fpn = upernet.extract_fpn(seg_model)
if fix_type == 'add':
fpn = addition_fixup(fpn)
model = FixUPAddModel(fpn, seg_model.decoder.object_head, tree.n_obj)
return model
elif fix_type == 'cat':
head, fpn = cat_fixup(fpn)
model = FixUPCatModel(head, fpn, seg_model.decoder.object_head, tree.n_obj, n_iter=n_iter)
return model
raise ValueError
def get_model(root, tree, op):
encoder_path, decoder_path = utils.upernet_ckpt(root)
instructor = Instructor(tree)
model = CSHead(instructor, tree, encoder_path, decoder_path, emb_op=op)
return model, instructor
def get_model(root, tree):
encoder_ckpt, decoder_ckpt = utils.upernet_ckpt(root)
fpn = get_fpn(tree, encoder_ckpt, decoder_ckpt)
outputs = {'object': tree.n_obj, 'part': tree.n_parts}
model = nn.Sequential(fpn, nnlayers.SplitHead(512, outputs))
return model