# !/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 math import random import os import datetime from optimization.training import train, evaluate from utils.load_data import load_dataset 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('-nc', '--no_cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--manual_seed', type=int, help='manual seed, if not given resorts to random seed.') parser.add_argument('-li', '--log_interval', type=int, default=20, metavar='LOG_INTERVAL', help='how many batches to wait before logging training status') parser.add_argument('--evaluate_interval_epochs', type=int, default=25, help='Evaluate per how many epochs') parser.add_argument('-od', '--out_dir', type=str, default='snapshots', metavar='OUT_DIR', help='output directory for model snapshots etc.') fp = parser.add_mutually_exclusive_group(required=False) fp.add_argument('-te', '--testing', action='store_true', dest='testing', help='evaluate on test set after training') fp.add_argument('-va', '--validation', action='store_false', dest='testing', help='only evaluate on validation set') parser.set_defaults(testing=True) # optimization settings parser.add_argument('-e', '--epochs', type=int, default=2000, metavar='EPOCHS', help='number of epochs to train (default: 2000)') parser.add_argument('-es', '--early_stopping_epochs', type=int, default=300, metavar='EARLY_STOPPING', help='number of early stopping epochs') parser.add_argument('-bs', '--batch_size', type=int, default=256, metavar='BATCH_SIZE', help='input batch size for training (default: 100)') parser.add_argument('-lr', '--learning_rate', type=float, default=0.001, metavar='LEARNING_RATE', help='learning rate') parser.add_argument('--warmup', type=int, default=10, help='number of warmup epochs') parser.add_argument('--data_augmentation_level', type=int, default=2, help='data augmentation level') parser.add_argument('--variable_type', type=str, default='discrete', help='variable type of data distribution: discrete/continuous', choices=['discrete', 'continuous']) parser.add_argument('--distribution_type', type=str, default='logistic', choices=['logistic', 'normal', 'steplogistic'], help='distribution type: logistic/normal') parser.add_argument('--n_flows', type=int, default=8, help='number of flows per level') parser.add_argument('--n_levels', type=int, default=3, help='number of levels') parser.add_argument('--n_bits', type=int, default=8, help='') # ---------------- SETTINGS CONCERNING NETWORKS ------------- parser.add_argument('--densenet_depth', type=int, default=8, help='Depth of densenets') parser.add_argument('--n_channels', type=int, default=512, help='number of channels in coupling and splitprior') # ---------------- ----------------------------- ------------- # ---------------- SETTINGS CONCERNING COUPLING LAYERS ------------- parser.add_argument('--coupling_type', type=str, default='shallow', choices=['shallow', 'resnet', 'densenet'], help='Type of coupling layer') parser.add_argument('--splitfactor', default=0, type=int, help='Split factor for coupling layers.') parser.add_argument('--split_quarter', dest='split_quarter', action='store_true', help='Split coupling layer on quarter') parser.add_argument('--no_split_quarter', dest='split_quarter', action='store_false') parser.set_defaults(split_quarter=True) # ---------------- ----------------------------------- ------------- # ---------------- SETTINGS CONCERNING SPLITPRIORS ------------- parser.add_argument('--splitprior_type', type=str, default='shallow', choices=['none', 'shallow', 'resnet', 'densenet'], help='Type of splitprior. Use \'none\' for no splitprior') # ---------------- ------------------------------- ------------- # ---------------- SETTINGS CONCERNING PRIORS ------------- parser.add_argument('--n_mixtures', type=int, default=1, help='number of mixtures') # ---------------- ------------------------------- ------------- parser.add_argument('--hard_round', dest='hard_round', action='store_true', help='Rounding of translation in discrete models. Weird ' 'probabilistic implications, only for experimental phase') parser.add_argument('--no_hard_round', dest='hard_round', action='store_false') parser.set_defaults(hard_round=True) parser.add_argument('--round_approx', type=str, default='smooth', choices=['smooth', 'stochastic']) parser.add_argument('--lr_decay', default=0.999, type=float, help='Learning rate') parser.add_argument('--temperature', default=1.0, type=float, help='Temperature used for BackRound. It is used in ' 'the the SmoothRound module. ' '(default=1.0') # gpu/cpu parser.add_argument('--gpu_num', type=int, default=0, metavar='GPU', help='choose GPU to run on.') args = parser.parse_args() args.cuda = not args.no_cuda and torch.cuda.is_available() if args.manual_seed is None: args.manual_seed = random.randint(1, 100000) random.seed(args.manual_seed) torch.manual_seed(args.manual_seed) np.random.seed(args.manual_seed) kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {} def run(args, kwargs): print('\nMODEL SETTINGS: \n', args, '\n') print("Random Seed: ", args.manual_seed) if 'imagenet' in args.dataset and args.evaluate_interval_epochs > 5: args.evaluate_interval_epochs = 5 # ================================================================================================================== # SNAPSHOTS # ================================================================================================================== args.model_signature = str(datetime.datetime.now())[0:19].replace(' ', '_') args.model_signature = args.model_signature.replace(':', '_') snapshots_path = os.path.join(args.out_dir, args.variable_type + '_' + args.distribution_type + args.dataset) snap_dir = snapshots_path snap_dir += '_' + 'flows_' + str(args.n_flows) + '_levels_' + str(args.n_levels) snap_dir = snap_dir + '__' + args.model_signature + '/' args.snap_dir = snap_dir if not os.path.exists(snap_dir): os.makedirs(snap_dir) with open(snap_dir + 'log.txt', 'a') as ff: print('\nMODEL SETTINGS: \n', args, '\n', file=ff) # SAVING torch.save(args, snap_dir + '.config') # ================================================================================================================== # LOAD DATA # ================================================================================================================== train_loader, val_loader, test_loader, args = load_dataset(args, **kwargs) # ================================================================================================================== # SELECT MODEL # ================================================================================================================== # flow parameters and architecture choice are passed on to model through args print(args.input_size) import models.Model as Model model = Model.Model(args) args.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.set_temperature(args.temperature) model.enable_hard_round(args.hard_round) model_sample = model # ==================================== # INIT # ==================================== # data dependend initialization on CPU for batch_idx, (data, _) in enumerate(train_loader): model(data) break if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = torch.nn.DataParallel(model, dim=0) model.to(args.device) def lr_lambda(epoch): return min(1., (epoch+1) / args.warmup) * np.power(args.lr_decay, epoch) optimizer = optim.Adamax(model.parameters(), lr=args.learning_rate, eps=1.e-7) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1) # ================================================================================================================== # TRAINING # ================================================================================================================== train_bpd = [] val_bpd = [] # for early stopping best_val_bpd = np.inf best_train_bpd = np.inf epoch = 0 train_times = [] model.eval() model.train() for epoch in range(1, args.epochs + 1): t_start = time.time() scheduler.step() tr_loss, tr_bpd = train(epoch, train_loader, model, optimizer, args) 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 < 25 or epoch % args.evaluate_interval_epochs == 0: v_loss, v_bpd = evaluate( train_loader, val_loader, model, model_sample, args, epoch=epoch, file=snap_dir + 'log.txt') val_bpd.append(v_bpd) # Model save based on TRAIN performance (is heavily correlated with validation performance.) if np.mean(tr_bpd) < best_train_bpd: best_train_bpd = np.mean(tr_bpd) best_val_bpd = v_bpd torch.save(model.module, snap_dir + 'a.model') torch.save(optimizer, snap_dir + 'a.optimizer') print('->model saved<-') print('(BEST: train bpd {:.4f}, test bpd {:.4f})\n'.format( best_train_bpd, best_val_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)) if __name__ == "__main__": run(args, kwargs)