Add PyTorch hyperparameter tuning integ test

tests/data/pytorch_mnist/mnist.py

@@ -39,14 +39,14 @@ def forward(self, x):
         return F.log_softmax(x, dim=1)
 
 
-def _get_train_data_loader(training_dir, is_distributed, **kwargs):
+def _get_train_data_loader(training_dir, is_distributed, batch_size, **kwargs):
     logger.info('Get train data loader')
     dataset = datasets.MNIST(training_dir, train=True, transform=transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
     ]))
     train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) if is_distributed else None
-    train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=train_sampler is None,
+    train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=train_sampler is None,
                                                sampler=train_sampler, **kwargs)
     return train_sampler, train_loader
 
@@ -94,7 +94,7 @@ def train(args):
     if use_cuda:
         torch.cuda.manual_seed(seed)
 
-    train_sampler, train_loader = _get_train_data_loader(args.data_dir, is_distributed, **kwargs)
+    train_sampler, train_loader = _get_train_data_loader(args.data_dir, is_distributed, args.batch_size, **kwargs)
     test_loader = _get_test_data_loader(args.data_dir, **kwargs)
 
     logger.debug('Processes {}/{} ({:.0f}%) of train data'.format(
@@ -142,9 +142,11 @@ def train(args):
                 logger.debug('Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(
                     epoch, batch_idx * len(data), len(train_loader.sampler),
                     100. * batch_idx / len(train_loader), loss.item()))
-        test(model, test_loader, device)
+        accuracy = test(model, test_loader, device)
     save_model(model, args.model_dir)
 
+    logger.debug('Overall test accuracy: {}'.format(accuracy))
+
 
 def test(model, test_loader, device):
     model.eval()
@@ -159,9 +161,12 @@ def test(model, test_loader, device):
             correct += pred.eq(target.view_as(pred)).sum().item()
 
     test_loss /= len(test_loader.dataset)
+    accuracy = 100. * correct / len(test_loader.dataset)
+
     logger.debug('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
-        test_loss, correct, len(test_loader.dataset),
-        100. * correct / len(test_loader.dataset)))
+        test_loss, correct, len(test_loader.dataset), accuracy))
+
+    return accuracy
 
 
 def model_fn(model_dir):
@@ -181,6 +186,7 @@ def save_model(model, model_dir):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--epochs', type=int, default=1, metavar='N')
+    parser.add_argument('--batch-size', type=int, default=64, metavar='N')
 
     # Container environment
     parser.add_argument('--hosts', type=list, default=json.loads(os.environ['SM_HOSTS']))

tests/integ/test_tuner.py

@@ -31,6 +31,7 @@
 from sagemaker.estimator import Estimator
 from sagemaker.mxnet.estimator import MXNet
 from sagemaker.predictor import json_deserializer
+from sagemaker.pytorch import PyTorch
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.tuner import IntegerParameter, ContinuousParameter, CategoricalParameter, HyperparameterTuner
 from tests.integ import DATA_DIR
@@ -314,6 +315,47 @@ def test_tuning_chainer(sagemaker_session):
         assert len(output) == batch_size
 
 
+@pytest.mark.continuous_testing
+def test_attach_tuning_pytorch(sagemaker_session):
+    mnist_dir = os.path.join(DATA_DIR, 'pytorch_mnist')
+    mnist_script = os.path.join(mnist_dir, 'mnist.py')
+
+    estimator = PyTorch(entry_point=mnist_script, role='SageMakerRole', train_instance_count=1,
+                        train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_session)
+
+    with timeout(minutes=15):
+        objective_metric_name = 'evaluation-accuracy'
+        metric_definitions = [{'Name': 'evaluation-accuracy', 'Regex': 'Overall test accuracy: (\d+)'}]
+        hyperparameter_ranges = {'batch-size': IntegerParameter(50, 100)}
+
+        tuner = HyperparameterTuner(estimator, objective_metric_name, hyperparameter_ranges, metric_definitions,
+                                    max_jobs=2, max_parallel_jobs=2)
+
+        training_data = estimator.sagemaker_session.upload_data(path=os.path.join(mnist_dir, 'training'),
+                                                                key_prefix='integ-test-data/pytorch_mnist/training')
+        tuner.fit({'training': training_data})
+
+        tuning_job_name = tuner.latest_tuning_job.name
+
+        print('Started hyperparameter tuning job with name:' + tuning_job_name)
+
+        time.sleep(15)
+        tuner.wait()
+
+    attached_tuner = HyperparameterTuner.attach(tuning_job_name, sagemaker_session=sagemaker_session)
+    best_training_job = tuner.best_training_job()
+    with timeout_and_delete_endpoint_by_name(best_training_job, sagemaker_session, minutes=20):
+        predictor = attached_tuner.deploy(1, 'ml.c4.xlarge')
+        data = np.zeros(shape=(1, 1, 28, 28), dtype=np.float32)
+        predictor.predict(data)
+
+        batch_size = 100
+        data = np.random.rand(batch_size, 1, 28, 28).astype(np.float32)
+        output = predictor.predict(data)
+
+        assert output.shape == (batch_size, 10)
+
+
 @pytest.mark.continuous_testing
 def test_tuning_byo_estimator(sagemaker_session):
     """Use Factorization Machines algorithm as an example here.