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
# Bookkeeping
if (self.step_counter % self.log_interval == 1):
self.store_loss('weighted_compression_loss',
compression_model_loss.item())
if return_intermediates is True:
return losses, intermediates
else:
return losses
if __name__ == '__main__':
logger = utils.logger_setup(logpath=os.path.join(directories.experiments,
'logs'),
filepath=os.path.abspath(__file__))
device = utils.get_device()
logger.info(f'Using device {device}')
storage_train = defaultdict(list)
storage_test = defaultdict(list)
model = Model(hific_args,
logger,
storage_train,
storage_test,
model_type=ModelTypes.COMPRESSION_GAN)
model.to(device)
logger.info(model)
transform_param_names = list()
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,
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()
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))
from LightningDataset import LightningOpenImages
from default_config import mse_lpips_args
from LightningTrainer import HiFiC
from src.helpers import utils
import os
import pytorch_lightning as pl
args = mse_lpips_args
args = utils.setup_generic_signature(args, special_info=args.model_type)
logger = utils.logger_setup(logpath=os.path.join(args.snapshot, 'logs'),
filepath=os.path.abspath("train.py"))
dm = LightningOpenImages(args)
model = HiFiC(args)
trainer = pl.Trainer(gpus=[0],
amp_level='O2',
auto_scale_batch_size=True,
max_epochs=20,
limit_val_batches=2,
val_check_interval=0.5,
progress_bar_refresh_rate=20,
automatic_optimization=False,
precision=16,
benchmark=True,
resume_from_checkpoint=None,
amp_backend='apex',
num_sanity_val_steps=0)
trainer.fit(model, dm)