Revert "Removed unnecessary mnist scripts."

This reverts commit 7d84091.

Author: nadiaya

test-toolkit/integration/__init__.py

@@ -24,9 +24,12 @@
 eia_sub_dir = 'model_eia'
 
 model_cpu_dir = os.path.join(mnist_path, cpu_sub_dir)
+mnist_cpu_script = os.path.join(model_cpu_dir, 'mnist.py')
 model_cpu_1d_dir = os.path.join(model_cpu_dir, '1d')
 mnist_1d_script = os.path.join(model_cpu_1d_dir, 'mnist_1d.py')
 model_gpu_dir = os.path.join(mnist_path, gpu_sub_dir)
+mnist_gpu_script = os.path.join(model_gpu_dir, 'mnist.py')
+model_gpu_1d_dir = os.path.join(model_gpu_dir, '1d')
 model_eia_dir = os.path.join(mnist_path, eia_sub_dir)
 mnist_eia_script = os.path.join(model_eia_dir, 'mnist.py')
 call_model_fn_once_script = os.path.join(resources_path, 'call_model_fn_once.py')

test-toolkit/integration/sagemaker/test_mnist.py

@@ -19,22 +19,22 @@
 import sagemaker
 from sagemaker.pytorch import PyTorchModel
 
-from integration import model_cpu_dir, mnist_script, model_eia_dir, mnist_eia_script
+from integration import model_cpu_dir, mnist_cpu_script, mnist_gpu_script, model_eia_dir, mnist_eia_script
 from integration.sagemaker.timeout import timeout_and_delete_endpoint
 
 
 @pytest.mark.cpu_test
 def test_mnist_cpu(sagemaker_session, image_uri, instance_type):
     instance_type = instance_type or 'ml.c4.xlarge'
     model_dir = os.path.join(model_cpu_dir, 'model_mnist.tar.gz')
-    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_script)
+    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_cpu_script)
 
 
 @pytest.mark.gpu_test
 def test_mnist_gpu(sagemaker_session, image_uri, instance_type):
     instance_type = instance_type or 'ml.p2.xlarge'
     model_dir = os.path.join(model_cpu_dir, 'model_mnist.tar.gz')
-    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_script)
+    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_gpu_script)
 
 
 @pytest.mark.eia_test

test-toolkit/resources/mnist/model_cpu/call_model_fn_once.py

@@ -0,0 +1,37 @@
+# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import os
+
+
+def model_fn(model_dir):
+    lock_file = os.path.join(model_dir, 'model_fn.lock.{}'.format(os.getpid()))
+    if os.path.exists(lock_file):
+        raise RuntimeError('model_fn called more than once (lock: {})'.format(lock_file))
+
+    open(lock_file, 'a').close()
+
+    return 'model'
+
+
+def input_fn(data, content_type):
+    return data
+
+
+def predict_fn(data, model):
+    return b'output'
+
+
+def output_fn(prediction, accept):
+    return prediction

test-toolkit/resources/mnist/model_cpu/mnist.py

@@ -0,0 +1,57 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import logging
+import os
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data
+import torch.utils.data.distributed
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+# Based on https://github.com/pytorch/examples/blob/master/mnist/main.py
+class Net(nn.Module):
+    def __init__(self):
+        logger.info("Create neural network module")
+
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def model_fn(model_dir):
+    logger.info('model_fn')
+    model = torch.nn.DataParallel(Net())
+    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
+        model.load_state_dict(torch.load(f))
+    return model

test-toolkit/resources/mnist/model_eia/mnist.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"). You
 # may not use this file except in compliance with the License. A copy of

test-toolkit/resources/mnist/model_gpu/mnist.py

@@ -0,0 +1,57 @@
+# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import logging
+import os
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data
+import torch.utils.data.distributed
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+# Based on https://github.com/pytorch/examples/blob/master/mnist/main.py
+class Net(nn.Module):
+    def __init__(self):
+        logger.info("Create neural network module")
+
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def model_fn(model_dir):
+    logger.info('model_fn')
+    model = torch.nn.DataParallel(Net())
+    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
+        model.load_state_dict(torch.load(f))
+    return model

test-toolkit/utils/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2018-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"). You
 # may not use this file except in compliance with the License. A copy of