Add files via upload

Author: ZZUTK

Adapt_Pix2Pix/checkGPU.py

@@ -0,0 +1,78 @@
+# check GPU status
+from pynvml import *
+from time import sleep
+from datetime import datetime
+
+
+def gpu_memory_usage(is_print=True):
+    try:
+        nvmlInit()
+        # version = nvmlSystemGetDriverVersion()
+        deviceCount = nvmlDeviceGetCount()
+        GPU = {}
+        for i in range(deviceCount):
+            GPU[i] = {}
+            handle = nvmlDeviceGetHandleByIndex(i)
+            info = nvmlDeviceGetMemoryInfo(handle)
+            GPU[i]['total'] = info.total / 1024.0 / 1024.0 / 1024.0
+            GPU[i]['free'] = info.free / 1024.0 / 1024.0 / 1024.0
+            if is_print:
+                print("\nGPU #%d Memory Usage:"
+                      "\n\tTotal:\t%4.2fGB\n\tFree:\t%4.2fGB" %
+                      (i, GPU[i]['total'], GPU[i]['free']))
+                print datetime.now()
+        nvmlShutdown()
+        return GPU
+    except:
+        print "Fail to check GPU status!"
+        exit(0)
+
+
+def auto_queue(gpu_memory_require=3.2, interval=1, schedule=None):
+    # input arg: schedule = datetime(year, month, day, hour, minute, second)
+    if schedule is None:
+        schedule = datetime.today()
+    else:
+        print '\nScheduled time: ', schedule
+
+    # wait until the scheduled time
+    now = datetime.today()
+    while now.year < schedule.year or now.month < schedule.month or now.day < schedule.day or \
+          now.hour < schedule.hour or now.minute < schedule.minute or now.second < schedule.second:
+        now = datetime.today()
+        sleep(interval)
+
+    gpu_stat = gpu_memory_usage()
+    if gpu_stat[0]['total'] < gpu_memory_require:
+        print 'Memory requirement is larger than GPU total memory'
+        exit(1)
+    while gpu_stat[0]['free'] < gpu_memory_require:
+        sleep(interval)  # second
+        gpu_stat = gpu_memory_usage()
+    return gpu_stat
+
+
+def set_memory_usage(usage=12.0, allow_growth=True):
+    auto_queue(gpu_memory_require=usage)
+    try:
+        import tensorflow as tf
+        assert type(usage) is int or float
+        assert usage >= 0
+       
+        config = tf.ConfigProto()
+        gpu_stat = gpu_memory_usage()
+        total_memory = gpu_stat[0]['total']
+        if usage > total_memory:
+            usage_percentage = 1.0
+        else:
+            usage_percentage = usage / total_memory
+        config.gpu_options.allow_growth = allow_growth
+        config.gpu_options.per_process_gpu_memory_fraction = usage_percentage
+        return config
+    except:
+        print 'Failed to set memory usage!'
+        return None
+
+
+if __name__ == '__main__':
+    gpu_memory_usage()

Adapt_Pix2Pix/datasets/facades/test/1.jpg

@@ -0,0 +1,9 @@
+mkdir datasets
+FILE=$1
+URL=https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/$FILE.tar.gz
+TAR_FILE=./datasets/$FILE.tar.gz
+TARGET_DIR=./datasets/$FILE/
+wget -N $URL -O $TAR_FILE
+mkdir $TARGET_DIR
+tar -zxvf $TAR_FILE -C ./datasets/
+rm $TAR_FILE

Adapt_Pix2Pix/datasets/facades/test/2.jpg

@@ -0,0 +1,82 @@
+import argparse
+import os
+import scipy.misc
+import numpy as np
+import checkGPU
+from model import pix2pix
+import tensorflow as tf
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--dataset_name', dest='dataset_name', default='facades', help='name of the dataset')
+parser.add_argument('--epoch', dest='epoch', type=int, default=200, help='# of epoch')
+parser.add_argument('--batch_size', dest='batch_size', type=int, default=1, help='# images in batch')
+parser.add_argument('--train_size', dest='train_size', type=int, default=1e8, help='# images used to train')
+parser.add_argument('--load_size', dest='load_size', type=int, default=286, help='scale images to this size')
+parser.add_argument('--fine_size', dest='fine_size', type=int, default=256, help='then crop to this size')
+parser.add_argument('--ngf', dest='ngf', type=int, default=64, help='# of gen filters in first conv layer')
+parser.add_argument('--ndf', dest='ndf', type=int, default=64, help='# of discri filters in first conv layer')
+parser.add_argument('--input_nc', dest='input_nc', type=int, default=3, help='# of input image channels')
+parser.add_argument('--output_nc', dest='output_nc', type=int, default=3, help='# of output image channels')
+parser.add_argument('--niter', dest='niter', type=int, default=200, help='# of iter at starting learning rate')
+parser.add_argument('--lr', dest='lr', type=float, default=0.0002, help='initial learning rate for adam')
+parser.add_argument('--beta1', dest='beta1', type=float, default=0.5, help='momentum term of adam')
+parser.add_argument('--flip', dest='flip', type=bool, default=True, help='if flip the images for data argumentation')
+parser.add_argument('--which_direction', dest='which_direction', default='AtoB', help='AtoB or BtoA')
+parser.add_argument('--phase', dest='phase', default='train', help='train, test')
+parser.add_argument('--save_epoch_freq', dest='save_epoch_freq', type=int, default=50, help='save a model every save_epoch_freq epochs (does not overwrite previously saved models)')
+parser.add_argument('--save_latest_freq', dest='save_latest_freq', type=int, default=5000, help='save the latest model every latest_freq sgd iterations (overwrites the previous latest model)')
+parser.add_argument('--print_freq', dest='print_freq', type=int, default=50, help='print the debug information every print_freq iterations')
+parser.add_argument('--continue_train', dest='continue_train', type=bool, default=False, help='if continue training, load the latest model: 1: true, 0: false')
+parser.add_argument('--serial_batches', dest='serial_batches', type=bool, default=False, help='f 1, takes images in order to make batches, otherwise takes them randomly')
+parser.add_argument('--serial_batch_iter', dest='serial_batch_iter', type=bool, default=True, help='iter into serial image list')
+parser.add_argument('--checkpoint_dir', dest='checkpoint_dir', default='checkpoint', help='models are saved here')
+parser.add_argument('--sample_dir', dest='sample_dir', default='sample', help='sample are saved here')
+parser.add_argument('--test_dir', dest='test_dir', default='test', help='test sample are saved here')
+parser.add_argument('--log_dir', dest='log_dir', default='log', help='log are saved here')
+parser.add_argument('--save_dir', dest='save_dir', default='save', help='save all')
+
+args = parser.parse_args()
+
+gpu_memory_require = 4.0
+
+
+def main(_):
+    args.checkpoint_dir = os.path.join(args.save_dir, args.checkpoint_dir)
+    args.sample_dir = os.path.join(args.save_dir, args.sample_dir)
+    args.test_dir = os.path.join(args.save_dir, args.test_dir)
+    args.log_dir = os.path.join(args.save_dir, args.log_dir)
+    if not os.path.exists(args.checkpoint_dir):
+        os.makedirs(args.checkpoint_dir)
+    if not os.path.exists(args.sample_dir):
+        os.makedirs(args.sample_dir)
+    if not os.path.exists(args.test_dir):
+        os.makedirs(args.test_dir)
+
+    checkGPU.auto_queue(
+        gpu_memory_require=gpu_memory_require,
+        interval=1,
+    )
+    config = checkGPU.set_memory_usage(
+        usage=gpu_memory_require,
+        allow_growth=True
+    )
+
+    with tf.Session(config=config) as sess:
+        model = pix2pix(sess, image_size=args.fine_size, batch_size=args.batch_size,
+                        output_size=args.fine_size, dataset_name=args.dataset_name,
+                        checkpoint_dir=args.checkpoint_dir, sample_dir=args.sample_dir)
+
+        if args.phase == 'train':
+            model.train(args)
+        else:
+            model.test(args)
+
+if __name__ == '__main__':
+
+    if 0:
+        print 'Run on CPU'
+        with tf.device("/cpu:0"):
+            gpu_memory_require = 0.0
+            tf.app.run()
+
+    tf.app.run()