Script changes for tensorflow (#1)

* Basic Tensorflow Support Added some initial tf tools Implemented UI Fixes for tensorflow 0.10 Removed tf-slim as its not part of the 0.10 master Added the lmdb reader with a tf.cond that needs replacement Implemented train and val seperation with a templating Fixed issue with dequeueing both runners by pulling both graphs Implemented training and validation rythm Added support for both png and jpg and added 16 bit support Implemented mean subtraction - but needs rework to load as constant Added an optimized implementation of mean subtraction Further optimized the mean loading by using a shared constant Wrapped the data loader in a factory to easily support more data types Implemented cropping Implemented floating point support. Implemented seperate LMDB database. Implemented regression support. Added some brief nosetests. Need to invoke accuracy only on classification though. Implemented variable restoration. Needs thorough testing Implemented inferencing, not entirely polished Moved some code into functions, started on modularization a bit Implemented digits custom helper functions Implemented custom printing ops Implemented autoencoder total rewrite of summaries Implemented output to console from scalar summaries Fixes for summary outputs: only simple scalar values are parsed to console Implemented binary segmentation and necessary fixes Some updates on binary seg Implemented all possible optimizers and started work on learning rate shaper Started work on the lr policies Fixes for learning_rates, implemented optimizers, tested variable summary output to UI Implemented and tested all learning rates and optimizers Introduces new model definition and improvements in loss handling and graph layout Major refactoring of main code. Implemented new model description. Implemented and tested inferencing. Implemented and tested weight/snapshot loading. All-round minor updates and fixes Fixes in summary cumulator and implemented an RNN model Fixes for mean subtraction in tf and tf-ui, implemented data order selection in image-view extension Implemented support for mean file of format: png, jpg, binaryproto - the latter being the fault that DIGITS will provide. Added support for runtime statistics and some allround fixes Added static tensorboard style network visualization for tensorflow. Added output of traces (no vis yet). Added a loader while waiting for network vis. Minor syntax cleanups. Implemented alexnet standard network Pulled in updates for travis build and added tensorflow install Added two more files for Mr Travis Implemented tensorflow configuration Added tf config to doc Fixes for ubuntu deployment of tf. Moved tf tools Fixes for tf ubuntu Fixes for tf ubuntu Some fixes and updates for TF in Travis Fix in network viz test Implemented default sinlge-gpu support and some nosetests Fixes for inference Added siamese network, bugfixes, minor features, some utility tf functions Added siamese network and example png Better error-ui format for network viz Added an alternative simpler siamese network that doesnt need a seperate db, minor error update Preliminary version of hdf5 implemented Implemented fine-tuning by renaming variables Implemented visualisation of variables and the activations of the Ops they belong to. Fix in inf vis naming Fixes in visualualisation shapes and naming Implemented softmax upon classification Implemented all nosetests for tf classification, and many allround bugfixes Implemented generic nosetests - some need work Fix for travis to find python exe Implemented a better file format deducer, and implemented a bare minimal TFRecord-reader Added top_n accuracy shortcut Implemented on-line data augmentation for TF, 5 types. Some minor bugfixes. Need to do something with image whitening though during validation and inf.. Added tensorflow data augmentation test Minor fixes and improvements from linter Implemented minimal and bare multigpu and fixes to get it running for greg Preliminary version of tfrecord writer for classification Some changes to optimize dataloading for tfr More fixes for tfrecrods Fix generic data loading Minor breaking changes but updates in namescoping Implemented new model structure. Improvements to multi-gpu handling. Updates to namespaces. Implemented accounting for regularization. Many allround updates Implemented proper visualisation for gpu devices Minor updates and converted alexnet and vgg16 to new format Fix in tfrecord shape WIP on timeline traces Finalized support for tensorflow timeline traces Fixed alexnet for tf Fix merge errors Minify tf-graph-basic.build.js * bAbI data plug-in Add utils Add inference form to bAbI dataset Allow inference without answer Allow unknown words in BaBI data plug-in Fix bAbI plugin Lint errors * Tensorflow integration updates Use TFRecords for TF inference TF: Don't rescale inputs Fix some TF classification tests Remove unnecessary print Fix TF imports when uninstalled Fix mean image scale Fix generic model tests Fix Torch single image inference Fix inference TMP TF Lint Revert changes in digits-lint script Lint: ignore tensorflow standard examples More Lint fixes * Add gradient hook * Add memn2n model * Update memn2n with gradient hooks * GAN example * Make batch size variable * Training/inference paths * Small update to TF 0.12 * Snapshot names, float inference, restore all vars * Do not restore global_step or optimizer variables * Add TB link * Update GAN network * Dynamically select inference form * TF inference: convert images to float * Update GAN z-gen network * Small Update model view layout * Add GAN plug-ins * Update GAN plug-in to create CelebA dataset * Add ability to show input in ImageOutput extension * Add all data to raw data view extension * Add model for CelebA dataset * Update GAN data plug-in * Update all losses in one session * Remove conversion to .png in GAN data plug-in * TF Slim Lenet example Divide input by 255 * Update GAN data plug-in * Fix TF model snapshot * Reduce scheduler delays to speed up inference * Update GAN plugins * Fix TF tests * Add API to LmdbReader (used by gan_features.py) * Save animated gif * Add GAN walk-through * Update GAN walkthrough with embeddings video * Fix GAN view for list encoding * Add animation task to GAN plugins * Add view task to see image attributes * Add comments to GAN models * Update README * Fix GAN features script * GAN app * Fix DIGITS inference * Adjust GAN window size automatically * Add attributes to GAN app * Move gandisplay.py * Remove wxpython 3.0 selection * Fix call to model * Adding disclaimer * Ported DIGITS to using tensorflow 1.1.0. * Ported DIGITS to using tensorflow 1.1.0. Got master branch working * updated gitignore * first cherrypick for installation scripts * Tf install experimental (#2) * Fix visualization when palette is None (NVIDIA#1177) The palette may be `None`when working with grayscale labels. Fix NVIDIA#1147 * Bugfix for customizing previous models (NVIDIA#1202) * [Packaging] Disable tests (NVIDIA#1227) * [Tests] Skip if extension not installed (NVIDIA#1263) * [Docs] Fix spelling errors in comments * [Docs] Add note about torch pkg and cusparse (NVIDIA#1303) * [Docs] Add note about torch pkg and cusparse (NVIDIA#1303) * [Caffe] Fix batch accumulation bug (NVIDIA#1307) * Use official NVIDIA model store by default (NVIDIA#1308) * Mark v5.0.0 * [Packaging] Pull latest docker image before build * Add .pgm to list of supported image file formats * Restrict usage of cmap to labels DB in generic dataset exploration fix NVIDIA#1322 * Update Object Detection example doc (NVIDIA#1323) * Update Object Detection example doc (NVIDIA#1323) * [TravisCI] Cache local OpenBLAS build This fixes a Torch bug we've been having on Travis for a while now. We had only been building OpenBLAS from source when there was no cached torch build present on the build machine. That meant you could get a cached build of Torch which was built against one version of OpenBLAS, but the system actually installed an older version. This led to memory corruption and segmentation faults. * [Tests] Skip if extension not installed (part 2) (NVIDIA#1337) * [TravisCI] Install all plugins by default Also test no plugins * [Tests] Skip if extension not installed (NVIDIA#1337) * [Docs] Update model store documentation (NVIDIA#1346) TODO: add a screenshot of the official model store once approved * [Docs] Update model store documentation (NVIDIA#1346) TODO: add a screenshot of the official model store once approved * Add steps to specify the Python layer file (NVIDIA#1347) * Add steps to specify the Python layer file (NVIDIA#1347) * [Docs] Install minimal boost libs for caffe * Remove the selenium walkthrough * Update copyright year for 2017 * Add a few missing copyright notices * Fix Siamese example Broadcast -1 into all elements that equal 0 in original label. * Fix Siamese example (NVIDIA#1405) Broadcast -1 into all elements that equal 0 in original label. * [Packaging] Make nginx site easier to customize * Fix documentation typo. train.txt and test.txt was swapped and shown in the wrong folders for mnist and cifar10 data sets. * Document a cuDNN workaround for text example (NVIDIA#1422) * Document a cuDNN workaround for text example (NVIDIA#1422) * Correct shebang for prepare_pascal_voc_data.sh (NVIDIA#1450) * [Docs] Document workaround for torch+hdf5 error * Fix typo in ModelStore.md * Fix typo in medical-imaging/README.md * Fix bash lint with shellcheck * Fix bugs when visiting nested image folder * Fix shellcheck-related bug in PPA upload script * Copy labels.txt inside the dataset Move import to the top * Fix Distribution Graph Move backwards-compatibility to setstate * Fix typo in Sunnybrook plug-in * Fix a bug introduced when fixing shellcheck lint * Fix another shellcheck-related bug * Fix table formatting in README.md Fix table formatting * Clamp distance values from segementation boundaries before begin converted to uint8. That was causing banding in the image because of wrapping at V % 256 * lint * [Docs] 5.0 debs and Ubuntu 16.04 support * WIP lint fix * Linted most of what I can lint prior to asking for context * updated the model store urls in the readme * added debugs in build scripts to understand the point of failure * added travis wait to install openblas * removed tensorflow to the build process to see if affects openblas * removed suppressing log contents * added set -x * fixed control * re-enabling tensorflow to see if travis builds * updated the version of numpy to ensure a stable build for travis wrt to open issue 8653 on numpy github * forcing numpy to v 1.8.1 * added the official store image and updated the documentation (NVIDIA#1650) * [TravisCI] Add `git fetch --unshallow` for DIST Useful for TravisCI builds in forks. * Got travis script to work for tensorflow installation * Cleaning installation to work with Numpy 1.3 upgrade removed the open blas stuff that somehow made it into here embarassing merge residue force install specific numpy version because 1.13 was being installed asdf trying changing the tensorflow install reodered the installation order to see if it builds due to TF using numpy 1.13 now * Tf example (#3) * inital work on autoencoder TF example * Moved the example files to its proper location * atempting to get autoencoder to work * autoencoder work * validated tensorflow autoencoder example * updated gitignore * disabled comments in the segmentation-model.lua script to prevent crashing * commiting the changes made to binary segmentation tf * adding work to do something else * I am seriously wayy too tired to write this commit message, it's just random bits of stuff * got binary seg and siamese working * started to work on the regression network * milestone * got regression for TF working * Got fine tuning to work in TF * changed the code to the format that is wanted by tim and greg * Finished all the work for examples inital work on autoencoder TF example Moved the example files to its proper location atempting to get autoencoder to work autoencoder work validated tensorflow autoencoder example updated gitignore disabled comments in the segmentation-model.lua script to prevent crashing commiting the changes made to binary segmentation tf adding work to do something else I am seriously wayy too tired to write this commit message, it's just random bits of stuff got binary seg and siamese working rebase rebase started to work on the regression network milestone got regression for TF working Got fine tuning to work in TF changed the code to the format that is wanted by tim and greg got fine tuning working * Some small fixes * changes WRT PR trying renaming the weights tested renaming variables * fixed api problem for multi gpus * git removed installing tests * updated most of linting * Removed unused block of code as per suggestion by Greg * Removing spaces...
ethantang95 · Jun 21, 2017 · e0511e4 · e0511e4
1 parent 6e50347
commit e0511e4
Show file tree

Hide file tree

Showing 171 changed files with 22,169 additions and 290 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,6 +1,7 @@
 # Temporary files
 *.swp
 *~
+.DS_Store
 TAGS
 
 # Compiled / optimized files
@@ -16,4 +17,13 @@ TAGS
 /build/
 /dist/
 *.egg-info/
+
+#Intellij files
+.idea/
+
+#vscode
+.vscode/
+
+#.project
+.project
 /.project
diff --git a/.gjslintrc b/.gjslintrc
@@ -1,3 +1,3 @@
 --max_line_length=120
---exclude_directories=3rdparty
+--exclude_directories=3rdparty,tb
 --disable=0121,0220
diff --git a/.travis.yml b/.travis.yml
@@ -8,14 +8,16 @@ python: 2.7
 
 env:
     global:
-        - OPENBLAS_ROOT=~/openblas
         - CAFFE_ROOT=~/caffe
         - TORCH_ROOT=~/torch
+        - OMP_NUM_THREADS=1
+        - OPENBLAS_MAIN_FREE=1
         - secure: "WSqrE+PQm76DdoRLRGKTK6fRWfXZjIb0BWCZm3IgHgFO7OE6fcK2tBnpDNNw4XQjmo27FFWlEhxN32g18P84n5PvErHaH65IuS9Nv6FkLlPXZlVqGNxbPmEA4oTkD/6Y6kZyZWZtLh2+/1ijuzQAPnIy/4BEuL8pdO+PsoJ9hYM="
     matrix:
         - DIGITS_TEST_FRAMEWORK=caffe CAFFE_FORK=NVIDIA
         - DIGITS_TEST_FRAMEWORK=caffe CAFFE_FORK=BVLC
         - DIGITS_TEST_FRAMEWORK=torch
+        - DIGITS_TEST_FRAMEWORK=tensorflow
         - DIGITS_TEST_FRAMEWORK=none
         - DIGITS_TEST_FRAMEWORK=none WITH_PLUGINS=false
 
@@ -72,7 +74,6 @@ matrix:
 cache:
     apt: true
     directories:
-        - $OPENBLAS_ROOT
         - $CAFFE_ROOT
         - $TORCH_ROOT
 
@@ -83,6 +84,7 @@ addons:
             - cmake
             - cython
             - git
+            - gfortran
             - graphviz
             - libboost-filesystem-dev
             - libboost-python-dev
@@ -93,6 +95,7 @@ addons:
             - libhdf5-serial-dev
             - libleveldb-dev
             - liblmdb-dev
+            - libopenblas-dev
             - libopencv-dev
             - libprotobuf-dev
             - libsnappy-dev
@@ -126,14 +129,13 @@ before_install:
 install:
     - mkdir -p ~/.config/matplotlib
     - echo "backend:agg" > ~/.config/matplotlib/matplotlibrc
-    - ./scripts/travis/install-openblas.sh $OPENBLAS_ROOT
     - ./scripts/travis/install-caffe.sh $CAFFE_ROOT
-    - if [ "$DIGITS_TEST_FRAMEWORK" == "torch" ]; then ./scripts/travis/install-torch.sh $TORCH_ROOT; else unset TORCH_ROOT; fi
-    - pip install -r ./requirements.txt
+    - if [ "$DIGITS_TEST_FRAMEWORK" == "torch" ]; then travis_wait ./scripts/travis/install-torch.sh $TORCH_ROOT; else unset TORCH_ROOT; fi
+    - pip install -r ./requirements.txt --force-reinstall
+    - if [ "$DIGITS_TEST_FRAMEWORK" == "tensorflow" ]; then travis_wait ./scripts/travis/install-tensorflow.sh; fi
     - pip install -r ./requirements_test.txt
     - pip install -e .
     - if [ "$WITH_PLUGINS" != "false" ]; then find ./plugins/*/* -maxdepth 0 -type d | xargs -n1 pip install -e; fi
 
 script:
     - ./digits-test -v
-
diff --git a/digits-lint b/digits-lint
@@ -5,9 +5,9 @@ set -e
 
 echo "=== Checking for Python lint ..."
 if which flake8 >/dev/null 2>&1; then
-    python2 `which flake8` .
+    python2 `which flake8` --exclude ./examples,./digits/standard-networks/tensorflow .
 else
-    python2 -m flake8 .
+    python2 -m flake8 --exclude ./examples,./digits/standard-networks/tensorflow .
 fi
 
 echo "=== Checking for JavaScript lint ..."

diff --git a/digits/config/__init__.py b/digits/config/__init__.py
@@ -12,6 +12,7 @@
     torch,
     server_name,
     store_option,
+    tensorflow,
 )
 
 

diff --git a/digits/config/tensorflow.py b/digits/config/tensorflow.py
@@ -0,0 +1,47 @@
+# Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
+from __future__ import absolute_import
+
+import os
+import platform
+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?
+
+if platform.system() == 'Darwin':
+    # DYLD_LIBRARY_PATH and LD_LIBRARY_PATH is sometimes stripped, and the cuda libraries might need it
+    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)
+
+if VARNAME_ENV_TFPY in os.environ:
+    tf_python_exe = os.environ[VARNAME_ENV_TFPY]
+else:
+    tf_python_exe = DEFAULT_PYTHON_EXE
+
+tf_enabled, err = test_tf_import(tf_python_exe)
+
+if not tf_enabled:
+    print('Tensorflow support disabled.')
+#   print('Failed importing Tensorflow with python executable "%s"\n%s' % (tf_python_exe, err))
+
+if tf_enabled:
+    option_list['tensorflow'] = {
+        'enabled': True,
+        'executable': tf_python_exe,
+    }
+else:
+    option_list['tensorflow'] = {
+        'enabled': False,
+    }
diff --git a/digits/dataset/images/classification/forms.py b/digits/dataset/images/classification/forms.py
@@ -21,13 +21,16 @@ class ImageClassificationDatasetForm(ImageDatasetForm):
                                   choices=[
                                       ('lmdb', 'LMDB'),
                                       ('hdf5', 'HDF5'),
+                                      ('tfrecords', 'TFRecords'),
                                   ],
                                   default='lmdb',
                                   )
 
     def validate_backend(form, field):
         if field.data == 'lmdb':
             form.compression.data = 'none'
+        elif field.data == 'tfrecords':
+            form.compression.data = 'none'
         elif field.data == 'hdf5':
             form.encoding.data = 'none'
 

diff --git a/digits/dataset/tasks/create_db.py b/digits/dataset/tasks/create_db.py
@@ -50,8 +50,9 @@ def __init__(self, input_file, db_name, backend, image_dims, **kwargs):
         self.input_file = input_file
         self.db_name = db_name
         self.backend = backend
-        if backend == 'hdf5':
+        if backend == 'hdf5' or backend == 'tfrecords':
             # the list of hdf5 files is stored in a textfile
+            # tfrecords can be sharded as well
             self.textfile = os.path.join(self.db_name, 'list.txt')
         self.image_dims = image_dims
         if image_dims[2] == 3:

diff --git a/digits/dataset/views.py b/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
diff --git a/digits/extensions/view/imageOutput/config_template.html b/digits/extensions/view/imageOutput/config_template.html
@@ -12,8 +12,20 @@
     {{ form.channel_order(class='form-control') }}
 </div>
 
+<div class="form-group{{mark_errors([form.data_order])}}">
+    {{ form.data_order.label }}
+    {{ form.data_order.tooltip }}
+    {{ form.data_order(class='form-control') }}
+</div>
+
 <div class="form-group{{mark_errors([form.pixel_conversion])}}">
     {{ form.pixel_conversion.label }}
     {{ 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>
diff --git a/digits/extensions/view/imageOutput/forms.py b/digits/extensions/view/imageOutput/forms.py
@@ -23,6 +23,18 @@ class ConfigForm(Form):
                 'is ignored in the case of a grayscale image)'
     )
 
+    data_order = utils.forms.SelectField(
+        'Data order',
+        choices=[
+            ('chw', 'CHW'),
+            ('hwc', 'HWC'),
+            ],
+        default='chw',
+        tooltip="Set the order of the data. For Caffe and Torch models this "
+                "is often CHW, for Tensorflow it's HWC."
+                "W=Width, H=Height, C=Channels"
+        )
+
     pixel_conversion = utils.forms.SelectField(
         'Pixel conversion',
         choices=[
@@ -33,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'
+    )
diff --git a/digits/extensions/view/imageOutput/view.py b/digits/extensions/view/imageOutput/view.py
@@ -29,7 +29,9 @@ def __init__(self, dataset, **kwargs):
 
         # view options
         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():
@@ -69,15 +71,31 @@ 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
         """
-        # assume the only output is a CHW image
-        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]
+
+    def process_image(self, data_order, data):
+        if data_order == 'HWC':
+            data = (data.transpose((2, 0, 1)))
+
+        # assume CHW at this point
         channels = data.shape[0]
         if channels == 3 and self.channel_order == 'BGR':
             data = data[[2, 1, 0], ...]  # BGR to RGB

diff --git a/digits/extensions/view/imageOutput/view_template.html b/digits/extensions/view/imageOutput/view_template.html
@@ -1,3 +1,6 @@
 {# Copyright (c) 2016-2017, 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%;" />
diff --git a/digits/extensions/view/rawData/header_template.html b/digits/extensions/view/rawData/header_template.html
@@ -0,0 +1,7 @@
+{# Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved. #}
+
+{% from "helper.html" import print_flashes %}
+{% from "helper.html" import print_errors %}
+{% from "helper.html" import mark_errors %}
+
+{{data}}
diff --git a/digits/frameworks/__init__.py b/digits/frameworks/__init__.py
@@ -12,13 +12,20 @@
     'TorchFramework',
 ]
 
+if config_value('tensorflow')['enabled']:
+    from .tensorflow_framework import TensorflowFramework
+    __all__.append('TensorflowFramework')
+
 #
 #  create framework instances
 #
 
 # torch is optional
 torch = TorchFramework() if config_value('torch')['enabled'] else None
 
+# tensorflow is optional
+tensorflow = TensorflowFramework() if config_value('tensorflow')['enabled'] else None
+
 # caffe is mandatory
 caffe = CaffeFramework()
 
@@ -35,6 +42,8 @@ def get_frameworks():
     frameworks = [caffe]
     if torch:
         frameworks.append(torch)
+    if tensorflow:
+        frameworks.append(tensorflow)
     return frameworks
 
 

diff --git a/digits/frameworks/caffe_framework.py b/digits/frameworks/caffe_framework.py
@@ -34,6 +34,8 @@ class CaffeFramework(Framework):
 
     # whether this framework can shuffle data during training
     CAN_SHUFFLE_DATA = False
+    SUPPORTS_PYTHON_LAYERS_FILE = True
+    SUPPORTS_TIMELINE_TRACING = False
 
     if config_value('caffe')['flavor'] == 'NVIDIA':
         if parse_version(config_value('caffe')['version']) > parse_version('0.14.0-alpha'):
@@ -132,10 +134,11 @@ def get_network_from_path(self, path):
         return network
 
     @override
-    def get_network_visualization(self, desc):
+    def get_network_visualization(self, **kwargs):
         """
         return visualization of network
         """
+        desc = kwargs['desc']
         net = caffe_pb2.NetParameter()
         text_format.Merge(desc, net)
         # Throws an error if name is None
@@ -150,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.')