105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
from __future__ import print_function
|
|
import argparse
|
|
import time
|
|
import torch
|
|
import torch.utils.data
|
|
import torch.optim as optim
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from torchvision.utils import make_grid
|
|
|
|
import os
|
|
from optimization.training import train, evaluate
|
|
from utils.load_data import load_dataset
|
|
from utils.plotting import plot_training_curve
|
|
import imageio
|
|
|
|
|
|
parser = argparse.ArgumentParser(description='PyTorch Discrete Normalizing flows')
|
|
|
|
parser.add_argument('-d', '--dataset', type=str, default='cifar10', choices=['cifar10', 'imagenet32', 'imagenet64'],
|
|
metavar='DATASET',
|
|
help='Dataset choice.')
|
|
|
|
parser.add_argument('-bs', '--batch_size', type=int, default=256, metavar='BATCH_SIZE',
|
|
help='input batch size for training (default: 100)')
|
|
|
|
parser.add_argument('--data_augmentation_level', type=int, default=2,
|
|
help='data augmentation level')
|
|
|
|
parser.add_argument('-nc', '--no_cuda', action='store_true', default=False,
|
|
help='disables CUDA training')
|
|
|
|
|
|
args = parser.parse_args()
|
|
args.cuda = not args.no_cuda and torch.cuda.is_available()
|
|
|
|
kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
|
|
|
|
|
|
def run(args, kwargs):
|
|
|
|
args.snap_dir = snap_dir = \
|
|
'snapshots/discrete_logisticcifar10_flows_2_levels_3__2019-09-27_13_08_49/'
|
|
|
|
# ==================================================================================================================
|
|
# SNAPSHOTS
|
|
# ==================================================================================================================
|
|
|
|
# ==================================================================================================================
|
|
# LOAD DATA
|
|
# ==================================================================================================================
|
|
train_loader, val_loader, test_loader, args = load_dataset(args, **kwargs)
|
|
|
|
final_model = torch.load(snap_dir + 'a.model')
|
|
if hasattr(final_model, 'module'):
|
|
final_model = final_model.module
|
|
|
|
from models.backround import SmoothRound
|
|
for module in final_model.modules():
|
|
if isinstance(module, SmoothRound):
|
|
module._round_decay = 1.
|
|
|
|
exp_dir = snap_dir + 'partials/'
|
|
os.makedirs(exp_dir, exist_ok=True)
|
|
|
|
images = []
|
|
with torch.no_grad():
|
|
for data, _ in test_loader:
|
|
|
|
if args.cuda:
|
|
data = data.cuda()
|
|
|
|
for i in range(len(data)):
|
|
_, _, _, pz, z, pys, ys, ldj = final_model.forward(data[i:i+1])
|
|
|
|
for j in range(len(ys) + 1):
|
|
x_recon = final_model.inverse(
|
|
z,
|
|
ys[len(ys) - j:])
|
|
|
|
images.append(x_recon.float())
|
|
|
|
if i == 10:
|
|
break
|
|
break
|
|
|
|
for j in range(len(ys) + 1):
|
|
|
|
grid = make_grid(
|
|
torch.stack(images[j::len(ys) + 1], dim=0).squeeze(),
|
|
nrow=11, padding=0,
|
|
normalize=True, range=None,
|
|
scale_each=False, pad_value=0)
|
|
|
|
imageio.imwrite(
|
|
exp_dir + 'loaded{j}.png'.format(j=j),
|
|
grid.cpu().numpy().transpose(1, 2, 0))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
run(args, kwargs)
|