# fetch dataloaders
params_transfer.encoding = params_transfer.encoding_source
train_dl = dataloader.fetch_dataloader(
args.data_dir, args.txt_train, 'train', params_transfer)
val_dl_source = dataloader.fetch_dataloader(
args.data_dir, args.txt_val_source, 'val', params_transfer)
params_transfer.encoding = params_transfer.encoding_target
val_dl_target = dataloader.fetch_dataloader(
args.data_dir, args.txt_val_target, 'val', params_transfer)
logging.info("- done.")
# Define the model and optimizer
model_source = get_network(params_source).to(params_transfer.device)
model_target = get_network(params_target).to(params_transfer.device)
transfer = get_transfer(params_transfer).to(params_transfer.device)
#load source and target model before training and extract backbones
ckpt_source_file_path = os.path.join(args.checkpoint_dir_source, best_ckpt_filename)
print(ckpt_source_file_path)
if os.path.exists(ckpt_source_file_path):
model_source = utils.load_checkpoint(model_source, ckpt_dir=args.checkpoint_dir_source, filename=best_ckpt_filename, is_best=True)[0]
print("=> loaded source model checkpoint form {}".format(ckpt_source_file_path))
else:
print("=> Initializing source model from scratch")
ckpt_target_file_path = os.path.join(args.checkpoint_dir_target, best_ckpt_filename)
if os.path.exists(ckpt_target_file_path):
model_target = utils.load_checkpoint(model_target, ckpt_dir=args.checkpoint_dir_target, filename=best_ckpt_filename, is_best=True)[0]
print("=> loaded target model checkpoint form {}".format(ckpt_target_file_path))
args = parser.parse_args()
json_path = os.path.join(args.model_dir, 'params.json')
assert os.path.isfile(
json_path), "No json configuration file found at {}".format(json_path)
params = utils.Params(json_path)
# use GPU if available
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# Set the random seed for reproducible experiments
seed = 42
torch.manual_seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
np.random.seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
train_dl = data_loader.fetch_dataloader(args.data_dir, args.txt_train, 'train', params)
# Define the model and optimizer
model = get_network(params).to(device)
opt = optim.AdamW(model.parameters(), lr=params.learning_rate)
loss_fn = get_loss_fn(loss_name=params.loss_fn)
if args.checkpoint_dir:
model = utils.load_checkpoint(model, is_best=False, ckpt_dir=args.checkpoint_dir)[0]
log_lrs, losses = find_lr(train_dl, opt, model, loss_fn, device)
plot_lr(log_lrs, losses)