def generate(model_file_name, corpus, device, input_wsc=None, model=None):
model = model or load_model(model_file_name, device)
use_data_paralellization = True if type(
model).__name__ == 'CustomDataParallel' else False
model.eval()
batch_size = 1
hidden = model.init_hidden(batch_size)
word_frequency = torch.tensor(list(corpus.dictionary.word_count.values()),
dtype=torch.float)
if input_wsc is not None:
input_wsc_words = input_wsc.split()
input_word_id = (torch.tensor(
[[corpus.dictionary.word2idx[input_wsc_words[0]]]]).to(device))
else:
input_word_id = (torch.tensor(
[[torch.multinomial(word_frequency, 1)[0]]]).to(device))
input_words_probs = [(
corpus.dictionary.word_count[corpus.dictionary.idx2word[input_word_id]]
/ word_frequency.sum()).item()]
input_words = [corpus.dictionary.idx2word[input_word_id]]
number_of_words = (WORDS_TO_GENERATE
if input_wsc is None else len(input_wsc_words)) - 1
with torch.no_grad(): # no tracking history
for i in range(number_of_words):
if use_data_paralellization:
hidden, input_word_id = permute_for_parallelization(
hidden, input_word_id)
results = model(input_word_id, hidden)
outputs, hidden = get_results_from_data_parallelized_forward(
results, device)
hidden = permute_for_parallelization(hidden)
output = outputs[0]
else:
output, hidden = model(input_word_id, hidden)
word_probs = F.softmax(output.squeeze().div(TEMPERATURE), dim=0)
if input_wsc is None:
new_word_id = torch.multinomial(word_probs, 1)[0]
else:
new_word_id = corpus.dictionary.word2idx[input_wsc_words[i +
1]]
input_word_id.fill_(new_word_id)
input_words.append(corpus.dictionary.idx2word[new_word_id])
input_words_probs.append(word_probs[new_word_id].item())
return input_words, input_words_probs
def prepare_model(ckpt_path, input_dir):
make_deterministic()
device = utils.get_device()
logger = utils.logger_setup(logpath=os.path.join(input_dir, f'logs_{time.time()}'), filepath=os.path.abspath(__file__))
loaded_args, model, _ = utils.load_model(ckpt_path, logger, device, model_mode=ModelModes.EVALUATION,
current_args_d=None, prediction=True, strict=False, silent=True)
model.logger.info('Model loaded from disk.')
# Build probability tables
model.logger.info('Building hyperprior probability tables...')
model.Hyperprior.hyperprior_entropy_model.build_tables()
model.logger.info('All tables built.')
return model, loaded_args
args = utils.Struct(**args_d)
args = utils.setup_generic_signature(args, special_info=args.model_type)
args.target_rate = args.target_rate_map[args.regime]
args.lambda_A = args.lambda_A_map[args.regime]
args.n_steps = int(args.n_steps)
storage = defaultdict(list)
storage_test = defaultdict(list)
logger = utils.logger_setup(logpath=os.path.join(args.snapshot, 'logs'), filepath=os.path.abspath(__file__))
if args.warmstart is True:
assert args.warmstart_ckpt is not None, 'Must provide checkpoint to previously trained AE/HP model.'
logger.info('Warmstarting discriminator/generator from autoencoder/hyperprior model.')
if args.model_type != ModelTypes.COMPRESSION_GAN:
logger.warning('Should warmstart compression-gan model.')
args, model, optimizers = utils.load_model(args.warmstart_ckpt, logger, device,
model_type=args.model_type, current_args_d=dictify(args), strict=False, prediction=False)
else:
model = create_model(args, device, logger, storage, storage_test)
model = model.to(device)
amortization_parameters = itertools.chain.from_iterable(
[am.parameters() for am in model.amortization_models])
hyperlatent_likelihood_parameters = model.Hyperprior.hyperlatent_likelihood.parameters()
amortization_opt = torch.optim.Adam(amortization_parameters,
lr=args.learning_rate)
hyperlatent_likelihood_opt = torch.optim.Adam(hyperlatent_likelihood_parameters,
lr=args.learning_rate)
optimizers = dict(amort=amortization_opt, hyper=hyperlatent_likelihood_opt)
if model.use_discriminator is True:
def compress_and_decompress(args):
# Reproducibility
make_deterministic()
perceptual_loss_fn = ps.PerceptualLoss(model='net-lin', net='alex', use_gpu=torch.cuda.is_available())
# Load model
device = utils.get_device()
logger = utils.logger_setup(logpath=os.path.join(args.image_dir, 'logs'), filepath=os.path.abspath(__file__))
loaded_args, model, _ = utils.load_model(args.ckpt_path, logger, device, model_mode=ModelModes.EVALUATION,
current_args_d=None, prediction=True, strict=False)
# Override current arguments with recorded
dictify = lambda x: dict((n, getattr(x, n)) for n in dir(x) if not (n.startswith('__') or 'logger' in n))
loaded_args_d, args_d = dictify(loaded_args), dictify(args)
loaded_args_d.update(args_d)
args = utils.Struct(**loaded_args_d)
logger.info(loaded_args_d)
# Build probability tables
logger.info('Building hyperprior probability tables...')
model.Hyperprior.hyperprior_entropy_model.build_tables()
logger.info('All tables built.')
eval_loader = datasets.get_dataloaders('evaluation', root=args.image_dir, batch_size=args.batch_size,
logger=logger, shuffle=False, normalize=args.normalize_input_image)
n, N = 0, len(eval_loader.dataset)
input_filenames_total = list()
output_filenames_total = list()
bpp_total, q_bpp_total, LPIPS_total = torch.Tensor(N), torch.Tensor(N), torch.Tensor(N)
utils.makedirs(args.output_dir)
logger.info('Starting compression...')
start_time = time.time()
with torch.no_grad():
for idx, (data, bpp, filenames) in enumerate(tqdm(eval_loader), 0):
data = data.to(device, dtype=torch.float)
B = data.size(0)
input_filenames_total.extend(filenames)
if args.reconstruct is True:
# Reconstruction without compression
reconstruction, q_bpp = model(data, writeout=False)
else:
# Perform entropy coding
compressed_output = model.compress(data)
if args.save is True:
assert B == 1, 'Currently only supports saving single images.'
compression_utils.save_compressed_format(compressed_output,
out_path=os.path.join(args.output_dir, f"{filenames[0]}_compressed.hfc"))
reconstruction = model.decompress(compressed_output)
q_bpp = compressed_output.total_bpp
if args.normalize_input_image is True:
# [-1., 1.] -> [0., 1.]
data = (data + 1.) / 2.
perceptual_loss = perceptual_loss_fn.forward(reconstruction, data, normalize=True)
for subidx in range(reconstruction.shape[0]):
if B > 1:
q_bpp_per_im = float(q_bpp.cpu().numpy()[subidx])
else:
q_bpp_per_im = float(q_bpp.item()) if type(q_bpp) == torch.Tensor else float(q_bpp)
fname = os.path.join(args.output_dir, "{}_RECON_{:.3f}bpp.png".format(filenames[subidx], q_bpp_per_im))
torchvision.utils.save_image(reconstruction[subidx], fname, normalize=True)
output_filenames_total.append(fname)
bpp_total[n:n + B] = bpp.data
q_bpp_total[n:n + B] = q_bpp.data if type(q_bpp) == torch.Tensor else q_bpp
LPIPS_total[n:n + B] = perceptual_loss.data
n += B
df = pd.DataFrame([input_filenames_total, output_filenames_total]).T
df.columns = ['input_filename', 'output_filename']
df['bpp_original'] = bpp_total.cpu().numpy()
df['q_bpp'] = q_bpp_total.cpu().numpy()
df['LPIPS'] = LPIPS_total.cpu().numpy()
df_path = os.path.join(args.output_dir, 'compression_metrics.h5')
df.to_hdf(df_path, key='df')
pprint(df)
logger.info('Complete. Reconstructions saved to {}. Output statistics saved to {}'.format(args.output_dir, df_path))
delta_t = time.time() - start_time
logger.info('Time elapsed: {:.3f} s'.format(delta_t))
logger.info('Rate: {:.3f} Images / s:'.format(float(N) / delta_t))
def main(training, generating, model_file_name, quiet, use_bert):
verbose = not quiet
setup_torch()
# code seems to run slower (~90ms/batch, with batch_size=40) when default GPU is not cuda:0
device = torch.device("cuda:" + str(MAIN_GPU_INDEX) if USE_CUDA else "cpu")
corpus = get_corpus()
ntokens = len(corpus.dictionary)
sanity_checks(corpus, ntokens)
if training:
model = (
RNNModel(
MODEL_TYPE,
ntokens,
EMBEDDINGS_SIZE,
HIDDEN_UNIT_COUNT,
LAYER_COUNT,
DROPOUT_PROB,
TIED
).to(device)
)
criterion = nn.CrossEntropyLoss()
if USE_DATA_PARALLELIZATION:
cuda_devices = [i for i in range(torch.cuda.device_count())]
device_ids = [MAIN_GPU_INDEX] + cuda_devices[:MAIN_GPU_INDEX] + cuda_devices[MAIN_GPU_INDEX + 1:]
model = CustomDataParallel(model, device_ids=device_ids)
criterion = DataParallelCriterion(criterion, device_ids=device_ids)
# optimizer = torch.optim.Adam(model.parameters(), lr=INITIAL_LEARNING_RATE, weight_decay=1.2e-6)
# optimizer = torch.optim.SGD(model.parameters(), lr=INITIAL_LEARNING_RATE, weight_decay=1.2e-6)
optimizer = None
if verbose:
summary(model, criterion)
train(model, corpus, criterion, optimizer, device, USE_DATA_PARALLELIZATION)
else:
if not use_bert:
if model_file_name is None:
model_file_name = get_latest_model_file()
model = load_model(model_file_name, device)
if verbose:
log_loaded_model_info(model_file_name, model, device)
tokenizer = None
else:
# model_file_name = 'bert-base-multilingual-cased' if PORTUGUESE else 'bert-base-cased'
# model_file_name = 'bert-base-multilingual-cased' if PORTUGUESE else 'bert-large-cased'
# model_file_name = 'models/neuralmind/bert-base-portuguese-cased' if PORTUGUESE else 'bert-large-cased'
model_file_name = 'models/neuralmind/bert-large-portuguese-cased' if PORTUGUESE else 'bert-large-cased'
tokenizer = BertTokenizer.from_pretrained(model_file_name)
model = BertForNextSentencePrediction.from_pretrained(model_file_name)
if not generating:
logger.info('Generating WSC set, using model: {}'.format(model_file_name))
df = generate_df_from_json()
df = winograd_test(
df, corpus, model_file_name, device, model, tokenizer,
english=not PORTUGUESE, use_bert=use_bert
)
else:
logger.info('Generating text, using model: {}'.format(model_file_name))
words, words_probs = generate(model_file_name, corpus, device, model=model)
logger.info('Generated text: {}'.format((' ').join(words)))
