def main():
import argparse, os
parser = argparse.ArgumentParser()
# Common arguments
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--encode_bs', type=int, default=64)
parser.add_argument('--black_box_jacobian_bs', type=int, default=None)
parser.add_argument('--cpu', action='store_true')
# Dataset arguments
parser.add_argument('--dataset', type=str, required=True)
parser.add_argument('--limit_dataset_size', type=int, default=None)
parser.add_argument('--cifar10_data_path', type=str, default='data')
parser.add_argument('--imagenet32_data_path',
type=str,
default='~/data/imagenet-small/valid_32x32.npy')
parser.add_argument('--imagenet64_data_path',
type=str,
default='~/data/imagenet-small/valid_64x64.npy')
# Model arguments
parser.add_argument('--cifar10_model',
type=str,
default='~/data/flowpp_cifar_model.npz')
parser.add_argument('--imagenet32_model',
type=str,
default='~/data/flowpp_imagenet32_model.npz')
parser.add_argument('--imagenet64_model',
type=str,
default='~/data/flowpp_imagenet64_model.npz')
# Script mode
parser.add_argument('--mode', type=str, required=True)
parser.add_argument('--test_output_filename', type=str, default=None)
parser.add_argument('--timing_test_count', type=int, default=6)
# Default compression options
parser.add_argument('--neg_log_noise_scale', type=int, default=14)
parser.add_argument('--disc_bits', type=int, default=32)
parser.add_argument('--disc_range', type=int, default=256)
parser.add_argument('--ans_init_bits', type=int, default=10000000)
parser.add_argument('--ans_num_streams', type=int, default=16)
parser.add_argument('--ans_mass_bits', type=int,
default=60) # probably never need to change this
args = parser.parse_args()
setup(seed=args.seed)
# Load data
if args.dataset == 'imagenet32':
model_ctor = compression.models.load_imagenet32_model
model_filename = os.path.expanduser(args.imagenet32_model)
dataset = load_imagenet_data(
os.path.expanduser(args.imagenet32_data_path))
elif args.dataset == 'imagenet64':
model_ctor = compression.models.load_imagenet64_model
model_filename = os.path.expanduser(args.imagenet64_model)
dataset = load_imagenet_data(
os.path.expanduser(args.imagenet64_data_path))
elif args.dataset == 'cifar10':
model_ctor = compression.models.load_cifar_model
model_filename = os.path.expanduser(args.cifar10_model)
dataset = CIFAR10WithoutLabels(
root=os.path.expanduser(args.cifar10_data_path),
train=False,
download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Lambda(
lambda x_: (x_ * 255).to(dtype=torch.int64)),
]))
else:
raise NotImplementedError(args.dataset)
dataloader, dataset = make_testing_dataloader(
dataset,
seed=args.seed,
limit_dataset_size=args.limit_dataset_size,
bs=args.encode_bs)
# Load model
device = torch.device('cpu' if args.cpu else 'cuda')
model = model_ctor(model_filename,
force_float32_cond=True).to(device=device)
# Dispatch to the chosen mode's main function
if args.mode == 'val_only':
num_datapoints_processed = main_val(dataloader, model, device)
assert num_datapoints_processed == len(dataset)
else:
def _make_stream(total_init_bits_=None):
return Bitstream(
device=device,
noise_scale=2**(-args.neg_log_noise_scale),
disc_bits=args.disc_bits,
disc_range=args.disc_range,
ans_mass_bits=args.ans_mass_bits,
ans_init_seed=0,
ans_init_bits=(
int(np.ceil(total_init_bits_ / args.ans_num_streams))
if total_init_bits_ is not None else args.
ans_init_bits # the --ans_init_bits argument value is the default
),
ans_num_streams=args.ans_num_streams)
if args.mode == 'test':
assert args.test_output_filename is not None
output = {
'args':
vars(args),
'results':
main_compression_test(
stream=_make_stream(),
model=model,
dataloader=dataloader,
device=device,
)
}
with open(args.test_output_filename, 'w') as f:
f.write(json.dumps(output) + '\n')
elif args.mode == 'timing_test_compositional':
batches = []
for (x_raw, ) in dataloader:
batches.append(x_raw)
if len(batches) >= args.timing_test_count:
break
main_timing_test(batches,
model=model,
stream=_make_stream(),
device=device)
elif args.mode == 'timing_test_blackbox':
assert args.black_box_jacobian_bs is not None
datapoints, = next(iter(dataloader))
datapoints = datapoints[:args.timing_test_count]
assert len(datapoints) == args.timing_test_count
from compression.logistic import force_accurate_mixlogistic_invcdf
with force_accurate_mixlogistic_invcdf():
main_timing_test_blackbox(
datapoints=datapoints.to(device=device),
model=model,
bbstream=BlackBoxBitstream(
model_ctor=model_ctor,
model_filename=model_filename,
device=device,
jacobian_bs=args.black_box_jacobian_bs))
else:
raise NotImplementedError(args.mode)
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset
from torchvision.transforms import transforms
from tqdm import tqdm
from encode import compress
from compression.utils import load_imagenet_data
from torchvision.utils import save_image
input = '/home/crhistyan/data/imagenet-small/valid_64x64.npy'
out_dir = 'real_image_size'
os.makedirs(out_dir, exist_ok=True)
dataset = load_imagenet_data(os.path.expanduser(input))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=False,
drop_last=False)
sizes = []
webp_sizes = []
flif_sizes = []
for i_batch, (x_raw, ) in enumerate(tqdm(dataloader)):
batch_size = 0
webp_batch_size = 0
flif_batch_size = 0
for i, tensor in enumerate(x_raw):
image_fn = os.path.join(out_dir, f'{i}.png')
save_image(tensor.type(torch.float32) / 255, image_fn)
def run(args):
print('\nMODEL SETTINGS: \n', args, '\n')
print("Random Seed: ", args.manual_seed)
# ==================================================================================================================
# SNAPSHOTS
# ==================================================================================================================
args.model_signature = str(datetime.datetime.now())[0:19].replace(' ', '_')
args.model_signature = args.model_signature.replace(':', '_')
os.makedirs(args.out_dir, exist_ok=True)
snap_dir = args.out_dir
with open(os.path.join(snap_dir, 'log.txt'), 'a') as ff:
print('\nMODEL SETTINGS: \n', args, '\n', file=ff)
# SAVING
torch.save(args, snap_dir + '.config')
# Load snapshot parameters
parameters_dict = None
if args.state_parameters is not None:
assert os.path.isfile(args.state_parameters)
parameters_dict = json.load(open(args.state_parameters))
args.learning_rate = parameters_dict['scheduler']['_last_lr'][0]
args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print('Device:', args.device)
# ==================================================================================================================
# LOAD DATA
# ==================================================================================================================
dataset = load_imagenet_data(os.path.expanduser(args.imagenet64_data_path))
validation_dataset = load_imagenet_data(os.path.expanduser(args.imagenet64_valid_data_path))
train_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, drop_last=False)
val_loader = torch.utils.data.DataLoader(validation_dataset, batch_size=args.batch_size, shuffle=True,
drop_last=False)
# test_loader = torch.utils.data.DataLoader(
# dataset,
# batch_size=args.batch_size,
# shuffle=False,
# **kwargs)
args.input_size = [3, 64, 64]
# ==================================================================================================================
# SELECT MODEL
# ==================================================================================================================
# flow parameters and architecture choice are passed on to model through args
print(args.input_size)
from compression.models.load_flowpp_imagenet64 import Imagenet64Model
# Load model
if args.imagenet64_model is None:
model = Imagenet64Model(force_float32_cond=True).eval()
else:
model_ctor = compression.models.load_imagenet64_model
model_filename = os.path.expanduser(args.imagenet64_model)
model = model_ctor(model_filename, force_float32_cond=True, from_torch=args.from_torch)
model.to(device=args.device)
model_sample = model
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma)
# ==================================================================================================================
# TRAINING
# ==================================================================================================================
train_bpd = []
val_bpd = []
# for early stopping
best_val_bpd = np.inf
best_val_loss = np.inf
if args.state_parameters is None:
last_epoch = 1
run_number = 1
else:
last_epoch = parameters_dict['epoch']
run_number = parameters_dict['run_number'] + 1
scheduler.load_state_dict(parameters_dict['scheduler'])
train_times = []
model.double()
for epoch in range(last_epoch, args.epochs + 1):
t_start = time.time()
if parameters_dict is not None:
tr_loss, tr_bpd = train(epoch, train_loader, model, optimizer, args, scheduler,
True, parameters_dict['batch_idx'], run_number)
else:
tr_loss, tr_bpd = train(epoch, train_loader, model, optimizer, args, scheduler, False)
train_bpd.append(tr_bpd)
train_times.append(time.time() - t_start)
print('One training epoch took %.2f seconds' % (time.time() - t_start))
if epoch < 5 or epoch % args.evaluate_interval_epochs == 0:
v_loss, v_bpd = evaluate(
val_loader, model, model_sample, args,
epoch=epoch, file=snap_dir + 'log.txt')
val_bpd.append(v_bpd)
best_val_bpd = min(v_bpd, best_val_bpd)
best_val_loss = min(v_loss, best_val_loss)
print('(BEST: val bpd {:.4f}, val loss {:.4f})\n'.format(best_val_bpd, best_val_loss))
print(f'VALIDATION: loss: {v_loss}, bpd: {v_bpd}')
if math.isnan(v_loss):
raise ValueError('NaN encountered!')
train_bpd = np.hstack(train_bpd)
val_bpd = np.array(val_bpd)
# training time per epoch
train_times = np.array(train_times)
mean_train_time = np.mean(train_times)
std_train_time = np.std(train_times, ddof=1)
print('Average train time per epoch: %.2f +/- %.2f' % (mean_train_time, std_train_time))
# ==================================================================================================================
# EVALUATION
# ==================================================================================================================
final_model = torch.load(snap_dir + 'a.model')
test_loss, test_bpd = evaluate(
train_loader, test_loader, final_model, final_model, args,
epoch=epoch, file=snap_dir + 'test_log.txt')
print('Test loss / bpd: %.2f / %.2f' % (test_loss, test_bpd))
def main():
parser = argparse.ArgumentParser()
# Common arguments
parser.add_argument('input', type=str)
parser.add_argument('--cpu', action='store_true')
parser.add_argument('--imagenet64_model', type=str, default='~/data/flowpp_imagenet64_model.npz')
parser.add_argument('-o', '--out_dir', type=str, default='lbb_compression')
parser.add_argument('--encode_out', type=str, default='image')
parser.add_argument('--decode_out', type=str, default='image.png')
parser.add_argument('--write_crops', action='store_true')
parser.add_argument('--single_image', action='store_true')
parser.add_argument('--batch_size', type=int, default=1)
# Default compression options
parser.add_argument('--neg_log_noise_scale', type=int, default=14)
parser.add_argument('--disc_bits', type=int, default=32)
parser.add_argument('--disc_range', type=int, default=256)
parser.add_argument('--ans_init_bits', type=int, default=10000000)
parser.add_argument('--ans_num_streams', type=int, default=1)
parser.add_argument('--ans_mass_bits', type=int, default=60) # probably never need to change this
args = parser.parse_args()
device = torch.device('cpu' if args.cpu else 'cuda')
args.device = device
print(f'Using {device}')
# Load model
model_fn = os.path.expanduser(args.imagenet64_model)
assert os.path.exists(model_fn), 'Model file not found'
model_ctor = compression.models.load_imagenet64_model
model = model_ctor(model_fn, force_float32_cond=False)
# Prepare output
args.out_dir = os.path.expanduser(args.out_dir)
os.makedirs(args.out_dir, exist_ok=True)
args.encode_out = os.path.join(args.out_dir, args.encode_out)
args.decode_out = os.path.join(args.out_dir, args.decode_out)
# Input files
input_fn = os.path.expanduser(args.input)
if args.single_image:
assert os.path.exists(input_fn), 'Input image not found'
print(f'Reading image from: {input_fn}')
dataset, bs = load_image_and_crop(input_fn)
dataloader, dataset = make_testing_dataloader(
dataset, seed=0, limit_dataset_size=0, bs=args.batch_size
)
else:
dataset = load_imagenet_data(os.path.expanduser(args.input))
print('Number of images:', len(dataset))
dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=False, drop_last=False)
for param in model.parameters():
param.requires_grad = False
# Prepare output file
out_fn = os.path.join(args.out_dir, 'disk_write_evaluation.csv')
file = open(out_fn, 'w+')
columns = ['iteration', 'image_number', 'file_size', 'bpd']
file.write(','.join(columns) + '\n')
file.close()
stream = make_stream(vars(args))
ok = encode(stream=stream, model=model, dataloader=dataloader, device=device,
write_crops=args.write_crops, output_fn=args.encode_out, eval_fn=out_fn, batch_size=args.batch_size)
if ok:
print('Encode success')
ok = decode(args.encode_out, model, write_crops=args.write_crops,
stream=make_stream({'ans_init_bits': 2}))
if ok:
print('Decode success')
def run(args):
print('\nMODEL SETTINGS: \n', args, '\n')
print("Random Seed: ", args.manual_seed)
# ==================================================================================================================
# SNAPSHOTS
# ==================================================================================================================
args.model_signature = str(datetime.datetime.now())[0:19].replace(' ', '_')
args.model_signature = args.model_signature.replace(':', '_')
os.makedirs(args.out_dir, exist_ok=True)
snap_dir = args.out_dir
with open(os.path.join(snap_dir, 'log.txt'), 'a') as ff:
print('\nMODEL SETTINGS: \n', args, '\n', file=ff)
# SAVING
torch.save(args, snap_dir + '.config')
# ==================================================================================================================
# LOAD DATA
# ==================================================================================================================
validation_dataset = load_imagenet_data(
os.path.expanduser(args.imagenet64_valid_data_path))
val_loader = torch.utils.data.DataLoader(validation_dataset,
batch_size=args.batch_size,
shuffle=True,
drop_last=False)
args.input_size = [3, 64, 64]
# ==================================================================================================================
# SELECT MODEL
# ==================================================================================================================
# flow parameters and architecture choice are passed on to model through args
print(args.input_size)
from compression.models.load_flowpp_imagenet64 import Imagenet64Model
args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
models_dir = os.path.expanduser(args.imagenet64_model_dir)
filenames = glob.glob(os.path.join(models_dir, '*.npz'))
out_fn = 'validation.csv'
file = open(out_fn, 'w+')
columns = ['val_bpd', 'val_loss']
file.write(','.join(columns) + '\n')
file.close()
val_bpd = []
val_loss = []
for model_filename in filenames:
# Load model
model_ctor = compression.models.load_imagenet64_model
model_filename = os.path.expanduser(model_filename)
model = model_ctor(model_filename,
force_float32_cond=True,
from_torch=True)
model.to(device=args.device)
print('Device:', args.device)
model_sample = model
model.double()
v_loss, v_bpd = evaluate(val_loader,
model,
model_sample,
args,
file=snap_dir + 'log.txt')
val_loss.append(v_loss)
val_bpd.append(v_bpd)
file = open(out_fn, 'a+')
v_bpd_str = f'{v_bpd}'
v_loss_str = f'{v_loss}'
columns = [v_bpd_str, v_loss_str]
file.write(','.join(columns) + '\n')
file.close()
print(f'VALIDATION: loss: {v_loss}, bpd: {v_bpd}')
if math.isnan(v_loss):
raise ValueError('NaN encountered!')
val_bpd = np.array(val_bpd)
val_loss = np.array(val_loss)
name = 'Validation BPD'
fig, ax = plt.subplots()
t = range(val_bpd.shape[0])
ax.plot(t, val_bpd)
ax.set(xlabel='batch_id', ylabel='', title=name)
plt.show()
fig.savefig(os.path.join(args.out_dir, f"{name}.png"))
name = 'Validation Loss'
fig, ax = plt.subplots()
t = range(val_loss.shape[0])
ax.plot(t, val_loss)
ax.set(xlabel='batch_id', ylabel='', title=name)
plt.show()
fig.savefig(os.path.join(args.out_dir, f"{name}.png"))
def main():
parser = argparse.ArgumentParser()
# Common arguments
parser.add_argument('input', type=str)
parser.add_argument('--cpu', action='store_true')
parser.add_argument('--from_torch', action='store_true')
parser.add_argument('--batch_size', type=int, default=2)
parser.add_argument('--imagenet64_model',
type=str,
default='~/data/flowpp_imagenet64_model.npz')
parser.add_argument('-o', '--out_dir', type=str, default='lbb_evaluation')
args = parser.parse_args()
device = torch.device('cpu' if args.cpu else 'cuda')
print(f'Using {device}')
# Load model
model_fn = os.path.expanduser(args.imagenet64_model)
assert os.path.exists(model_fn), 'Model file not found'
model_ctor = compression.models.load_imagenet64_model
model = model_ctor(model_fn,
force_float32_cond=False,
from_torch=args.from_torch)
for param in model.parameters():
param.requires_grad = False
# Prepare output
args.out_dir = os.path.expanduser(args.out_dir)
os.makedirs(args.out_dir, exist_ok=True)
# Input files
input_fn = os.path.expanduser(args.input)
assert os.path.exists(input_fn), 'Input file not found'
print(f'Reading images from: {input_fn}')
out_fn = os.path.join('evaluation.csv')
file = open(out_fn, 'w+')
columns = ['iteration', 'min_bits', 'stream_len']
file.write(','.join(columns) + '\n')
file.close()
batch_test = False
if batch_test:
# Read input image
dataset, num_crops = load_image_and_crop(input_fn)
for bs in range(1, 20):
dataloader, dataset = make_testing_dataloader(dataset,
seed=0,
limit_dataset_size=0,
bs=bs)
stream = lbb.make_stream({'ans_init_bits': 10000000})
ok = lbb.encode(stream=stream,
model=model,
dataloader=dataloader,
device=device,
output_fn='image',
write=False,
run_once=True)
dict_stream = stream.to_dict(orig=True)
file = open(out_fn, 'a+')
columns = [
str(bs),
str(dict_stream['ans'][0]['min_size']),
str(len(stream))
]
print(columns)
file.write(','.join(columns) + '\n')
file.close()
else:
dataset = load_imagenet_data(os.path.expanduser(args.input))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
drop_last=False)
stream = lbb.make_stream({'ans_init_bits': 10000000})
for i_batch, (x_raw, ) in enumerate(dataloader):
ok = lbb.encode_image(stream, model, device, i_batch, x_raw)
dict_stream = stream.to_dict(orig=True)
file = open(out_fn, 'a+')
columns = [
str(args.batch_size * (i_batch + 1)),
str(dict_stream['ans'][0]['min_size']),
str(len(stream))
]
print(columns)
file.write(','.join(columns) + '\n')
file.close()