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 __init__(self, tree, weights_encoder='', weights_decoder=''):
super().__init__()
self.fpn = get_fpn(tree,
weights_encoder=weights_encoder,
weights_decoder=weights_decoder)
fpn_dim = 512
self.obj_branch = nn.ModuleList([
nnlayers.conv_layer(fpn_dim, fpn_dim),
conv2d(fpn_dim, tree.n_obj, ks=1, bias=True)
])
self.bu = nn.ModuleList([
nnlayers.conv_layer(tree.n_obj, fpn_dim // 2),
nnlayers.conv_layer(fpn_dim, fpn_dim)
])
self.part_branch = nn.Sequential(
nnlayers.conv_layer(fpn_dim, fpn_dim),
conv2d(fpn_dim, tree.n_parts, ks=1, bias=True))
self.lateral = nnlayers.conv_layer(fpn_dim, fpn_dim // 2)
def __init__(self,
instructor,
tree,
weights_encoder='',
weights_decoder='',
emb_op=torch.mul):
super().__init__()
self.fpn = get_fpn(tree,
weights_encoder=weights_encoder,
weights_decoder=weights_decoder)
fpn_dim = 512
self.td = nn.Sequential(
nnlayers.conv_layer(fpn_dim, fpn_dim // 4),
nnlayers.conv_layer(fpn_dim // 4, fpn_dim // 8),
fv.conv2d(fpn_dim // 8, 1, ks=1, bias=True))
self.embedding = fv.embedding(tree.n_obj, fpn_dim)
self.instructor = instructor
self.emb_op = emb_op
def __init__(self, tree, weights_encoder='', weights_decoder='', hidden=2):
super().__init__()
self.fpn = get_fpn(tree,
weights_encoder=weights_encoder,
weights_decoder=weights_decoder)
fpn_dim = 512
self.embedding = fv.embedding(tree.n_obj_with_parts + 1, fpn_dim)
self.td = nn.ModuleList(
[nnlayers.conv_layer(fpn_dim, fpn_dim) for _ in range(hidden)])
dims = tree.sections + [tree.n_obj]
self.heads = nn.ModuleList(
[fv.conv2d(fpn_dim, dim, ks=1, bias=True) for dim in dims])
self.bu_start = nn.ModuleList(
[fv.conv2d(dim, fpn_dim // 2) for dim in dims])
self.bu_lateral = nn.ModuleList([
nnlayers.conv_layer(fpn_dim, fpn_dim // 2) for _ in range(hidden)
])
self.bu = nn.ModuleList([
nnlayers.conv_layer(fpn_dim, fpn_dim // 2)
for _ in range(hidden - 1)
] + [nnlayers.conv_layer(fpn_dim, fpn_dim)])
self.obj_inst = tree.n_obj_with_parts
self.tree = tree
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