{"id":805,"date":"2021-09-05T15:52:51","date_gmt":"2021-09-05T15:52:51","guid":{"rendered":"https:\/\/salarydistribution.com\/machine-learning\/2021\/09\/05\/define-and-run-machine-learning-pipelines-on-step-functions-using-python-workflow-studio-or-states-language\/"},"modified":"2021-09-05T15:52:51","modified_gmt":"2021-09-05T15:52:51","slug":"define-and-run-machine-learning-pipelines-on-step-functions-using-python-workflow-studio-or-states-language","status":"publish","type":"post","link":"https:\/\/salarydistribution.com\/machine-learning\/2021\/09\/05\/define-and-run-machine-learning-pipelines-on-step-functions-using-python-workflow-studio-or-states-language\/","title":{"rendered":"Define and run Machine Learning pipelines on Step Functions using Python, Workflow Studio, or States Language"},"content":{"rendered":"
\n

You can use various tools to define and run machine learning (ML) pipelines or DAGs (Directed Acyclic Graphs). Some popular options include AWS Step Functions<\/a>, Apache Airflow, KubeFlow Pipelines (KFP), TensorFlow Extended (TFX), Argo, Luigi, and Amazon SageMaker Pipelines<\/a>. All these tools help you compose pipelines in various languages (JSON, YAML, Python, and more), followed by viewing and running them using a workflow orchestrator. A deep comparison of each of these options is out of scope for this post, and involves appropriately selecting and benchmarking tools for your specific use case.<\/p>\n

In this post, we discuss how to author end-to-end ML pipelines in Step Functions using three different methods:<\/p>\n

Solution overview<\/h2>\n

In this post, we create a simple workflow that involves a training step, creating a model, configuring an endpoint, and deploying the model<\/p>\n

\"\"<\/a><\/p>\n

You can also create more complex workflows involving other steps such as Amazon SageMaker Processing<\/a>, or automatic model tuning (HPO). You can also use Step Functions to integrate with other AWS services such as AWS Lambda<\/a>, Amazon DynamoDB<\/a>, AWS Glue<\/a>, Amazon EMR<\/a>, Amazon Athena<\/a>, Amazon Elastic Kubernetes Service<\/a> (Amazon EKS), and AWS Fargate<\/a>. For more information on supported services, see Supported AWS Service Integrations for Step Functions<\/a>. We provide guidance on other similar pipeline tools later in this post.<\/p>\n

In this post, we use the MNIST dataset, which is a widely used dataset for handwritten digit classification. It consists of 70,000 labeled 28\u00d728 pixel grayscale images of hand-written digits. The dataset is split into 60,000 training images and 10,000 test images. There are 10 classes (one for each of the 10 digits). The code used here closely follows a similar use case<\/a> where the task is to classify each input image as one of the 10 digits (0\u20139).<\/p>\n

The main training code uses a class from the standard PyTorch example<\/a> for the model definition:<\/p>\n

\n
class Net(nn.Module):\n    def __init__(self):\n        super(Net, self).__init__()\n        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n        self.conv2_drop = nn.Dropout2d()\n        self.fc1 = nn.Linear(320, 50)\n        self.fc2 = nn.Linear(50, 10)\n\n    def forward(self, x):\n        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n        x = x.view(-1, 320)\n        x = F.relu(self.fc1(x))\n        x = F.dropout(x, training=self.training)\n        x = self.fc2(x)\n        return F.log_softmax(x, dim=1)<\/code><\/pre>\n<\/p><\/div>\n
The main training function works for both single instance and distributed training, and does so by checking ARGs:<\/p>\n
\n
is_distributed = len(args.hosts) > 1 and args.backend is not None\nif is_distributed:\n        # Initialize the distributed environment.\n        world_size = len(args.hosts)\n        os.environ['WORLD_SIZE'] = str(world_size)\n        host_rank = args.hosts.index(args.current_host)\n        os.environ['RANK'] = str(host_rank)\n        dist.init_process_group(backend=args.backend, rank=host_rank, world_size=world_size)\n        logger.info('Initialized the distributed environment: '{}' backend on {} nodes. '.format(\n            args.backend, dist.get_world_size()) + 'Current host rank is {}. Number of gpus: {}'.format(\n            dist.get_rank(), args.num_gpus))            <\/code><\/pre>\n<\/p><\/div>\n
The number of hosts is conveniently stored in an Amazon SageMaker<\/a> environment variable, which can also be passed in as an argument:<\/p>\n
\n
parser.add_argument('--hosts', type=list, default=json.loads(os.environ['SM_HOSTS']))<\/code><\/pre>\n<\/p><\/div>\n
Next, we load the datasets from the default data directory:<\/p>\n
\n
(os.environ['SM_MODEL_DIR']):\n\ntrain_loader = _get_train_data_loader(args.batch_size, args.data_dir, is_distributed, **kwargs)\ntest_loader = _get_test_data_loader(args.test_batch_size, args.data_dir, **kwargs)\n<\/code><\/pre>\n<\/p><\/div>\n
We then enter the main training loop:<\/p>\n
\n
for\u00a0epoch in\u00a0range(1, args.epochs + 1):\nmodel.train()<\/code><\/pre>\n<\/p><\/div>\n
Finally, we save the model:<\/p>\n
\n
def save_model(model, model_dir):\n    logger.info(\"Saving the model.\")\n    path = os.path.join(model_dir, 'model.pth')\n    # recommended way from http:\/\/pytorch.org\/docs\/master\/notes\/serialization.html\n    torch.save(model.cpu().state_dict(), path)<\/code><\/pre>\n<\/p><\/div>\n
This code is stored in a file called mnist.py<\/code> and used in later steps (see the full code on GitHub<\/a>). The following are two important takeaways in connection to the pipeline:<\/p>\n