def get_vae(model_encoder,
model_decoder,
tokenizer_encoder,
tokenizer_decoder,
beta=1):
ArgsObj = namedtuple("Args", ["latent_size", "device", "fb_mode", "beta"])
args = ArgsObj(latent_size=LATENT_SIZE_LARGE,
device=get_device(),
fb_mode=0,
beta=beta)
checkpoint_full_dir = os.path.join(OUTPUT_DIR, "checkpoint-full-31250")
if not torch.cuda.is_available():
checkpoint = torch.load(os.path.join(checkpoint_full_dir,
"training.bin"),
map_location="cpu")
else:
checkpoint = torch.load(
os.path.join(checkpoint_full_dir, "training.bin"))
model_vae = VAE(model_encoder, model_decoder, tokenizer_encoder,
tokenizer_decoder, args)
model_vae.load_state_dict(checkpoint["model_state_dict"])
# logger.info("Pre-trained Optimus is successfully loaded")
model_vae.to(args.device)
return model_vae
def val_test(args):
writer = SummaryWriter('./logs/{0}'.format(args.output_folder))
save_filename = './models/{0}'.format(args.output_folder)
train_loader, valid_loader, test_loader = train_util.get_dataloaders(args)
recons_input_img = train_util.log_input_img_grid(test_loader, writer)
input_dim = 3
model = VAE(input_dim, args.hidden_size, args.enc_type, args.dec_type)
# if torch.cuda.device_count() > 1 and args.device == "cuda":
# model = torch.nn.DataParallel(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
discriminators = {}
if args.recons_loss == "gan":
recons_disc = Discriminator(input_dim, args.img_res,
args.input_type).to(args.device)
recons_disc_opt = torch.optim.Adam(recons_disc.parameters(),
lr=args.disc_lr,
amsgrad=True)
discriminators["recons_disc"] = [recons_disc, recons_disc_opt]
model.to(args.device)
for disc in discriminators:
discriminators[disc][0].to(args.device)
if args.weights == "load":
start_epoch = train_util.load_state(save_filename, model, optimizer,
discriminators)
else:
start_epoch = 0
stop_patience = args.stop_patience
best_loss = torch.tensor(np.inf)
for epoch in tqdm(range(start_epoch, 4), file=sys.stdout):
val_loss_dict, z = train_util.test(get_losses, model, valid_loader,
args, discriminators, True)
# if args.weights == "init" and epoch==1:
# epoch+=1
# break
# print(z.shape)
train_util.log_recons_img_grid(recons_input_img, model, epoch + 1,
args.device, writer)
train_util.log_interp_img_grid(recons_input_img, model, epoch + 1,
args.device, writer)
train_util.log_losses("val", val_loss_dict, epoch + 1, writer)
train_util.log_latent_metrics("val", z, epoch + 1, writer)
train_util.save_state(model, optimizer, discriminators,
val_loss_dict["recons_loss"], best_loss,
args.recons_loss, epoch, save_filename)
print(val_loss_dict)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--train_data_file",
default=None,
type=str,
required=True,
help="The input training data file (a text file).")
parser.add_argument(
"--eval_data_file",
default=None,
type=str,
help=
"An input evaluation data file to evaluate the perplexity on (a text file)."
)
parser.add_argument("--checkpoint_dir",
default=None,
type=str,
required=True,
help="The directory where checkpoints are saved.")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
parser.add_argument("--dataset",
default='Snli',
type=str,
help="The dataset.")
## Variational auto-encoder
parser.add_argument("--latent_size",
default=32,
type=int,
help="Latent space dimension.")
parser.add_argument("--total_sents",
default=10,
type=int,
help="Total sentences to test recontruction.")
parser.add_argument("--num_interpolation_steps",
default=10,
type=int,
help="Total sentences to test recontruction.")
parser.add_argument("--play_mode",
default="interpolation",
type=str,
help="interpolation or reconstruction.")
## Encoder options
parser.add_argument(
"--encoder_model_type",
default="bert",
type=str,
help="The encoder model architecture to be fine-tuned.")
parser.add_argument(
"--encoder_model_name_or_path",
default="bert-base-cased",
type=str,
help="The encoder model checkpoint for weights initialization.")
parser.add_argument(
"--encoder_config_name",
default="",
type=str,
help=
"Optional pretrained config name or path if not the same as model_name_or_path"
)
parser.add_argument(
"--encoder_tokenizer_name",
default="",
type=str,
help=
"Optional pretrained tokenizer name or path if not the same as model_name_or_path"
)
## Decoder options
parser.add_argument(
"--decoder_model_type",
default="gpt2",
type=str,
help="The decoder model architecture to be fine-tuned.")
parser.add_argument(
"--decoder_model_name_or_path",
default="bert-base-cased",
type=str,
help="The decoder model checkpoint for weights initialization.")
parser.add_argument(
"--decoder_config_name",
default="",
type=str,
help=
"Optional pretrained config name or path if not the same as model_name_or_path"
)
parser.add_argument(
"--decoder_tokenizer_name",
default="",
type=str,
help=
"Optional pretrained tokenizer name or path if not the same as model_name_or_path"
)
parser.add_argument("--per_gpu_train_batch_size",
default=1,
type=int,
help="Batch size per GPU/CPU for training.")
parser.add_argument("--per_gpu_eval_batch_size",
default=1,
type=int,
help="Batch size per GPU/CPU for evaluation.")
parser.add_argument('--gloabl_step_eval',
type=int,
default=661,
help="Evaluate the results at the given global step")
parser.add_argument(
"--max_seq_length",
default=512,
type=int,
help=
"Optional input sequence length before tokenization. The sequence will be dropped if it is longer the max_seq_length"
)
# Interact with users
parser.add_argument("--interact_with_user_input",
action='store_true',
help="Use user input to interact_with.")
parser.add_argument("--sent_source", type=str, default="")
parser.add_argument("--sent_target", type=str, default="")
parser.add_argument("--sent_input", type=str, default="")
parser.add_argument("--degree_to_target", type=float, default="1.0")
## Variational auto-encoder
parser.add_argument("--nz",
default=32,
type=int,
help="Latent space dimension.")
parser.add_argument("--prompt", type=str, default="")
parser.add_argument("--padding_text", type=str, default="")
parser.add_argument("--length", type=int, default=20)
parser.add_argument("--temperature", type=float, default=1.0)
parser.add_argument("--top_k", type=int, default=0)
parser.add_argument("--top_p", type=float, default=1.0)
parser.add_argument("--no_cuda",
action='store_true',
help="Avoid using CUDA when available")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
"--block_size",
default=-1,
type=int,
help="Optional input sequence length after tokenization."
"The training dataset will be truncated in block of this size for training."
"Default to the model max input length for single sentence inputs (take into account special tokens)."
)
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--use_philly",
action='store_true',
help="Use Philly for computing.")
args = parser.parse_args()
args.device = torch.device(
"cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
set_seed(args)
args.encoder_model_type = args.encoder_model_type.lower()
args.decoder_model_type = args.decoder_model_type.lower()
global_step = args.gloabl_step_eval
output_encoder_dir = os.path.join(
args.checkpoint_dir, 'checkpoint-encoder-{}'.format(global_step))
output_decoder_dir = os.path.join(
args.checkpoint_dir, 'checkpoint-decoder-{}'.format(global_step))
checkpoints = [[output_encoder_dir, output_decoder_dir]]
logger.info("Evaluate the following checkpoints: %s", checkpoints)
# Load a trained Encoder model and vocabulary that you have fine-tuned
encoder_config_class, encoder_model_class, encoder_tokenizer_class = MODEL_CLASSES[
args.encoder_model_type]
model_encoder = encoder_model_class.from_pretrained(
output_encoder_dir, latent_size=args.latent_size)
tokenizer_encoder = encoder_tokenizer_class.from_pretrained(
args.encoder_tokenizer_name
if args.encoder_tokenizer_name else args.encoder_model_name_or_path,
do_lower_case=args.do_lower_case)
model_encoder.to(args.device)
if args.block_size <= 0:
args.block_size = tokenizer_encoder.max_len_single_sentence # Our input block size will be the max possible for the model
args.block_size = min(args.block_size,
tokenizer_encoder.max_len_single_sentence)
# Load a trained Decoder model and vocabulary that you have fine-tuned
decoder_config_class, decoder_model_class, decoder_tokenizer_class = MODEL_CLASSES[
args.decoder_model_type]
model_decoder = decoder_model_class.from_pretrained(
output_decoder_dir, latent_size=args.latent_size)
tokenizer_decoder = decoder_tokenizer_class.from_pretrained(
args.decoder_tokenizer_name
if args.decoder_tokenizer_name else args.decoder_model_name_or_path,
do_lower_case=args.do_lower_case)
model_decoder.to(args.device)
if args.block_size <= 0:
args.block_size = tokenizer_decoder.max_len_single_sentence # Our input block size will be the max possible for the model
args.block_size = min(args.block_size,
tokenizer_decoder.max_len_single_sentence)
# Load full model
output_full_dir = os.path.join(args.checkpoint_dir,
'checkpoint-full-{}'.format(global_step))
checkpoint = torch.load(os.path.join(output_full_dir, 'training.bin'),
map_location=torch.device('cpu'))
# Chunyuan: Add Padding token to GPT2
special_tokens_dict = {
'pad_token': '<PAD>',
'bos_token': '<BOS>',
'eos_token': '<EOS>'
}
num_added_toks = tokenizer_decoder.add_special_tokens(special_tokens_dict)
print('We have added', num_added_toks, 'tokens to GPT2')
model_decoder.resize_token_embeddings(
len(tokenizer_decoder)
) # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e. the length of the tokenizer.
assert tokenizer_decoder.pad_token == '<PAD>'
# Evaluation
model_vae = VAE(model_encoder, model_decoder, tokenizer_encoder,
tokenizer_decoder, args)
model_vae.load_state_dict(checkpoint['model_state_dict'])
logger.info("Pre-trained Optimus is successfully loaded")
model_vae.to(args.device)
if args.interact_with_user_input:
if args.play_mode == 'interpolation':
if len(args.sent_source) > 0 and len(args.sent_source) > 0:
result = interpolate(model_vae, tokenizer_encoder,
tokenizer_decoder, args)
else:
print('Please check: specify the source and target sentences!')
if args.play_mode == 'analogy':
if len(args.sent_source) > 0 and len(args.sent_source) > 0 and len(
args.sent_input) > 0:
result = analogy(model_vae, tokenizer_encoder,
tokenizer_decoder, args)
else:
print(
'Please check: specify the source, target and input analogy sentences!'
)
else:
result = evaluate_latent_space(args,
model_vae,
tokenizer_encoder,
tokenizer_decoder,
prefix=global_step)
def main(args):
train_loader, val_loader, test_loader = train_util.get_dataloaders(args)
input_dim = 3
model = VAE(input_dim, args.hidden_size, args.enc_type, args.dec_type)
opt = torch.optim.Adam(model.parameters(), lr=LR, amsgrad=True)
discriminators = {}
if args.recons_loss != "mse":
if args.recons_loss == "gan":
recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device)
elif args.recons_loss == "comp":
recons_disc = AnchorComparator(input_dim*2, args.img_res, args.input_type).to(args.device)
elif "comp_2" in args.recons_loss:
recons_disc = ClubbedPermutationComparator(input_dim*2, args.img_res, args.input_type).to(args.device)
elif "comp_6" in args.recons_loss:
recons_disc = FullPermutationComparator(input_dim*2, args.img_res, args.input_type).to(args.device)
recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True)
discriminators["recons_disc"] = [recons_disc, recons_disc_opt]
if torch.cuda.device_count() > 1:
model = train_util.ae_data_parallel(model)
for disc in discriminators:
discriminators[disc][0] = torch.nn.DataParallel(discriminators[disc][0])
model.to(args.device)
model_name = f"vae_{args.recons_loss}"
if args.output_folder is None:
args.output_folder = os.path.join(model_name, args.dataset, f"depth_{args.enc_type}_{args.dec_type}_hs_{args.img_res}_{args.hidden_size}")
log_save_path = os.path.join("./logs", args.output_folder)
model_save_path = os.path.join("./models", args.output_folder)
if not os.path.exists(log_save_path):
os.makedirs(log_save_path)
print(f"log:{log_save_path}", file=sys.stderr)
sys.stderr.flush()
if not os.path.exists(model_save_path):
os.makedirs(model_save_path)
writer = SummaryWriter(log_save_path)
print(f"train loader length:{len(train_loader)}", file=sys.stderr)
best_loss = torch.tensor(np.inf)
if args.weights == "load":
start_epoch = train_util.load_state(model_save_path, model, opt, discriminators)
else:
start_epoch = 0
recons_input_img = train_util.log_input_img_grid(test_loader, writer)
train_util.save_recons_img_grid("val", recons_input_img, model, 0, args)
for epoch in range(1, args.num_epochs):
print("Epoch {}:".format(epoch))
train(model, opt, train_loader)
curr_loss = val(model, val_loader)
# val_loss_dict, z = train_util.test(get_losses, model, val_loader, args, discriminators)
print(f"epoch val loss:{curr_loss}", file=sys.stderr)
sys.stderr.flush()
train_util.save_recons_img_grid("val", recons_input_img, model, epoch+1, args)
train_util.save_interp_img_grid("val", recons_input_img, model, epoch+1, args)
'Training mode selection. Choices: mnist, synthetic_timeseries, cell_timeseries. (default: mnist)'
)
args = parser.parse_args()
model_filepath = "model-{}.pth".format(args.train_mode)
root_path = "results/{}".format(args.train_mode)
try:
os.makedirs(root_path)
except:
pass
is_cuda = not args.no_cuda
device = torch.device("cuda" if is_cuda else "cpu")
model = VAE(dropout=args.dropout,
input_dim=input_dims[args.train_mode]).to(device)
try:
model = load_model(model_filepath, model)
logger.info("Loading model from {}".format(model_filepath))
except:
logger.info("Creating VAE model from scratch")
model = VAE(dropout=args.dropout,
input_dim=input_dims[args.train_mode]).to(device)
if args.train_mode == 'mnist':
train_mnist(model, device, args.epochs, root_path)
elif args.train_mode == "synthetic_timeseries":
model.decoder.sigmoid = False # disable sigmoid from the final decoder layer
train_synthetic_timeseries(model, device, args.epochs, root_path)
elif args.train_mode == "cell_timeseries":
model.decoder.sigmoid = False # disable sigmoid from the final decoder layer
train_cell_timeseries(model, device, args.epochs, root_path)
model.to(torch.device("cpu"))
store_model(model_filepath, model)
def main(args):
input_dim = 3
model = VAE(input_dim, args.hidden_size, args.enc_type, args.dec_type)
opt = torch.optim.Adam(model.parameters(), lr=args.lr, amsgrad=True)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, "min", patience=args.lr_patience, factor=0.5,
# threshold=args.threshold, threshold_mode="abs", min_lr=1e-6)
# ae_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, "min", patience=args.lr_patience, factor=0.5,
# threshold=args.threshold, threshold_mode="abs", min_lr=1e-7)
discriminators = {}
if args.recons_loss != "mse":
if args.recons_loss == "gan":
recons_disc = Discriminator(input_dim, args.img_res,
args.input_type).to(args.device)
elif args.recons_loss == "comp":
recons_disc = AnchorComparator(input_dim * 2, args.img_res,
args.input_type).to(args.device)
elif "comp_2" in args.recons_loss:
recons_disc = ClubbedPermutationComparator(
input_dim * 2, args.img_res, args.input_type).to(args.device)
elif "comp_6" in args.recons_loss:
recons_disc = FullPermutationComparator(
input_dim * 2, args.img_res, args.input_type).to(args.device)
recons_disc_opt = torch.optim.Adam(recons_disc.parameters(),
lr=args.disc_lr,
amsgrad=True)
discriminators["recons_disc"] = [recons_disc, recons_disc_opt]
if torch.cuda.device_count() > 1:
model = train_util.ae_data_parallel(model)
for disc in discriminators:
discriminators[disc][0] = torch.nn.DataParallel(
discriminators[disc][0])
model.to(args.device)
for disc in discriminators:
discriminators[disc][0].to(args.device)
print("model built", file=sys.stderr)
#print("model created")
train_loader, val_loader, test_loader = train_util.get_dataloaders(args)
print("loaders acquired", file=sys.stderr)
#print("loaders acquired")
model_name = f"vae_{args.recons_loss}"
if args.output_folder is None:
args.output_folder = os.path.join(
model_name, args.dataset,
f"depth_{args.enc_type}_{args.dec_type}_hs_{args.img_res}_{args.hidden_size}"
)
log_save_path = os.path.join("./logs", args.output_folder)
model_save_path = os.path.join("./models", args.output_folder)
if not os.path.exists(log_save_path):
os.makedirs(log_save_path)
print(f"log:{log_save_path}", file=sys.stderr)
sys.stderr.flush()
if not os.path.exists(model_save_path):
os.makedirs(model_save_path)
writer = SummaryWriter(log_save_path)
print(f"train loader length:{len(train_loader)}", file=sys.stderr)
best_loss = torch.tensor(np.inf)
if args.weights == "load":
start_epoch = train_util.load_state(model_save_path, model, opt,
discriminators)
else:
start_epoch = 0
recons_input_img = train_util.log_input_img_grid(test_loader, writer)
train_util.log_recons_img_grid(recons_input_img, model, 0, args.device,
writer)
stop_patience = args.stop_patience
for epoch in range(start_epoch, args.num_epochs):
try:
train(model, train_loader, opt, epoch, writer, args,
discriminators)
except RuntimeError as err:
print("".join(
traceback.TracebackException.from_exception(err).format()),
file=sys.stderr)
print("*******", file=sys.stderr)
print(err, file=sys.stderr)
exit(0)
val_loss_dict, z = train_util.test(get_losses, model, val_loader, args,
discriminators)
print(f"epoch loss:{val_loss_dict['recons_loss'].item()}")
train_util.save_recons_img_grid("test", recons_input_img, model,
epoch + 1, args)
train_util.save_interp_img_grid("test", recons_input_img, model,
epoch + 1, args)
train_util.log_losses("val", val_loss_dict, epoch + 1, writer)
train_util.log_latent_metrics("val", z, epoch + 1, writer)
train_util.save_state(model, opt, discriminators,
val_loss_dict["recons_loss"], best_loss,
args.recons_loss, epoch, model_save_path)
