if (cmd_args.gpu != 0) or (cmd_args.force_set_gpu is True):
torch.cuda.set_device(cmd_args.gpu)
if cmd_args.model_type == ModelTypes.COMPRESSION:
args = mse_lpips_args
elif cmd_args.model_type == ModelTypes.COMPRESSION_GAN:
args = hific_args
start_time = time.time()
device = utils.get_device()
# Override default arguments from config file with provided command line arguments
dictify = lambda x: dict((n, getattr(x, n)) for n in dir(x) if not (n.startswith('__') or 'logger' in n))
args_d, cmd_args_d = dictify(args), vars(cmd_args)
args_d.update(cmd_args_d)
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,
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))