Update fit() to work with frameworks (aws#27)

Author: laurenyu

src/sagemaker/session.py

@@ -12,18 +12,17 @@
 # language governing permissions and limitations under the License.
 from __future__ import print_function, absolute_import
 
+import json
 import logging
-import re
-
 import os
+import re
 import sys
 import time
 
 import boto3
-import json
+import botocore.config
 import six
 import yaml
-import botocore.config
 from botocore.exceptions import ClientError
 
 from sagemaker.user_agent import prepend_user_agent
@@ -257,22 +256,22 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
         LOGGER.debug('train request: {}'.format(json.dumps(train_request, indent=4)))
         self.sagemaker_client.create_training_job(**train_request)
 
-    def tune(self, job_name, strategy, objective, metric_name,
+    def tune(self, job_name, strategy, objective_type, objective_metric_name,
              max_jobs, max_parallel_jobs, parameter_ranges,
-             static_hp, image, input_mode, metric_definitions,
+             static_hyperparameters, image, input_mode, metric_definitions,
              role, input_config, output_config, resource_config, stop_condition):
-        """Create an Amazon SageMaker HPO job.
+        """Create an Amazon SageMaker hyperparameter tuning job
 
         Args:
             job_name (str): Name of the tuning job being created.
             strategy (str): Strategy to be used.
-            objective (str): Minimize/Maximize
-            metric_name (str): Name of the metric to use when evaluating training job.
+            objective_type (str): Minimize/Maximize
+            objective_metric_name (str): Name of the metric to use when evaluating training job.
             max_jobs (int): Maximum total number of jobs to start.
             max_parallel_jobs (int): Maximum number of parallel jobs to start.
             parameter_ranges (dict): Parameter ranges in a dictionary of types: Continuous, Integer, Categorical
-            static_hp (dict): Hyperparameters for model training. The hyperparameters are made accessible as
-                a dict[str, str] to the training code on SageMaker. For convenience, this accepts other types for
+            static_hyperparameters (dict): Hyperparameters for model training. The hyperparameters are made accessible
+                as a dict[str, str] to the training code on SageMaker. For convenience, this accepts other types for
                 keys and values, but ``str()`` will be called to convert them before training.
             image (str): Docker image containing training code.
             input_mode (str): The input mode that the algorithm supports. Valid modes:
@@ -293,30 +292,27 @@ def tune(self, job_name, strategy, objective, metric_name,
             instance_type (str): Type of EC2 instance to use for training, for example, 'ml.c4.xlarge'.
             stop_condition (dict): Defines when training shall finish. Contains entries that can be understood by the
                 service like ``MaxRuntimeInSeconds``.
-
-        Returns:
-
         """
-
         tune_request = {
             'HyperParameterTuningJobName': job_name,
             'HyperParameterTuningJobConfig': {
                 'Strategy': strategy,
                 'HyperParameterTuningJobObjective': {
-                    'Type': objective,
-                    'MetricName': metric_name,
+                    'Type': objective_type,
+                    'MetricName': objective_metric_name,
                 },
                 'ResourceLimits': {
                     'MaxNumberOfTrainingJobs': max_jobs,
-                    'MaxParallelTrainingJobs': max_parallel_jobs
+                    'MaxParallelTrainingJobs': max_parallel_jobs,
                 },
-                'ParameterRanges': parameter_ranges
+                'ParameterRanges': parameter_ranges,
             },
             'TrainingJobDefinition': {
-                'StaticHyperParameters': static_hp,
+                'StaticHyperParameters': static_hyperparameters,
                 'AlgorithmSpecification': {
                     'TrainingImage': image,
-                    'TrainingInputMode': input_mode
+                    'TrainingInputMode': input_mode,
+                    'MetricDefinitions': metric_definitions,
                 },
                 'RoleArn': role,
                 'InputDataConfig': input_config,
@@ -329,7 +325,7 @@ def tune(self, job_name, strategy, objective, metric_name,
         if metric_definitions is not None:
             tune_request['TrainingJobDefinition']['AlgorithmSpecification']['MetricDefinitions'] = metric_definitions
 
-        LOGGER.info('Creating tuning-job with name: {}'.format(job_name))
+        LOGGER.info('Creating hyperparameter tuning job with name: {}'.format(job_name))
         LOGGER.debug('tune request: {}'.format(json.dumps(tune_request, indent=4)))
         self.sagemaker_client.create_hyper_parameter_tuning_job(**tune_request)
 

src/sagemaker/tuner.py

@@ -13,7 +13,9 @@
 from __future__ import absolute_import
 
 import inspect
+import json
 
+from sagemaker.estimator import Framework
 from sagemaker.job import _Job
 from sagemaker.utils import base_name_from_image, name_from_base
 
@@ -51,6 +53,9 @@ def as_tuning_range(self, name):
         return {'Name': name,
                 'Values': self.values}
 
+    def as_json_range(self, name):
+        return {'Name': name, 'Values': [json.dumps(v) for v in self.values]}
+
 
 class IntegerParameter(_ParameterRange):
     __name__ = 'Integer'
@@ -60,31 +65,58 @@ def __init__(self, min_value, max_value):
 
 
 class HyperparameterTuner(object):
-    __objectives__ = ['Minimize', 'Maximize']
+    SAGEMAKER_ESTIMATOR_CLASS_NAME = 'sagemaker_estimator_class_name'
+    SAGEMAKER_ESTIMATOR_MODULE = 'sagemaker_estimator_module'
 
     def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metric_definitions, strategy='Bayesian',
                  objective_type='Maximize', max_jobs=1, max_parallel_jobs=1, base_tuning_job_name=None):
-        if objective_type not in HyperparameterTuner.__objectives__:
-            raise ValueError("Unsupported 'objective' values")
+        self._hyperparameter_ranges = hyperparameter_ranges
+        if self._hyperparameter_ranges is None or len(self._hyperparameter_ranges) == 0:
+            raise ValueError('Need to specify hyperparameter ranges')
 
         self.estimator = estimator
         self.objective_metric_name = objective_metric_name
-        self._hyperparameter_ranges = hyperparameter_ranges
+        self.metric_definitions = metric_definitions
+
         self.strategy = strategy
         self.objective_type = objective_type
+
         self.max_jobs = max_jobs
         self.max_parallel_jobs = max_parallel_jobs
         self.tuning_job_name = base_tuning_job_name
         self.metric_definitions = metric_definitions
         self.latest_tuning_job = None
         self._validate_parameter_ranges()
 
-    def fit(self, inputs):
-        """Create tuning job
+    def prepare_for_training(self):
+        # TODO: Change this so that it can handle unicode in Python 2
+        self.static_hyperparameters = {str(k): str(v) for (k, v) in self.estimator.hyperparameters().items()}
+        for hyperparameter_name in self._hyperparameter_ranges.keys():
+            self.static_hyperparameters.pop(hyperparameter_name, None)
+
+        # For attach() to know what estimator to use
+        self.static_hyperparameters[self.SAGEMAKER_ESTIMATOR_CLASS_NAME] = self.estimator.__class__.__name__
+        self.static_hyperparameters[self.SAGEMAKER_ESTIMATOR_MODULE] = self.estimator.__module__
+
+    def fit(self, inputs, job_name=None, **kwargs):
+        """Start a hyperparameter tuning job.
 
         Args:
-            inputs (str): Parameters used when called  :meth:`~sagemaker.estimator.EstimatorBase.fit`.
+            inputs (str): Parameters used when called :meth:`~sagemaker.estimator.EstimatorBase.fit`.
+            job_name (str): Job name
+            **kwargs: Other arguments
         """
+        # TODO: I think I have to move RecordSet to its own file
+        from sagemaker.amazon.amazon_estimator import RecordSet
+
+        # 1P estimators require a RecordSet object
+        if isinstance(inputs, RecordSet):
+            self.estimator.prepare_for_training(inputs, **kwargs)
+            inputs = inputs.data_channel()
+        else:
+            self.estimator.prepare_for_training(**kwargs)
+
+        self.prepare_for_training()
         self.latest_tuning_job = _TuningJob.start_new(self, inputs)
 
     def stop_tuning_job(self):
@@ -101,15 +133,20 @@ def hyperparameter_ranges(self):
         """Return collections of ``ParameterRanges``
 
         Returns:
-            dict: ParameterRanges suitable for tuning job.
+            dict: ParameterRanges suitable for a hyperparameter tuning job.
         """
         hyperparameter_ranges = dict()
         for range_type in _ParameterRange.__all_types__:
-            parameter_range = []
+            parameter_ranges = []
             for parameter_name, parameter in self._hyperparameter_ranges.items():
                 if parameter is not None and parameter.__name__ == range_type:
-                    parameter_range.append(parameter.as_tuning_range(parameter_name))
-            hyperparameter_ranges[range_type + 'ParameterRanges'] = parameter_range
+                    # Categorical parameters needed to be serialized as JSON for our framework containers
+                    if isinstance(parameter, CategoricalParameter) and isinstance(self.estimator, Framework):
+                        tuning_range = parameter.as_json_range(parameter_name)
+                    else:
+                        tuning_range = parameter.as_tuning_range(parameter_name)
+                    parameter_ranges.append(tuning_range)
+            hyperparameter_ranges[range_type + 'ParameterRanges'] = parameter_ranges
         return hyperparameter_ranges
 
     def _validate_parameter_ranges(self):
@@ -138,43 +175,38 @@ def _validate_parameter_ranges(self):
 
 
 class _TuningJob(_Job):
-    SAGEMAKER_ESTIMATOR_CLASS_NAME = 'sagemaker_estimator_class_name'
-    SAGEMAKER_ESTIMATOR_MODULE = 'sagemaker_estimator_module'
-
     def __init__(self, sagemaker_session, tuning_job_name):
         super(_TuningJob, self).__init__(sagemaker_session, tuning_job_name)
 
     @classmethod
     def start_new(cls, tuner, inputs):
-        """Create a new Amazon SageMaker tuning job from the HyperparameterTuner.
+        """Create a new Amazon SageMaker hyperparameter tuning job from the HyperparameterTuner.
 
         Args:
-            tuner (sagemaker.tuner.HyperparameterTuner): Tuner object created by the user.
-            inputs (str): Parameters used when called  :meth:`~sagemaker.estimator.EstimatorBase.fit`.
+            tuner (sagemaker.tuner.HyperparameterTuner): HyperparameterTuner object created by the user.
+            inputs (str): Parameters used when called :meth:`~sagemaker.estimator.EstimatorBase.fit`.
 
         Returns:
             sagemaker.tuner._TuningJob: Constructed object that captures all information about the started job.
         """
         config = _Job._load_config(inputs, tuner.estimator)
 
-        static_hyperparameters = {str(k): str(v) for (k, v) in tuner.estimator.hyperparameters().items()}
-        for hyperparameter_name in tuner._hyperparameter_ranges.keys():
-            static_hyperparameters.pop(hyperparameter_name, None)
-
-        static_hyperparameters[cls.SAGEMAKER_ESTIMATOR_CLASS_NAME] = tuner.estimator.__class__.__name__
-        static_hyperparameters[cls.SAGEMAKER_ESTIMATOR_MODULE] = tuner.estimator.__module__
-
         base_name = tuner.estimator.base_job_name or base_name_from_image(tuner.estimator.train_image())
         tuning_job_name = name_from_base(base_name)
 
+        # TODO: Update name generation so that the base name isn't limited to so few characters
+        if len(tuning_job_name) > 32:
+            raise ValueError('Tuning job name too long - must be 32 characters or fewer: {}'.format(tuning_job_name))
+
         tuner.estimator.sagemaker_session.tune(job_name=tuning_job_name, strategy=tuner.strategy,
-                                               objective=tuner.objective_type, metric_name=tuner.objective_metric_name,
+                                               objective_type=tuner.objective_type,
+                                               objective_metric_name=tuner.objective_metric_name,
                                                max_jobs=tuner.max_jobs, max_parallel_jobs=tuner.max_parallel_jobs,
                                                parameter_ranges=tuner.hyperparameter_ranges(),
-                                               static_hp=static_hyperparameters,
+                                               static_hyperparameters=tuner.static_hyperparameters,
                                                image=tuner.estimator.train_image(),
                                                input_mode=tuner.estimator.input_mode,
-                                               metric_definitions=tuner.estimator.metric_definitions,
+                                               metric_definitions=tuner.metric_definitions,
                                                role=(config['role']), input_config=(config['input_config']),
                                                output_config=(config['output_config']),
                                                resource_config=(config['resource_config']),

tests/data/mxnet_mnist/tuning.py

@@ -0,0 +1,73 @@
+import gzip
+import logging
+import os
+import struct
+
+import mxnet as mx
+import numpy as np
+
+
+def load_data(path):
+    with gzip.open(find_file(path, "labels.gz")) as flbl:
+        struct.unpack(">II", flbl.read(8))
+        labels = np.fromstring(flbl.read(), dtype=np.int8)
+    with gzip.open(find_file(path, "images.gz")) as fimg:
+        _, _, rows, cols = struct.unpack(">IIII", fimg.read(16))
+        images = np.fromstring(fimg.read(), dtype=np.uint8).reshape(len(labels), rows, cols)
+        images = images.reshape(images.shape[0], 1, 28, 28).astype(np.float32) / 255
+    return labels, images
+
+
+def find_file(root_path, file_name):
+    for root, dirs, files in os.walk(root_path):
+        if file_name in files:
+            return os.path.join(root, file_name)
+
+
+def build_graph():
+    data = mx.sym.var('data')
+    data = mx.sym.flatten(data=data)
+    fc1 = mx.sym.FullyConnected(data=data, num_hidden=128)
+    act1 = mx.sym.Activation(data=fc1, act_type="relu")
+    fc2 = mx.sym.FullyConnected(data=act1, num_hidden=64)
+    act2 = mx.sym.Activation(data=fc2, act_type="relu")
+    fc3 = mx.sym.FullyConnected(data=act2, num_hidden=10)
+    return mx.sym.SoftmaxOutput(data=fc3, name='softmax')
+
+
+def train(current_host, channel_input_dirs, hyperparameters, hosts, num_cpus, num_gpus):
+    (train_labels, train_images) = load_data(os.path.join(channel_input_dirs['train']))
+    (test_labels, test_images) = load_data(os.path.join(channel_input_dirs['test']))
+
+    # Alternatively to splitting in memory, the data could be pre-split in S3 and use ShardedByS3Key
+    # to do parallel training.
+    shard_size = len(train_images) // len(hosts)
+    for i, host in enumerate(hosts):
+        if host == current_host:
+            start = shard_size * i
+            end = start + shard_size
+            break
+
+    batch_size = 100
+    train_iter = mx.io.NDArrayIter(train_images[start:end], train_labels[start:end], batch_size, shuffle=True)
+    val_iter = mx.io.NDArrayIter(test_images, test_labels, batch_size)
+    logging.getLogger().setLevel(logging.DEBUG)
+    kvstore = 'local' if len(hosts) == 1 else 'dist_sync'
+    mlp_model = mx.mod.Module(
+        symbol=build_graph(),
+        context=get_train_context(num_cpus, num_gpus))
+    mlp_model.fit(train_iter,
+                  eval_data=val_iter,
+                  kvstore=kvstore,
+                  optimizer='sgd',
+                  optimizer_params={'learning_rate': float(hyperparameters.get("learning_rate", 0.1))},
+                  eval_metric='acc',
+                  batch_end_callback=mx.callback.Speedometer(batch_size, 100),
+                  num_epoch=25)
+    return mlp_model
+
+
+def get_train_context(num_cpus, num_gpus):
+    if num_gpus > 0:
+        return mx.gpu()
+    return mx.cpu()