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)))