/
test_mnist.py
91 lines (68 loc) · 3.45 KB
/
test_mnist.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
# 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 the License is located at
#
# http://www.apache.org/licenses/LICENSE-2.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
import numpy as np
import pytest
import sagemaker.utils
from sagemaker.chainer import Chainer
from sagemaker.chainer.model import ChainerModel
from sagemaker.utils import sagemaker_timestamp
from timeout import timeout, timeout_and_delete_endpoint_by_name
@pytest.mark.deploy_test
def test_chainer_mnist_single_machine_train(sagemaker_session, ecr_image, instance_type):
_test_mnist_train(sagemaker_session, ecr_image, instance_type, 1, 'single_machine_customer_script.py')
@pytest.mark.deploy_test
def test_chainer_deploy(sagemaker_session, instance_type):
_test_mnist_deploy(sagemaker_session, instance_type)
def test_chainer_mnist_distributed_train(sagemaker_session, ecr_image, instance_type):
_test_mnist_train(sagemaker_session, ecr_image, instance_type, 2, 'distributed_customer_script.py')
def _test_mnist_train(sagemaker_session, ecr_image, instance_type, instance_count, script):
source_dir = 'test/resources/mnist'
with timeout(minutes=15):
data_path = 'test/resources/mnist/data'
chainer = Chainer(entry_point=script,
source_dir=source_dir,
role='SageMakerRole',
train_instance_count=instance_count,
train_instance_type=instance_type,
sagemaker_session=sagemaker_session,
image_name=ecr_image,
hyperparameters={'batch-size': 10000, 'epochs': 1})
prefix = 'chainer_mnist/{}'.format(sagemaker_timestamp())
train_data_path = os.path.join(data_path, 'train')
key_prefix = prefix + '/train'
train_input = sagemaker_session.upload_data(path=train_data_path, key_prefix=key_prefix)
test_path = os.path.join(data_path, 'test')
test_input = sagemaker_session.upload_data(path=test_path, key_prefix=prefix + '/test')
chainer.fit({'train': train_input, 'test': test_input})
def _test_mnist_deploy(sagemaker_session, instance_type):
model_path = 'test/resources/mnist/model.tar.gz'
script_path = 'test/resources/mnist/mnist.py'
endpoint_name = sagemaker.utils.unique_name_from_base('sagemaker-chainer-test')
model_data = sagemaker_session.upload_data(
path=model_path,
key_prefix='sagemaker-chainer/models',
)
with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session, minutes=30):
chainer = ChainerModel(
model_data=model_data,
role='SageMakerRole',
entry_point=script_path,
sagemaker_session=sagemaker_session,
)
predictor = chainer.deploy(initial_instance_count=1, instance_type=instance_type)
batch_size = 100
data = np.zeros(shape=(batch_size, 1, 28, 28), dtype='float32')
output = predictor.predict(data)
assert len(output) == batch_size