def main():
# LOAD ARGS
args = get_arguments()
print('Called with args:')
print(args)
assert args.cfg is not None, 'Missing cfg file'
cfg_from_file(args.cfg)
# auto-generate exp name if not specified
if cfg.EXP_NAME == '':
cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}'
if args.exp_suffix:
cfg.EXP_NAME += f'_{args.exp_suffix}'
# auto-generate snapshot path if not specified
if cfg.TRAIN.SNAPSHOT_DIR == '':
cfg.TRAIN.SNAPSHOT_DIR = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.EXP_NAME)
os.makedirs(cfg.TRAIN.SNAPSHOT_DIR, exist_ok=True)
# tensorboard
if args.tensorboard:
if cfg.TRAIN.TENSORBOARD_LOGDIR == '':
cfg.TRAIN.TENSORBOARD_LOGDIR = osp.join(cfg.EXP_ROOT_LOGS,
'tensorboard',
cfg.EXP_NAME)
os.makedirs(cfg.TRAIN.TENSORBOARD_LOGDIR, exist_ok=True)
if args.viz_every_iter is not None:
cfg.TRAIN.TENSORBOARD_VIZRATE = args.viz_every_iter
else:
cfg.TRAIN.TENSORBOARD_LOGDIR = ''
print('Using config:')
pprint.pprint(cfg)
# INIT
_init_fn = None
if not args.random_train:
torch.manual_seed(cfg.TRAIN.RANDOM_SEED)
torch.cuda.manual_seed(cfg.TRAIN.RANDOM_SEED)
np.random.seed(cfg.TRAIN.RANDOM_SEED)
random.seed(cfg.TRAIN.RANDOM_SEED)
def _init_fn(worker_id):
np.random.seed(cfg.TRAIN.RANDOM_SEED + worker_id)
if os.environ.get('ADVENT_DRY_RUN', '0') == '1':
return
# LOAD SEGMENTATION NET
assert osp.exists(
cfg.TRAIN.RESTORE_FROM), f'Missing init model {cfg.TRAIN.RESTORE_FROM}'
if cfg.TRAIN.MODEL == 'DeepLabv2':
model = get_deeplab_v2(num_classes=cfg.NUM_CLASSES,
multi_level=cfg.TRAIN.MULTI_LEVEL)
saved_state_dict = torch.load(cfg.TRAIN.RESTORE_FROM)
if 'DeepLab_resnet_pretrained_imagenet' in cfg.TRAIN.RESTORE_FROM:
new_params = model.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
model.load_state_dict(new_params)
else:
model.load_state_dict(saved_state_dict)
else:
raise NotImplementedError(f"Not yet supported {cfg.TRAIN.MODEL}")
print('Model loaded')
# DATALOADERS
source_dataset = GTA5DataSet(root=cfg.DATA_DIRECTORY_SOURCE,
list_path=cfg.DATA_LIST_SOURCE,
set=cfg.TRAIN.SET_SOURCE,
max_iters=cfg.TRAIN.MAX_ITERS *
cfg.TRAIN.BATCH_SIZE_SOURCE,
crop_size=cfg.TRAIN.INPUT_SIZE_SOURCE,
mean=cfg.TRAIN.IMG_MEAN)
source_loader = data.DataLoader(source_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_SOURCE,
num_workers=cfg.NUM_WORKERS,
shuffle=True,
pin_memory=True,
worker_init_fn=_init_fn)
target_dataset = CityscapesDataSet(root=cfg.DATA_DIRECTORY_TARGET,
list_path=cfg.DATA_LIST_TARGET,
set=cfg.TRAIN.SET_TARGET,
info_path=cfg.TRAIN.INFO_TARGET,
max_iters=cfg.TRAIN.MAX_ITERS *
cfg.TRAIN.BATCH_SIZE_TARGET,
crop_size=cfg.TRAIN.INPUT_SIZE_TARGET,
mean=cfg.TRAIN.IMG_MEAN)
target_loader = data.DataLoader(target_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_TARGET,
num_workers=cfg.NUM_WORKERS,
shuffle=True,
pin_memory=True,
worker_init_fn=_init_fn)
with open(osp.join(cfg.TRAIN.SNAPSHOT_DIR, 'train_cfg.yml'),
'w') as yaml_file:
yaml.dump(cfg, yaml_file, default_flow_style=False)
# UDA TRAINING
train_domain_adaptation(model, source_loader, target_loader, cfg)
def main():
# LOAD ARGS
args = get_arguments()
print('Called with args:')
print(args)
assert args.cfg is not None, 'Missing cfg file'
cfg_from_file(args.cfg)
cfg.NUM_WORKERS = args.num_workers
if args.option is not None:
cfg.TRAIN.OPTION = args.option
cfg.TRAIN.LAMBDA_BOUNDARY = args.LAMBDA_BOUNDARY
cfg.TRAIN.LAMBDA_DICE = args.LAMBDA_DICE
## gan method settings
cfg.GAN = args.gan
if cfg.GAN == 'gan':
cfg.TRAIN.LAMBDA_ADV_MAIN = 0.001 # GAN
elif cfg.GAN == 'lsgan':
cfg.TRAIN.LAMBDA_ADV_MAIN = 0.01 # LS-GAN
else:
raise NotImplementedError(f"Not Supported gan method")
### dataset settings
cfg.SOURCE = args.source
cfg.TARGET = args.target
## source config
if cfg.SOURCE == 'GTA':
cfg.DATA_LIST_SOURCE = str(project_root /
'advent/dataset/gta5_list/{}.txt')
cfg.DATA_DIRECTORY_SOURCE = str(project_root / 'data/GTA5')
cfg.TRAIN.INPUT_SIZE_SOURCE = (1280, 720)
elif cfg.SOURCE == 'SYNTHIA':
raise NotImplementedError(f"Not yet supported {cfg.SOURCE} dataset")
else:
raise NotImplementedError(f"Not yet supported {cfg.SOURCE} dataset")
## target config
if cfg.TARGET == 'Cityscapes':
cfg.DATA_LIST_TARGET = str(project_root /
'advent/dataset/cityscapes_list/{}.txt')
cfg.DATA_DIRECTORY_TARGET = str(project_root / 'data/cityscapes')
cfg.EXP_ROOT = project_root / 'experiments_G2C'
cfg.EXP_ROOT_SNAPSHOT = osp.join(cfg.EXP_ROOT, 'snapshots_G2C')
cfg.EXP_ROOT_LOGS = osp.join(cfg.EXP_ROOT, 'logs_G2C')
cfg.TRAIN.INPUT_SIZE_TARGET = (1024, 512)
cfg.TRAIN.INFO_TARGET = str(project_root /
'advent/dataset/cityscapes_list/info.json')
cfg.TEST.INPUT_SIZE_TARGET = (1024, 512)
cfg.TEST.OUTPUT_SIZE_TARGET = (2048, 1024)
cfg.TEST.INFO_TARGET = str(project_root /
'advent/dataset/cityscapes_list/info.json')
elif cfg.TARGET == 'BDD':
cfg.DATA_LIST_TARGET = str(project_root /
'advent/dataset/compound_list/{}.txt')
cfg.DATA_DIRECTORY_TARGET = str(project_root / 'data/bdd/Compound')
cfg.EXP_ROOT = project_root / 'experiments'
cfg.EXP_ROOT_SNAPSHOT = osp.join(cfg.EXP_ROOT, 'snapshots')
cfg.EXP_ROOT_LOGS = osp.join(cfg.EXP_ROOT, 'logs')
cfg.TRAIN.INPUT_SIZE_TARGET = (960, 540)
cfg.TRAIN.INFO_TARGET = str(project_root /
'advent/dataset/compound_list/info.json')
else:
raise NotImplementedError(f"Not yet supported {cfg.TARGET} dataset")
# auto-generate exp name if not specified
if cfg.EXP_NAME == '':
cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}_{cfg.TRAIN.OCDA_METHOD}'
if args.exp_suffix:
cfg.EXP_NAME += f'_{args.exp_suffix}'
# auto-generate snapshot path if not specified
if cfg.TRAIN.SNAPSHOT_DIR == '':
cfg.TRAIN.SNAPSHOT_DIR = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.EXP_NAME)
os.makedirs(cfg.TRAIN.SNAPSHOT_DIR, exist_ok=True)
# tensorboard
if args.tensorboard:
if cfg.TRAIN.TENSORBOARD_LOGDIR == '':
cfg.TRAIN.TENSORBOARD_LOGDIR = osp.join(cfg.EXP_ROOT_LOGS,
'tensorboard',
cfg.EXP_NAME)
os.makedirs(cfg.TRAIN.TENSORBOARD_LOGDIR, exist_ok=True)
if args.viz_every_iter is not None:
cfg.TRAIN.TENSORBOARD_VIZRATE = args.viz_every_iter
else:
cfg.TRAIN.TENSORBOARD_LOGDIR = ''
print('Using config:')
pprint.pprint(cfg)
# INIT
_init_fn = None
if not args.random_train:
torch.manual_seed(cfg.TRAIN.RANDOM_SEED)
torch.cuda.manual_seed(cfg.TRAIN.RANDOM_SEED)
np.random.seed(cfg.TRAIN.RANDOM_SEED)
random.seed(cfg.TRAIN.RANDOM_SEED)
def _init_fn(worker_id):
np.random.seed(cfg.TRAIN.RANDOM_SEED + worker_id)
if os.environ.get('ADVENT_DRY_RUN', '0') == '1':
return
# LOAD SEGMENTATION NET
if cfg.TRAIN.MODEL == 'DeepLabv2':
model = get_deeplab_v2(num_classes=cfg.NUM_CLASSES,
multi_level=cfg.TRAIN.MULTI_LEVEL)
saved_state_dict = torch.load(cfg.TRAIN.RESTORE_FROM)
if 'DeepLab_resnet_pretrained_imagenet' in cfg.TRAIN.RESTORE_FROM:
new_params = model.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
model.load_state_dict(new_params)
else:
model.load_state_dict(saved_state_dict)
elif cfg.TRAIN.MODEL == 'DeepLabv2_VGG':
model = get_deeplab_v2_vgg(cfg=cfg,
num_classes=cfg.NUM_CLASSES,
pretrained_model=cfg.TRAIN_VGG_PRE_MODEL)
if cfg.TRAIN.SELF_TRAINING:
path = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.TRAIN.RESTORE_FROM_SELF)
saved_state_dict = torch.load(path)
model.load_state_dict(saved_state_dict, strict=False)
trg_list = cfg.DATA_LIST_TARGET_ORDER
print("self-training model loaded: {} ".format(path))
else:
trg_list = cfg.DATA_LIST_TARGET
else:
raise NotImplementedError(f"Not yet supported {cfg.TRAIN.MODEL}")
print("model: ")
print(model)
print('Model loaded')
######## DATALOADERS ########
# GTA5: 24,966: 274,626 / 24,966 = 11 epoch
# self-training : target data shuffle
shuffle = cfg.TRAIN.SHUFFLE
if cfg.TRAIN.SELF_TRAINING:
max_iteration = None
else:
max_iteration = cfg.TRAIN.MAX_ITERS * cfg.TRAIN.BATCH_SIZE_SOURCE
source_dataset = GTA5DataSet(root=cfg.DATA_DIRECTORY_SOURCE,
list_path=cfg.DATA_LIST_SOURCE,
set=cfg.TRAIN.SET_SOURCE,
max_iters=max_iteration,
crop_size=cfg.TRAIN.INPUT_SIZE_SOURCE,
mean=cfg.TRAIN.IMG_MEAN)
source_loader = data.DataLoader(source_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_SOURCE,
num_workers=cfg.NUM_WORKERS,
shuffle=True,
pin_memory=True,
worker_init_fn=_init_fn)
if cfg.TARGET == "BDD":
# GTA5: 14,697: 264,546 / 14,697 = 18 epoch
target_dataset = BDDdataset(root=cfg.DATA_DIRECTORY_TARGET,
list_path=trg_list,
set=cfg.TRAIN.SET_TARGET,
info_path=cfg.TRAIN.INFO_TARGET,
max_iters=max_iteration,
crop_size=cfg.TRAIN.INPUT_SIZE_TARGET,
mean=cfg.TRAIN.IMG_MEAN)
target_loader = data.DataLoader(target_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_TARGET,
num_workers=cfg.NUM_WORKERS,
shuffle=shuffle,
pin_memory=True,
worker_init_fn=_init_fn)
elif cfg.TARGET == 'Cityscapes':
target_dataset = CityscapesDataSet(
root=cfg.DATA_DIRECTORY_TARGET,
list_path=cfg.DATA_LIST_TARGET,
set=cfg.TRAIN.SET_TARGET,
info_path=cfg.TRAIN.INFO_TARGET,
max_iters=cfg.TRAIN.MAX_ITERS * cfg.TRAIN.BATCH_SIZE_TARGET,
crop_size=cfg.TRAIN.INPUT_SIZE_TARGET,
mean=cfg.TRAIN.IMG_MEAN)
target_loader = data.DataLoader(target_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_TARGET,
num_workers=cfg.NUM_WORKERS,
shuffle=True,
pin_memory=True,
worker_init_fn=_init_fn)
else:
raise NotImplementedError(f"Not yet supported {cfg.TARGET} datasets")
with open(osp.join(cfg.TRAIN.SNAPSHOT_DIR, 'train_cfg.yml'),
'w') as yaml_file:
yaml.dump(cfg, yaml_file, default_flow_style=False)
# UDA TRAINING
train_domain_adaptation(model, source_loader, target_loader, cfg)