now works with imagenet

Author: davidsketchdeck

.DS_Store

@@ -9,7 +9,7 @@ nohup*
 model
 samples
 results
-datasets
+datasets*
 logs
 dataset
 checkpoint

.gitignore

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2019 Junho Kim (1993.01.12)
+Copyright (c) 2019 David Mack
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

LICENSE

@@ -1,19 +1,30 @@
-# BigGAN-Tensorflow-TPU
+# BigGAN Tensorflow TPU
 
-**This is a half-finished TPU conversion of [Junho Kim's](https://github.com/taki0112/BigGAN-Tensorflow) implementation. Only 128x128 supported ATM**
-
-Simple Tensorflow implementation of ["Large Scale GAN Training for High Fidelity Natural Image Synthesis" (BigGAN)](https://arxiv.org/abs/1809.11096)
+Simple Tensorflow TPU implementation of ["Large Scale GAN Training for High Fidelity Natural Image Synthesis" (BigGAN)](https://arxiv.org/abs/1809.11096)
 
 ![main](./assets/main.png)
 
-## Issue
-* **The paper** used `orthogonal initialization`, but `I used random normal initialization.` The reason is, when using the orthogonal initialization, it did not train properly.
-* I have applied a hierarchical latent space, but **not** a class embeddedding.
+## Implementation notes/issues
+
+- This is a half-finished TPU conversion of [Junho Kim's](https://github.com/taki0112/BigGAN-Tensorflow) implementation. Only 128x128 supported
+- **The paper** used `orthogonal initialization`, but `I used random normal initialization.` The reason is, when using the orthogonal initialization, it did not train properly.
+- I have applied a hierarchical latent space, but **not** a class embeddedding.
 
 ## Usage
 
-### train
-* pipenv run ./launch_tpu_8.sh
+### Building the data
+
+For ImageNet, use [TensorFlow's build scripts](https://github.com/tensorflow/models/blob/master/research/inception/README.md#getting-started) to create TFRecord files of your chosen image size (e.g. 128x128). `--tfr-format inception`
+
+You can also use the data build script from [NVidia's Progressive Growing of GANs](https://github.com/tkarras/progressive_growing_of_gans). `--tfr-format progan`
+
+### Training
+
+You can train on a Google TPU by setting the name of your TPU as an env var and running one of the training scripts. For example,
+
+* `TPU_NAME=node-1 pipenv run ./launch_train_tpu_b128.sh`
+
+You need to have your training data stored on a Google cloud bucket.
 
 
 ## Architecture
@@ -28,5 +39,9 @@ Simple Tensorflow implementation of ["Large Scale GAN Training for High Fidelity
 ### 512x512
 <img src = './assets/512.png' width = '600px'> 
 
-## Author
+## Contributing
+
+You're very welcome to! Submit a PR or [contact the author(s)](https://octavian.ai)
+
+## Authors
 Junho Kim, David Mack

README.md

@@ -0,0 +1,137 @@
+
+
+from comet_ml import Experiment
+
+import tensorflow as tf
+
+import argparse
+import subprocess
+import os.path
+
+import logging
+import coloredlogs
+logger = logging.getLogger(__name__)
+
+from utils import *
+
+
+"""parsing and configuration"""
+def parse_args():
+	desc = "Tensorflow implementation of BigGAN"
+	parser = argparse.ArgumentParser(description=desc)
+	parser.add_argument('--tag'              , action="append" , default=[])
+	parser.add_argument('--phase'            , type=str        , default='train'                                           , help='train or test ?')
+	
+	parser.add_argument('--train-input-path' , type=str        , default='./datasets/imagenet/train*')
+	parser.add_argument('--eval-input-path'  , type=str        , default='./datasets/imagenet/validate*')
+	parser.add_argument('--tfr-format'       , type=str        , default='inception', choices=['inception', 'progan'])
+
+	parser.add_argument('--model-dir'        , type=str        , default='model')
+	parser.add_argument('--result-dir'       , type=str        , default='results')
+
+	# SAGAN
+	# batch_size = 256
+	# base channel = 64
+	# epoch = 100 (1M iterations)
+	# self-attn-res = [64]
+
+	parser.add_argument('--img-size'        , type=int             , default=128                               , help='The width and height of the input/output image')
+	parser.add_argument('--img-ch'          , type=int             , default=3                                 , help='The number of channels in the input/output image')
+
+	parser.add_argument('--epochs'          , type=int             , default=100                               , help='The number of training iterations')
+	parser.add_argument('--train-steps'     , type=int             , default=10000                             , help='The number of training iterations')
+	parser.add_argument('--eval-steps'      , type=int             , default=100                               , help='The number of eval iterations')
+	parser.add_argument('--batch-size'      , type=int             , default=2048  , dest="_batch_size"        , help='The size of batch across all GPUs')
+	parser.add_argument('--shuffle-buffer'  , type=int             , default=4000 )
+	
+
+	parser.add_argument('--ch'              , type=int             , default=96                                , help='base channel number per layer')
+	parser.add_argument('--layers'          , type=int             , default=5 )
+
+	parser.add_argument('--use-tpu'         , action='store_true')
+	parser.add_argument('--tpu-name'        , action='append'      , default=[] )
+	parser.add_argument('--tpu-zone'		, type=str, default='us-central1-f')
+	parser.add_argument('--num-shards'      , type=int             , default=8) # A single TPU has 8 shards
+	parser.add_argument('--steps-per-loop'  , type=int             , default=10000)
+
+	parser.add_argument('--disable-comet'   , action='store_false', dest='use_comet')
+
+	parser.add_argument('--self-attn-res'   , action='append', default=[] )
+
+	parser.add_argument('--g-lr'            , type=float           , default=0.00005                           , help='learning rate for generator')
+	parser.add_argument('--d-lr'            , type=float           , default=0.0002                            , help='learning rate for discriminator')
+
+	# if lower batch size
+	# g_lr = 0.0001
+	# d_lr = 0.0004
+
+	# if larger batch size
+	# g_lr = 0.00005
+	# d_lr = 0.0002
+
+	parser.add_argument('--beta1'          , type=float    , default=0.0           , help='beta1 for Adam optimizer')
+	parser.add_argument('--beta2'          , type=float    , default=0.9           , help='beta2 for Adam optimizer')
+	parser.add_argument('--moving-decay'   , type=float    , default=0.9999        , help='moving average decay for generator')
+
+	parser.add_argument('--z-dim'          , type=int      , default=128           , help='Dimension of noise vector')
+	parser.add_argument('--sn'             , type=str2bool , default=True          , help='using spectral norm')
+
+	parser.add_argument('--gan-type'       , type=str      , default='hinge'       , help='[gan / lsgan / wgan-gp / wgan-lp / dragan / hinge]')
+	parser.add_argument('--ld'             , type=float    , default=10.0          , help='The gradient penalty lambda')
+	parser.add_argument('--n-critic'       , type=int      , default=2             , help='The number of critic')
+
+	# IGoodfellow says sould be 50k
+	parser.add_argument('--inception-score-num'     , type=int      , default=512            , help='The number of sample images to use in inception score')
+	parser.add_argument('--sample-num'     , type=int      , default=36            , help='The number of sample images to save')
+	parser.add_argument('--test-num'       , type=int      , default=10            , help='The number of images generated by the test')
+
+	parser.add_argument('--verbosity', type=str, default='WARNING')
+
+	args = parser.parse_args()
+	return check_args(args)
+
+
+
+def check_args(args):
+	tf.gfile.MakeDirs(suffixed_folder(args, args.result_dir))
+	tf.gfile.MakeDirs("./temp/")
+
+	assert args.epochs >= 1, "number of epochs must be larger than or equal to one"
+	assert args._batch_size >= 1, "batch size must be larger than or equal to one"
+	assert args.ch >= 8, "--ch cannot be less than 8 otherwise some dimensions of the network will be size 0"
+
+	if args.use_tpu:
+		assert len(args.tpu_name) > 0, "Please provide at least one --tpu-name"
+
+	return args
+
+
+
+def model_dir(args):
+	return os.path.join(args.model_dir, *args.tag, model_name(args))
+
+
+
+
+
+def setup_logging(args):
+
+	coloredlogs.install(level='INFO', logger=logger)
+	coloredlogs.install(level='INFO', logger=logging.getLogger('main_tpu'))
+	coloredlogs.install(level='INFO', logger=logging.getLogger('ops'))
+	coloredlogs.install(level='INFO', logger=logging.getLogger('utils'))
+	coloredlogs.install(level='INFO', logger=logging.getLogger('BigGAN_128'))
+
+	tf.logging.set_verbosity(args.verbosity)
+
+	# log = logging.getLogger()
+	# log_path = os.path.join(suffixed_folder(args, args.result_dir), 'log.txt')
+	# stream = tf.gfile.Open(log_path, 'a')
+	# fh = logging.StreamHandler(stream=stream)
+	# fh.setLevel(logging.INFO)
+	# formatter = logging.Formatter('%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s')
+	# fh.setFormatter(formatter)
+	# log.addHandler(fh)
+
+	logger.info(f"cmd args: {vars(args)}")
+