GAN support for DIGITS

digits/config/__init__.py

@@ -21,4 +21,3 @@ def config_value(option):
     Return the current configuration value for the given option
     """
     return option_list[option]
-

digits/config/tensorflow.py

@@ -3,26 +3,27 @@
 
 import os
 import platform
-from subprocess import Popen,PIPE
+from subprocess import Popen, PIPE
 
 from . import option_list
 
 VARNAME_ENV_TFPY = 'TENSORFLOW_PYTHON'
-DEFAULT_PYTHON_EXE = 'python2' # @TODO(tzaman) - use the python executable that was used to launch digits?
+DEFAULT_PYTHON_EXE = 'python2'  # @TODO(tzaman) - use the python executable that was used to launch digits?
 
 if platform.system() == 'Darwin':
     # DYLD_LIBRARY_PATH and LD_LIBRARY_PATH is sometimes stripped, and the cuda libraries might need it
-    if not "DYLD_LIBRARY_PATH" in os.environ:
+    if "DYLD_LIBRARY_PATH" not in os.environ:
         if "CUDA_HOME" in os.environ:
             os.environ["DYLD_LIBRARY_PATH"] = str(os.environ["CUDA_HOME"] + '/lib')
 
+
 def test_tf_import(python_exe):
     """
     Tests if tensorflow can be imported, returns if it went okay and optional error.
     """
     p = Popen([python_exe, "-c", "import tensorflow"], stdout=PIPE, stderr=PIPE)
     (out, err) = p.communicate()
-    return p.returncode==0, str(err)
+    return p.returncode == 0, str(err)
 
 if VARNAME_ENV_TFPY in os.environ:
     tf_python_exe = os.environ[VARNAME_ENV_TFPY]

digits/dataset/views.py

@@ -6,6 +6,7 @@
 
 from . import images as dataset_images
 from . import generic
+from digits import extensions
 from digits.utils.routing import job_from_request, request_wants_json
 from digits.webapp import scheduler
 
@@ -54,3 +55,31 @@ def summary():
         return generic.views.summary(job)
     else:
         raise werkzeug.exceptions.BadRequest('Invalid job type')
+
+
+@blueprint.route('/inference-form/<extension_id>/<job_id>', methods=['GET'])
+def inference_form(extension_id, job_id):
+    """
+    Returns a rendering of an inference form
+    """
+    inference_form_html = ""
+
+    if extension_id != "all-default":
+        extension_class = extensions.data.get_extension(extension_id)
+        if not extension_class:
+            raise RuntimeError("Unable to find data extension with ID=%s"
+                               % job_id.dataset.extension_id)
+        job = scheduler.get_job(job_id)
+        if hasattr(job, 'extension_userdata'):
+            extension_userdata = job.extension_userdata
+        else:
+            extension_userdata = {}
+        extension_userdata.update({'is_inference_db': True})
+        extension = extension_class(**extension_userdata)
+
+        form = extension.get_inference_form()
+        if form:
+            template, context = extension.get_inference_template(form)
+            inference_form_html = flask.render_template_string(template, **context)
+
+    return inference_form_html

digits/extensions/view/imageOutput/config_template.html

@@ -23,3 +23,9 @@
     {{ form.pixel_conversion.tooltip }}
     {{ form.pixel_conversion(class='form-control') }}
 </div>
+
+<div class="form-group{{mark_errors([form.show_input])}}">
+    {{ form.show_input.label }}
+    {{ form.show_input.tooltip }}
+    {{ form.show_input(class='form-control') }}
+</div>

digits/extensions/view/imageOutput/forms.py

@@ -45,3 +45,13 @@ class ConfigForm(Form):
         tooltip='Select method to convert pixel values to the target bit '
                 'range'
     )
+
+    show_input = utils.forms.SelectField(
+        'Show input as image',
+        choices=[
+            ('yes', 'Yes'),
+            ('no', 'No'),
+        ],
+        default='no',
+        tooltip='Show input as image'
+    )

digits/extensions/view/imageOutput/view.py

@@ -31,6 +31,7 @@ def __init__(self, dataset, **kwargs):
         self.channel_order = kwargs['channel_order'].upper()
         self.data_order = kwargs['data_order'].upper()
         self.normalize = (kwargs['pixel_conversion'] == 'normalize')
+        self.show_input = (kwargs['show_input'] == 'yes')
 
     @staticmethod
     def get_config_form():
@@ -70,17 +71,28 @@ def get_view_template(self, data):
           - context is a dictionary of context variables to use for rendering
           the form
         """
-        return self.view_template, {'image': digits.utils.image.embed_image_html(data)}
+        return self.view_template, {'image_input': digits.utils.image.embed_image_html(data[0]),
+                                    'image_output': digits.utils.image.embed_image_html(data[1])}
 
     @override
     def process_data(self, input_id, input_data, output_data):
         """
         Process one inference and return data to visualize
         """
 
-        data = output_data[output_data.keys()[0]].astype('float32')
+        if self.show_input:
+            data_input = input_data.astype('float32')
+            image_input = self.process_image(self.data_order, data_input)
+        else:
+            image_input = None
+
+        data_output = output_data[output_data.keys()[0]].astype('float32')
+        image_output = self.process_image(self.data_order, data_output)
+
+        return [image_input, image_output]
 
-        if self.data_order == 'HWC':
+    def process_image(self, data_order, data):
+        if data_order == 'HWC':
             data = (data.transpose((2, 0, 1)))
 
         # assume CHW at this point

digits/extensions/view/imageOutput/view_template.html

@@ -1,3 +1,6 @@
 {# Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved. #}
 
-<img src="{{image}}" style="max-width:100%;" />
+{% if image_input %}
+    <img src="{{image_input}}" style="max-width:100%;" />
+{% endif %}
+<img src="{{image_output}}" style="max-width:100%;" />

digits/frameworks/__init__.py

@@ -3,7 +3,6 @@
 
 from .caffe_framework import CaffeFramework
 from .framework import Framework
-from .tensorflow_framework import TensorflowFramework
 from .torch_framework import TorchFramework
 from digits.config import config_value
 
@@ -13,6 +12,10 @@
     'TorchFramework',
 ]
 
+if config_value('tensorflow')['enabled']:
+    from .tensorflow_framework import TensorflowFramework
+    __all__.append('TensorflowFramework')
+
 #
 #  create framework instances
 #

digits/frameworks/caffe_framework.py

@@ -153,7 +153,6 @@ def can_accumulate_gradients(self):
         if config_value('caffe')['flavor'] == 'BVLC':
             return True
         elif config_value('caffe')['flavor'] == 'NVIDIA':
-            return (parse_version(config_value('caffe')['version'])
-                    > parse_version('0.14.0-alpha'))
+            return (parse_version(config_value('caffe')['version']) > parse_version('0.14.0-alpha'))
         else:
             raise ValueError('Unknown flavor.  Support NVIDIA and BVLC flavors only.')

digits/frameworks/tensorflow_framework.py

@@ -1,23 +1,20 @@
 # Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
 from __future__ import absolute_import
 
-import numpy as np
 import os
 import re
 import subprocess
-import time
 import tempfile
 
-import flask
-
-from .errors import Error, NetworkVisualizationError, BadNetworkError
+from .errors import NetworkVisualizationError
 from .framework import Framework
 import digits
 from digits import utils
 from digits.config import config_value
 from digits.model.tasks import TensorflowTrainTask
 from digits.utils import subclass, override, constants
 
+
 @subclass
 class TensorflowFramework(Framework):
     """
@@ -35,7 +32,7 @@ class TensorflowFramework(Framework):
     SUPPORTS_PYTHON_LAYERS_FILE = False
     SUPPORTS_TIMELINE_TRACING = True
 
-    SUPPORTED_SOLVER_TYPES = ['SGD','ADADELTA','ADAGRAD','ADAGRADDA','MOMENTUM','ADAM','FTRL','RMSPROP']
+    SUPPORTED_SOLVER_TYPES = ['SGD', 'ADADELTA', 'ADAGRAD', 'ADAGRADDA', 'MOMENTUM', 'ADAM', 'FTRL', 'RMSPROP']
 
     SUPPORTED_DATA_TRANSFORMATION_TYPES = ['MEAN_SUBTRACTION', 'CROPPING']
     SUPPORTED_DATA_AUGMENTATION_TYPES = ['FLIPPING', 'NOISE', 'CONTRAST', 'WHITENING', 'HSV_SHIFTING']
@@ -50,7 +47,7 @@ def create_train_task(self, **kwargs):
         """
         create train task
         """
-        return TensorflowTrainTask(framework_id = self.framework_id, **kwargs)
+        return TensorflowTrainTask(framework_id=self.framework_id, **kwargs)
 
     @override
     def get_standard_network_desc(self, network):
@@ -126,10 +123,10 @@ def get_network_visualization(self, **kwargs):
         # Another for the HTML
         _, temp_html_path = tempfile.mkstemp(suffix='.html')
 
-        try: # do this in a try..finally clause to make sure we delete the temp file
+        try:  # do this in a try..finally clause to make sure we delete the temp file
             # build command line
             args = [config_value('tensorflow')['executable'],
-                    os.path.join(os.path.dirname(digits.__file__),'tools','tensorflow','main.py'),
+                    os.path.join(os.path.dirname(digits.__file__), 'tools', 'tensorflow', 'main.py'),
                     '--network=%s' % os.path.basename(temp_network_path),
                     '--networkDirectory=%s' % os.path.dirname(temp_network_path),
                     '--visualizeModelPath=%s' % temp_graphdef_path,
@@ -141,7 +138,7 @@ def get_network_visualization(self, **kwargs):
 
             if use_mean and use_mean != 'none':
                 mean_file = dataset.get_mean_file()
-                assert mean_file != None, 'Failed to retrieve mean file.'
+                assert mean_file is not None, 'Failed to retrieve mean file.'
                 args.append('--subtractMean=%s' % use_mean)
                 args.append('--mean=%s' % dataset.path(mean_file))
 
@@ -163,15 +160,14 @@ def get_network_visualization(self, **kwargs):
 
             env = os.environ.copy()
             # make only a selected number of GPUs visible. The ID is not important for just the vis
-            env['CUDA_VISIBLE_DEVICES'] = ",".join([str(i) for i in range(0,int(num_gpus))])
+            env['CUDA_VISIBLE_DEVICES'] = ",".join([str(i) for i in range(0, int(num_gpus))])
 
             # execute command
             p = subprocess.Popen(args,
-                        stdout=subprocess.PIPE,
-                        stderr=subprocess.STDOUT,
-                        close_fds=True,
-                        env=env
-                        )
+                                 stdout=subprocess.PIPE,
+                                 stderr=subprocess.STDOUT,
+                                 close_fds=True,
+                                 env=env)
 
             stdout_log = ''
             while p.poll() is None:
@@ -181,7 +177,7 @@ def get_network_visualization(self, **kwargs):
                         stdout_log += line
             if p.returncode:
                 raise NetworkVisualizationError(stdout_log)
-            else: # Success!
+            else:  # Success!
                 return repr(str(open(temp_graphdef_path).read()))
         finally:
             os.remove(temp_network_path)

digits/model/forms.py

@@ -121,13 +121,14 @@ def validate_py_ext(form, field):
         tooltip="How many epochs of training between running through one pass of the validation data?"
     )
 
-    traces_interval = utils.forms.IntegerField('Tracing Interval (in steps)',
-            validators=[
-                validators.NumberRange(min=0)
-            ],
-            default=0,
-            tooltip="Generation of a timeline trace every few steps"
-            )
+    traces_interval = utils.forms.IntegerField(
+        'Tracing Interval (in steps)',
+        validators=[
+            validators.NumberRange(min=0)
+        ],
+        default=0,
+        tooltip="Generation of a timeline trace every few steps"
+    )
 
     random_seed = utils.forms.IntegerField(
         'Random seed',
@@ -311,10 +312,11 @@ def validate_lr_multistep_values(form, field):
     )
 
     def validate_custom_network_snapshot(form, field):
-        if form.method.data == 'custom':
-            for filename in field.data.strip().split(os.path.pathsep):
-                if filename and not os.path.exists(filename):
-                    raise validators.ValidationError('File "%s" does not exist' % filename)
+        pass
+        #if form.method.data == 'custom':
+        #    for filename in field.data.strip().split(os.path.pathsep):
+        #        if filename and not os.path.exists(filename):
+        #            raise validators.ValidationError('File "%s" does not exist' % filename)
 
     # Select one of several GPUs
     select_gpu = wtforms.RadioField(

digits/model/images/classification/job.py

@@ -28,7 +28,7 @@ def job_type(self):
     def download_files(self, epoch=-1):
         task = self.train_task()
 
-        snapshot_filename = task.get_snapshot(epoch)
+        snapshot_filename = task.get_snapshot(epoch, download=True)
 
         # get model files
         model_files = task.get_model_files()