AWS Sagemaker Training of a Binary Classification model using PyTorch
By: Santhosh Puttaraju
Tags: Amazon Sagemaker, Pytorch

In this post, we’ll learn how to train a binary classification model using PyTorch on AWS Sagemaker, where the model ready data resides in Snowflake. The focus here isn’t really on building a very good deep learning model; its rather on how we can train a PyTorch model on AWS Sagemaker.

You can use Amazon SageMaker to train and deploy a model using custom PyTorch code. The SageMaker Python SDK PyTorch estimators and models and the SageMaker open-source PyTorch container make writing a PyTorch script and running it in SageMaker easier.

The model ready data for this demo is stored in Snowflake. So, we’ll be connecting to a Snowflake instance to get the train and test data. The PyTorch container provided by Sagemaker doesn’t have the snowflake python connector. So, I’ll be using a custom container which has this connector (refer my previous blog on how to build one: https://bit.ly/3rrBc64 ). In your scenario, you might have data in a S3 bucket in which case you can use the sagemaker’s default container itself.

1. Log into AWS and select the Sagemaker service. Then select/create a notebook instance and create a python3-PyTorch notebook instance.

2. Importing all the required libraries and since I need to establish connection to Snowflake, I’ll set up a helper function. This function will help get an access token (using oauth) for snowflake connection. This depends on how you’ve setup access to your snowflake environment and if it’s just through a username and password this won’t be required.

3. Next, setup a S3 session, generate a token and save it in a file in S3. Also create a dummy file in the same S3 folder.

4. Prepare a training script. This script holds all the logic required for the model training and will be submitted to Sagemaker via a PyTorch estimator.

a. Import all libraries that would be required for the code in the script and define a helper function for snowflake connection.

b. Define a class for the neural network. It’s a simple network consisting of 3 layers. 1 input, 1 hidden and 1 output layer and applies Linear transformation to input data. The input data has 23 features. Hidden layer contains (a random) 7 neurons. And 2 output classes.

c. Next, the main method is written. All the parameters passed by the estimator will be parsed. Then, the snowflake connection is established by first retrieving the token stored in S3, saving it to a variable and then passing it to the helper function along the other environment details.

d. Now that the snowflake connection is set-up, we get the train and test data stored in the MRD tables on snowflake. Then, convert them to tensor data types.

e. The final piece in the training script would be to initialize the model, define the loss criterion and the optimizer and perform the training.

5. The final bit is defining the estimator. Before that the subnets and security group ids have been declared to make sure that we would be running the training job inside our own VPC to take care of the proxies. The snowflake token path (in S3) and the name of the custom container are also passed. Also, please note the path where the training script is saved and pass it. And then call the fit method to initiate the training job.

6. Once the above code is run, the training job will be initiated and once it is completed we can look at the CloudWatch logs (link will be found in the training job details) for any information we want. The same logs can be used in case the job runs into an error.

Please leave your comments and questions below.