def _test_mnist_distributed(sagemaker_session, ecr_image, instance_type):
model_dir = os.path.join(model_cpu_dir, 'model_mnist.tar.gz')
endpoint_name = sagemaker.utils.unique_name_from_base(
"sagemaker-pytorch-serving")
model_data = sagemaker_session.upload_data(
path=model_dir,
key_prefix="sagemaker-pytorch-serving/models",
)
pytorch = PyTorchModel(model_data, 'SageMakerRole', mnist_script,
ecr_image, sagemaker_session)
with timeout_and_delete_endpoint(endpoint_name,
sagemaker_session,
minutes=30):
predictor = pytorch.deploy(initial_instance_count=1,
instance_type=instance_type,
endpoint_name=endpoint_name)
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)
def _test_mnist_distributed(sagemaker_session, ecr_image, instance_type, model_dir, mnist_script,
accelerator_type=None):
endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-pytorch-serving")
model_data = sagemaker_session.upload_data(
path=model_dir,
key_prefix="sagemaker-pytorch-serving/models",
)
pytorch = PyTorchModel(model_data=model_data, role='SageMakerRole', entry_point=mnist_script,
image=ecr_image, sagemaker_session=sagemaker_session)
with timeout_and_delete_endpoint(endpoint_name, sagemaker_session, minutes=30):
# Use accelerator type to differentiate EI vs. CPU and GPU. Don't use processor value
if accelerator_type is not None:
predictor = pytorch.deploy(initial_instance_count=1, instance_type=instance_type,
accelerator_type=accelerator_type, endpoint_name=endpoint_name)
else:
predictor = pytorch.deploy(initial_instance_count=1, instance_type=instance_type,
endpoint_name=endpoint_name)
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)
def _test_mnist_distributed(sagemaker_session, ecr_image, instance_type,
dist_backend):
with timeout(minutes=DEFAULT_TIMEOUT):
pytorch = PyTorch(entry_point=mnist_script,
role='SageMakerRole',
train_instance_count=2,
train_instance_type=instance_type,
sagemaker_session=sagemaker_session,
image_name=ecr_image,
hyperparameters={
'backend': dist_backend,
'epochs': 1
})
training_input = pytorch.sagemaker_session.upload_data(
path=training_dir, key_prefix='pytorch/mnist')
pytorch.fit({'training': training_input})
with timeout_and_delete_endpoint(estimator=pytorch, minutes=30):
predictor = pytorch.deploy(initial_instance_count=1,
instance_type=instance_type)
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)