def __init__(
self, tasks, tasks_in={}, tasks_out={},
pretrained=True,
freeze_list=[], direct_edges={}, lazy=False,
model_class='resnet_based', models_dir="./models"
):
super().__init__()
self.tasks = tasks
self.tasks += [task.base for task in self.tasks if hasattr(task, "base")]
self.tasks_in, self.tasks_out = tasks_in, tasks_out
self.pretrained = pretrained
self.edges_in, self.edges_out, = {}, {}
self.direct_edges = direct_edges
self.freeze_list = freeze_list
self.edge_map = {}
print('Creating graph with tasks:', self.tasks)
self.params = {}
transfer_models = model_types[model_class]
for task in self.tasks_out.get("edges", None):
key = str((task.name, "LS"))
model_type, path = transfer_models.get(task.name, {})["down"]
path = os.path.join(models_dir, path)
if not os.path.isfile(path):
path = None
transfer = Transfer(
task, task_configs.tasks.LS,
model_type=model_type, path=path
)
transfer.freezed = task in self.tasks_out.get("freeze")
self.edges_out[task.name] = transfer
self.edge_map[key] = transfer
try:
if not lazy:
transfer.load_model()
except Exception as e:
print(e)
IPython.embed()
if transfer.freezed:
transfer.set_requires_grad(False)
else:
self.params[key] = transfer
for task in self.tasks_in.get("edges", None):
key = str(("LS", task.name))
model_type, path = transfer_models.get(task.name, {})["up"]
path = os.path.join(models_dir, path)
if not os.path.isfile(path):
path = None
transfer = Transfer(
task_configs.tasks.LS, task,
model_type=model_type, path=path
)
transfer.freezed = task in self.tasks_in.get("freeze")
self.edges_in[task.name] = transfer
self.edge_map[key] = transfer
try:
if not lazy:
transfer.load_model()
except Exception as e:
print(e)
IPython.embed()
if transfer.freezed:
transfer.set_requires_grad(False)
else:
self.params[key] = transfer
# construct transfer graph
for src_task, dest_task in itertools.product(self.tasks, self.tasks):
key = str((src_task.name, dest_task.name))
transfer = None
if src_task==dest_task: continue
if isinstance(dest_task, RealityTask): continue
if src_task==task_configs.tasks.LS or dest_task==task_configs.tasks.LS:
continue
if isinstance(src_task, RealityTask):
transfer = RealityTransfer(src_task, dest_task)
self.edge_map[key] = transfer
elif key in self.direct_edges:
transfer = Transfer(src_task, dest_task, pretrained=pretrained)
transfer.freezed = key in self.freeze_list
try:
if not lazy: transfer.load_model()
except Exception as e:
print(e)
IPython.embed()
if transfer.model_type is None:
continue
if not transfer.freezed:
self.params[key] = transfer
else:
print("Setting link: " + str(key) + " not trainable.")
transfer.set_requires_grad(False)
else: continue
self.edge_map[key] = transfer
self.params = nn.ModuleDict(self.params)
def __init__(
self,
tasks=tasks,
realities=None,
edges=None,
edges_exclude=None,
pretrained=True,
finetuned=False,
reality=[],
task_filter=[tasks.segment_semantic, tasks.class_scene],
freeze_list=[],
lazy=False,
):
super().__init__(tasks=[])
self.tasks = list(set(tasks) - set(task_filter))
self.tasks += [
task.base for task in self.tasks if hasattr(task, "base")
]
self.edge_list, self.edge_list_exclude = edges, edges_exclude
self.pretrained, self.finetuned = pretrained, finetuned
self.edges, self.adj, self.in_adj = [], defaultdict(list), defaultdict(
list)
self.edge_map, self.reality = {}, reality
print('graph tasks!', self.tasks)
self.params = nn.ModuleDict()
self.realities = realities
# RGB -> Normal
transfer = Transfer(TASKS.rgb,
TASKS.normal,
pretrained=pretrained,
finetuned=finetuned)
transfer.name = get_transfer_name(transfer)
self.params[str((TASKS.rgb.name, TASKS.normal.name))] = transfer
self.edge_map[str((TASKS.rgb.name, TASKS.normal.name))] = transfer
self.edges += [transfer]
try:
if not lazy: transfer.load_model()
except:
print('Cound not load model:',
str((TASKS.rgb.name, TASKS.normal.name)))
IPython.embed()
# RGB -> Depth
transfer = Transfer(TASKS.rgb,
TASKS.depth_zbuffer,
pretrained=pretrained,
finetuned=finetuned)
transfer.name = get_transfer_name(transfer)
self.params[str((TASKS.rgb.name, TASKS.depth_zbuffer.name))] = transfer
self.edge_map[str(
(TASKS.rgb.name, TASKS.depth_zbuffer.name))] = transfer
try:
if not lazy: transfer.load_model()
except:
print('Cound not load model:',
str((TASKS.rgb.name, TASKS.depth_zbuffer.name)))
IPython.embed()
# Depth -> Normals
src_task = (TASKS.depth_zbuffer.name, TASKS.FoV.name,
TASKS.normal.name)
target_task = TASKS.normal
transfer_name = str((src_task, target_task.name))
model_type, path = pretrained_transfers[(src_task, target_task.name)]
transfer = Transfer(src_task,
target_task,
pretrained=pretrained,
model_type=model_type,
path=path,
checkpoint=False,
finetuned=finetuned,
name=f"{src_task}2{target_task.name}")
transfer.name = transfer_name
self.params[transfer_name] = transfer
self.edge_map[transfer_name] = transfer
try:
if not lazy: transfer.load_model()
except:
print('Cound not load model:', transfer_name)
IPython.embed()
# Normal -> Depth
src_task = (TASKS.depth_zbuffer.name, TASKS.FoV.name,
TASKS.normal.name)
target_task = TASKS.depth_zbuffer
transfer_name = str((src_task, target_task.name))
model_type, path = pretrained_transfers[(src_task, target_task.name)]
transfer = Transfer(src_task,
target_task,
pretrained=pretrained,
model_type=model_type,
path=path,
checkpoint=False,
finetuned=finetuned,
name=f"{src_task}2{target_task.name}")
transfer.name = transfer_name
self.params[transfer_name] = transfer
self.edge_map[transfer_name] = transfer
try:
if not lazy: transfer.load_model()
except:
print('Cound not load model:', transfer_name)
IPython.embed()
for src_task, dest_task in itertools.product(self.realities,
self.tasks):
key = (src_task, dest_task)
if edges is not None and key not in edges: continue
if edges_exclude is not None and key in edges_exclude: continue
if src_task == dest_task: continue
if isinstance(dest_task, RealityTask): continue
transfer = None
if isinstance(src_task, RealityTask):
if dest_task not in src_task.tasks: continue
transfer = RealityTransfer(src_task, dest_task)
else:
transfer = Transfer(src_task,
dest_task,
pretrained=pretrained,
finetuned=finetuned)
transfer.name = get_transfer_name(transfer)
if transfer.model_type is None:
continue
self.edges += [transfer]
self.adj[src_task.name] += [transfer]
self.in_adj[dest_task.name] += [transfer]
self.edge_map[str((src_task.name, dest_task.name))] = transfer
if isinstance(transfer, nn.Module):
if str((src_task.name, dest_task.name)) not in freeze_list:
self.params[str(
(src_task.name, dest_task.name))] = transfer
else:
print("freezing " + str((src_task.name, dest_task.name)))
try:
if not lazy: transfer.load_model()
except:
print('Cound not load model:',
str((src_task.name, dest_task.name)))
IPython.embed()