Enable selection of GPU on inference page

digits/inference/job.py

@@ -13,7 +13,7 @@ class InferenceJob(Job):
     A Job that exercises the forward pass of a neural network
     """
 
-    def __init__(self, model, images, epoch, layers, resize=True, **kwargs):
+    def __init__(self, model, images, epoch, layers, resize=True, gpu=None, **kwargs):
         """
         Arguments:
         model   -- job object associated with model to perform inference on
@@ -40,6 +40,7 @@ def __init__(self, model, images, epoch, layers, resize=True, **kwargs):
             epoch=epoch,
             layers=layers,
             resize=resize,
+            gpu=gpu
         ))
 
     @override

digits/inference/tasks/inference.py

@@ -21,7 +21,7 @@ class InferenceTask(Task):
     A task for inference jobs
     """
 
-    def __init__(self, model, images, epoch, layers, resize, **kwargs):
+    def __init__(self, model, images, epoch, layers, resize, gpu=None, **kwargs):
         """
         Arguments:
         model  -- trained model to perform inference on
@@ -40,7 +40,7 @@ def __init__(self, model, images, epoch, layers, resize, **kwargs):
         self.inference_log_file = "inference.log"
 
         # resources
-        self.gpu = None
+        self.gpu = gpu
 
         # generated data
         self.inference_data_filename = None
@@ -180,6 +180,8 @@ def offer_resources(self, resources):
                 if resources[gpu_key]:
                     for resource in resources[gpu_key]:
                         if resource.remaining() >= 1:
+                            if self.gpu is not None and self.gpu != int(resource.identifier):
+                                continue
                             self.gpu = int(resource.identifier)
                             reserved_resources[gpu_key] = [(resource.identifier, 1)]
                             break

digits/model/forms.py

@@ -305,37 +305,30 @@ def validate_custom_network_snapshot(form, field):
                 if filename and not os.path.exists(filename):
                     raise validators.ValidationError('File "%s" does not exist' % filename)
 
+    # List of GPUs
+    gpu_list = [(
+        index,
+        '#%s - %s (%s memory)' % (
+            index,
+            get_device(index).name,
+            sizeof_fmt(
+                get_nvml_info(index)['memory']['total']
+                if get_nvml_info(index) and 'memory' in get_nvml_info(index)
+                else get_device(index).totalGlobalMem)
+        ),
+    ) for index in config_value('gpu_list').split(',') if index]
+
     # Select one of several GPUs
     select_gpu = wtforms.RadioField(
         'Select which GPU you would like to use',
-        choices=[('next', 'Next available')] + [(
-            index,
-            '#%s - %s (%s memory)' % (
-                index,
-                get_device(index).name,
-                sizeof_fmt(
-                    get_nvml_info(index)['memory']['total']
-                    if get_nvml_info(index) and 'memory' in get_nvml_info(index)
-                    else get_device(index).totalGlobalMem)
-            ),
-        ) for index in config_value('gpu_list').split(',') if index],
+        choices=[('next', 'Next available')] + gpu_list,
         default='next',
     )
 
     # Select N of several GPUs
     select_gpus = utils.forms.SelectMultipleField(
         'Select which GPU[s] you would like to use',
-        choices=[(
-            index,
-            '#%s - %s (%s memory)' % (
-                index,
-                get_device(index).name,
-                sizeof_fmt(
-                    get_nvml_info(index)['memory']['total']
-                    if get_nvml_info(index) and 'memory' in get_nvml_info(index)
-                    else get_device(index).totalGlobalMem)
-            ),
-        ) for index in config_value('gpu_list').split(',') if index],
+        choices=gpu_list,
         tooltip="The job won't start until all of the chosen GPUs are available."
     )
 

digits/model/images/classification/views.py

@@ -19,7 +19,7 @@
 from digits.pretrained_model.job import PretrainedModelJob
 from digits.status import Status
 from digits.utils import filesystem as fs
-from digits.utils.forms import fill_form_if_cloned, save_form_to_job
+from digits.utils.forms import fill_form_if_cloned, save_form_to_job, get_selected_gpu
 from digits.utils.routing import request_wants_json, job_from_request
 from digits.webapp import scheduler
 
@@ -333,14 +333,16 @@ def show(job, related_jobs=None):
     """
     Called from digits.model.views.models_show()
     """
+    form = ImageClassificationModelForm()
     return flask.render_template(
         'models/images/classification/show.html',
         job=job,
         framework_ids=[
             fw.get_id()
             for fw in frameworks.get_frameworks()
         ],
-        related_jobs=related_jobs
+        related_jobs=related_jobs,
+        gpu_list=form.gpu_list,
     )
 
 
@@ -384,14 +386,17 @@ def classify_one():
     if 'show_visualizations' in flask.request.form and flask.request.form['show_visualizations']:
         layers = 'all'
 
+    selected_gpu = get_selected_gpu(flask.request.form)
+
     # create inference job
     inference_job = ImageInferenceJob(
         username=utils.auth.get_username(),
         name="Classify One Image",
         model=model_job,
         images=[image_path],
         epoch=epoch,
-        layers=layers
+        layers=layers,
+        gpu=selected_gpu
     )
 
     # schedule tasks
@@ -477,6 +482,7 @@ def classify_many():
         epoch = float(flask.request.form['snapshot_epoch'])
 
     paths, ground_truths = read_image_list(image_list, image_folder, num_test_images)
+    selected_gpu = get_selected_gpu(flask.request.form)
 
     # create inference job
     inference_job = ImageInferenceJob(
@@ -485,7 +491,8 @@ def classify_many():
         model=model_job,
         images=paths,
         epoch=epoch,
-        layers='none'
+        layers='none',
+        gpu=selected_gpu
     )
 
     # schedule tasks
@@ -633,6 +640,7 @@ def top_n():
         num_test_images = None
 
     paths, _ = read_image_list(image_list, image_folder, num_test_images)
+    selected_gpu = get_selected_gpu(flask.request.form)
 
     # create inference job
     inference_job = ImageInferenceJob(
@@ -641,7 +649,8 @@ def top_n():
         model=model_job,
         images=paths,
         epoch=epoch,
-        layers='none'
+        layers='none',
+        gpu=selected_gpu
     )
 
     # schedule tasks

digits/model/images/generic/views.py

@@ -18,7 +18,7 @@
 from digits.status import Status
 from digits.utils import filesystem as fs
 from digits.utils import constants
-from digits.utils.forms import fill_form_if_cloned, save_form_to_job
+from digits.utils.forms import fill_form_if_cloned, save_form_to_job, get_selected_gpu
 from digits.utils.routing import request_wants_json, job_from_request
 from digits.webapp import scheduler
 
@@ -309,12 +309,14 @@ def show(job, related_jobs=None):
             template, context = extension.get_inference_template(form)
             inference_form_html = flask.render_template_string(template, **context)
 
+    generic_form = GenericImageModelForm()
     return flask.render_template(
         'models/images/generic/show.html',
         job=job,
         view_extensions=view_extensions,
         related_jobs=related_jobs,
         inference_form_html=inference_form_html,
+        gpu_list=generic_form.gpu_list,
     )
 
 
@@ -361,6 +363,8 @@ def infer_one():
     else:
         resize = True
 
+    selected_gpu = get_selected_gpu(flask.request.form)
+
     # create inference job
     inference_job = ImageInferenceJob(
         username=utils.auth.get_username(),
@@ -370,6 +374,7 @@ def infer_one():
         epoch=epoch,
         layers=layers,
         resize=resize,
+        gpu=selected_gpu
     )
 
     # schedule tasks
@@ -446,6 +451,8 @@ def infer_extension():
         if 'show_visualizations' in flask.request.form and flask.request.form['show_visualizations']:
             layers = 'all'
 
+        selected_gpu = get_selected_gpu(flask.request.form)
+
         # create inference job
         inference_job = ImageInferenceJob(
             username=utils.auth.get_username(),
@@ -455,6 +462,7 @@ def infer_extension():
             epoch=epoch,
             layers=layers,
             resize=False,
+            gpu=selected_gpu
         )
 
         # schedule tasks
@@ -539,6 +547,8 @@ def infer_db():
     else:
         resize = True
 
+    selected_gpu = get_selected_gpu(flask.request.form)
+
     # create inference job
     inference_job = ImageInferenceJob(
         username=utils.auth.get_username(),
@@ -548,6 +558,7 @@ def infer_db():
         epoch=epoch,
         layers='none',
         resize=resize,
+        gpu=selected_gpu
     )
 
     # schedule tasks
@@ -633,6 +644,8 @@ def infer_many():
     else:
         resize = True
 
+    selected_gpu = get_selected_gpu(flask.request.form)
+
     paths = []
 
     for line in image_list.readlines():
@@ -664,6 +677,7 @@ def infer_many():
         epoch=epoch,
         layers='none',
         resize=resize,
+        gpu=selected_gpu
     )
 
     # schedule tasks

digits/templates/models/images/classification/show.html

@@ -8,6 +8,7 @@
 <script src="{{ url_for('static', filename='js/model-graphs.js', ver=dir_hash) }}"></script>
 
 {% set task = job.train_task() %}
+{% set show_multi_gpu_form = gpu_list| length > 1 %}
 
 <div class="row">
     <div class="col-sm-6">
@@ -102,13 +103,9 @@ <h4 class='text-center'>Dataset</h4>
                 {% endif %}
                 >
                 <h2>Trained Models</h2>
-                <div class="row">
-                    <div class="col-sm-12">
-                        <label for="snapshot_epoch">Select Model</label>
-                    </div>
-                </div>
                 <div class="row">
                     <div class="col-sm-6">
+                        <label for="snapshot_epoch">Select Model</label>
                         <div class="form-group">
                             <select id="snapshot_epoch" name="snapshot_epoch" class="form-control">
                             </select>
@@ -139,9 +136,41 @@ <h2>Trained Models</h2>
 updateSnapshotList({% autoescape false %}{{task.snapshot_list()}}{% endautoescape %});
                             </script>
                         </div>
+                        {% if show_multi_gpu_form %}
+                        <button
+                            formaction="{{url_for('digits.model.views.download', job_id=job.id())}}"
+                            formmethod="post"
+                            formenctype="multipart/form-data"
+                            class="btn btn-info">
+                            Download Model
+                        </button>
+                        <button
+                            formaction="{{url_for('digits.model.views.to_pretrained', job_id=job.id())}}"
+                            formmethod="post"
+                            formenctype="multipart/form-data"
+                            class="btn btn-success">
+                            Make Pretrained Model
+                        </button>
+                        {% endif %}
+                    </div>
+                    {% if show_multi_gpu_form %}
+                    <div class="col-sm-6">
+                        <div class="form-group">
+                            <label for="select_one_of_gpus">Select which GPU you would like to use</label>
+                            <select class="form-control" id="select_one_of_gpus" name="select_one_of_gpus" size="4">
+                                <option selected="" value="next">Next available</option>
+                                {% for gpu_id, gpu_description in gpu_list %}
+                                    <option value="{{ gpu_id }}">{{ gpu_description }}</option>
+                                {% endfor %}
+                            </select>
+                        </div>
                     </div>
+                    {% else %}
                     <div class="col-sm-6">
-                        <button
+                        <label for="empty space">&nbsp;</label>
+                    </div>
+                    <div class="col-sm-6">
+                         <button
                             formaction="{{url_for('digits.model.views.download', job_id=job.id())}}"
                             formmethod="post"
                             formenctype="multipart/form-data"
@@ -156,6 +185,7 @@ <h2>Trained Models</h2>
                             Make Pretrained Model
                         </button>
                     </div>
+                    {% endif %}
                 </div>
                 {% if task.get_framework_id() in framework_ids %}
                 <div class="row">
@@ -184,6 +214,7 @@ <h3>Test a single image</h3>
                                 </div>
                             </div>
                         </div>
+
                         <script type="text/javascript">
                             // When you fill in one field, the other gets blanked out
                             $("#image_path").change(function() { $("#image_file").val(""); });