def main(opt : DictConfig):
os.chdir(hydra.utils.get_original_cwd())
# opt = CycleGANOptions()
tasks = ['./datasets/cityscapes', './datasets/maps', './datasets/facades', './datasets/vangogh2photo']
torch.manual_seed(0)
np.random.seed(0)
torch.cuda.manual_seed(0)
torch.cuda.manual_seed_all(0)
OmegaConf.set_struct(opt, False)
if opt.train:
start_task = 0
end_task = len(tasks)
opt.world_size = len(opt.gpu_ids) * opt.nodes
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '8888'
for task_idx in range(start_task, end_task):
# Create Task folder
opt.task_folder_name = "Task_"+str(task_idx+1)+"_"+tasks[task_idx][41:]+"_"+"cycleGAN"
opt.image_folder_name = "Intermediate_train_images"
if not os.path.exists(os.path.join(opt.checkpoints_dir, opt.task_folder_name, opt.image_folder_name)):
os.makedirs(os.path.join(opt.checkpoints_dir, opt.task_folder_name, opt.image_folder_name))
opt.ckpt_save_path = os.path.join(opt.checkpoints_dir, opt.task_folder_name)
opt.img_save_path = os.path.join(opt.checkpoints_dir, opt.task_folder_name, opt.image_folder_name)
if task_idx == 0:
netG_A_filter_list = []
netG_B_filter_list = []
weights_A = []
weights_B = []
else:
old_task_folder_name = "Task_"+str(task_idx)+"_"+tasks[task_idx-1][41:]+"_"+"cycleGAN"
print("Loading ", os.path.join(opt.checkpoints_dir, old_task_folder_name)+'/filters.pt')
filters = torch.load(os.path.join(opt.checkpoints_dir, old_task_folder_name)+'/filters.pt')
netG_A_filter_list = filters["netG_A_filter_list"]
netG_B_filter_list = filters["netG_B_filter_list"]
weights_A = filters["weights_A"]
weights_B = filters["weights_B"]
opt.netG_A_filter_list = netG_A_filter_list
opt.netG_B_filter_list = netG_B_filter_list
opt.weights_A = weights_A
opt.weights_B = weights_B
opt.dataroot = tasks[task_idx]
opt.task_num = task_idx+1
mp.spawn(train, nprocs=len(opt.gpu_ids), args=(opt,))
else:
'''
We will load the unconstrained filters and the weights ONLY from the last task.
This is because, after every task we store the unconstrined filter and weight
matrix of that task and all the previous ones. So we will only load from the last one
which will contain everything we need.
'''
print("In Testing mode")
start_task = 0
end_task = len(tasks)
load_filter_path = opt.checkpoints_dir+f"/Task_{len(tasks)}_{tasks[-1][41:]}_cycleGAN/filters.pt"
opt.load_filter_path = load_filter_path
filters = torch.load(opt.load_filter_path)
opt.netG_A_filter_list = filters["netG_A_filter_list"]
opt.netG_B_filter_list = filters["netG_B_filter_list"]
opt.weights_A = filters["weights_A"]
opt.weights_B = filters["weights_B"]
opt.image_folder_name = "Test_images"
for task_idx in range(start_task, end_task):
print(f"Task {task_idx+1}")
opt.task_folder_name = "Task_"+str(task_idx+1)+"_"+tasks[task_idx][41:]+"_"+"cycleGAN"
opt.img_save_path = os.path.join(opt.checkpoints_dir, opt.task_folder_name, opt.image_folder_name)
if not os.path.exists(os.path.join(opt.checkpoints_dir, opt.task_folder_name, opt.image_folder_name)):
os.makedirs(os.path.join(opt.checkpoints_dir, opt.task_folder_name, opt.image_folder_name))
opt.dataroot = tasks[task_idx]
opt.task_num = task_idx+1
test(opt, task_idx)
