AWS PyTorch: Build, Train, and Deploy Deep Learning Models in the Cloud

A developer published a walkthrough showing how to train a small PyTorch digit classifier on the MNIST dataset and extend it into an AWS pipeline using Amazon S3 for dataset and artifact storage, SageMaker AI for managed training jobs, and CloudWatch for monitoring latency and resource usage. The guide starts with a local training script using a simple two-layer fully connected network and then outlines the path from local experimentation to scalable cloud deployment.

PyTorch makes it easier to experiment with neural networks, computer vision, natural language processing, and deep learning. But building a model is only one part of a machine learning project. When it's time to move from experimentation to production, you need reliable compute, data storage, repeatable training, model artifacts, deployment, monitoring, and cost controls. This is where AWS can help. By combining PyTorch, Amazon SageMaker AI, Amazon S3, and Amazon CloudWatch , developers can build machine learning workflows that move from local experiments toward scalable cloud deployments. In this article, we'll build a simple image classifier and explore how to extend the project into an AWS-based machine learning pipeline. PyTorch is an open-source deep learning framework that provides tensor operations, automatic differentiation, GPU acceleration, and tools for creating neural networks. AWS provides several services that can support PyTorch workloads: You don't need all these services for every project. Start with what your workload requires and add complexity when it provides a clear benefit. A typical workflow looks like this: Training Dataset | v Amazon S3 | v Amazon SageMaker AI Training Job | v PyTorch Model | v Model Artifacts in S3 | v Deployment / Inference | v Application Predictions Amazon S3 provides durable storage for datasets and artifacts. SageMaker AI can run managed training jobs, while a suitable inference service can expose the trained model to applications. CloudWatch and related AWS monitoring capabilities help you investigate failures, latency, and resource usage. For small experiments, you can train locally first. For larger workloads, managed training can reduce infrastructure-management overhead. We'll train a small neural network on the MNIST handwritten-digit dataset. Our model will learn to classify grayscale images into one of ten categories, representing the digits from 0 to 9. You'll need: The first part runs locally and does not require AWS resources. Create a project directory and enter it: mkdir aws-pytorch-demo cd aws-pytorch-demo python -m venv .venv Activate the virtual environment on Windows: .venv\Scripts\Activate.ps1 On Linux or macOS: source .venv/bin/activate Install the dependencies: pip install torch torchvision For GPU training, install a PyTorch build compatible with your GPU and CUDA configuration. Follow the official installation instructions: https://pytorch.org/get-started/locally/ https://pytorch.org/get-started/locally/ . Create train.py : python import torch from torch import nn from torch.utils.data import DataLoader from torchvision import datasets, transforms device = torch.device "cuda" if torch.cuda.is available else "cpu" print f"Training device: {device}" transform = transforms.ToTensor train data = datasets.MNIST root="data", train=True, download=True, transform=transform, test data = datasets.MNIST root="data", train=False, download=True, transform=transform, train loader = DataLoader train data, batch size=64, shuffle=True, test loader = DataLoader test data, batch size=256, shuffle=False, class DigitClassifier nn.Module : def init self : super . init self.network = nn.Sequential nn.Flatten , nn.Linear 28 28, 128 , nn.ReLU , nn.Linear 128, 10 , def forward self, x : return self.network x model = DigitClassifier .to device criterion = nn.CrossEntropyLoss optimizer = torch.optim.Adam model.parameters , lr=0.001, Training loop for epoch in range 3 : model.train total loss = 0.0 for images, labels in train loader: images = images.to device labels = labels.to device optimizer.zero grad outputs = model images loss = criterion outputs, labels loss.backward optimizer.step total loss += loss.item average loss = total loss / len train loader print f"Epoch {epoch + 1}/3, " f"Loss: {average loss:.4f}" Evaluation model.eval correct = 0 total = 0 with torch.no grad : for images, labels in test loader: images = images.to device labels = labels.to device outputs = model images predictions = outputs.argmax dim=1 correct += predictions == labels .sum .item total += labels.size 0 accuracy = 100 correct / total print f"Test accuracy: {accuracy:.2f}%" Save model parameters torch.save model.state dict , "digit classifier.pth", print "Saved digit classifier.pth" Run: python train.py The script downloads MNIST, trains the neural network, evaluates it against the test dataset, and saves the learned parameters. Your output should include: Training device: cpu Epoch 1/3, Loss: ... Epoch 2/3, Loss: ... Epoch 3/3, Loss: ... Test accuracy: ... Saved digit classifier.pth The loss and accuracy values depend on the training run. If a compatible GPU is available, the script can use it automatically. Create predict.py : python import torch from torch import nn from torchvision import datasets, transforms class DigitClassifier nn.Module : def init self : super . init self.network = nn.Sequential nn.Flatten , nn.Linear 28 28, 128 , nn.ReLU , nn.Linear 128, 10 , def forward self, x : return self.network x device = torch.device "cuda" if torch.cuda.is available else "cpu" model = DigitClassifier .to device state = torch.load "digit classifier.pth", map location=device, weights only=True, model.load state dict state model.eval dataset = datasets.MNIST root="data", train=False, download=True, transform=transforms.ToTensor , image, actual label = dataset 0 with torch.no grad : logits = model image.unsqueeze 0 .to device predicted label = logits.argmax dim=1 .item print f"Predicted digit: {predicted label}" print f"Actual digit: {actual label}" Run the prediction: python predict.py Example output: Predicted digit: 7 Actual digit: 7 These values are illustrative. The actual result depends on your trained model. You now have a basic end-to-end deep learning workflow: train, evaluate, save, reload, and predict. Once the model works locally, you can move the workflow to AWS. Choose a globally unique bucket name: aws s3 mb s3://YOUR-UNIQUE-PYTORCH-BUCKET \ --region us-east-1 Replace the placeholder with your actual bucket name. Upload the training script: aws s3 cp train.py \ s3://YOUR-UNIQUE-PYTORCH-BUCKET/code/train.py For this example, the script downloads MNIST directly. A production pipeline should manage dataset preparation, storage, and versioning more explicitly. Amazon SageMaker AI supports PyTorch training jobs using supported framework environments and training configurations. At a high level, you need to: Use the current official guide for the SDK version you install: https://docs.aws.amazon.com/sagemaker/latest/dg/pytorch.html https://docs.aws.amazon.com/sagemaker/latest/dg/pytorch.html SageMaker SDK interfaces evolve, so use examples matching your installed SDK version rather than mixing older estimator APIs with newer SDK workflows. For a locally trained model, you can upload the saved parameters: aws s3 cp digit classifier.pth \ s3://YOUR-UNIQUE-PYTORCH-BUCKET/models/digit classifier.pth For a managed training job, configure the output location so SageMaker AI saves the model artifacts automatically. Keep a record of the dataset version, source-code version, training parameters, framework version, and evaluation metrics. This makes it easier to reproduce and audit experiments. Training and inference are separate stages. For an application to request predictions, you need a serving interface that loads the model, validates input, executes inference, and returns a response. With Amazon SageMaker AI, you can deploy supported PyTorch models to a managed endpoint. The required model packaging and inference code depend on the selected deployment approach. A typical inference flow is: For the MNIST model, input should represent a grayscale image with dimensions of 28 × 28 pixels. Validate the request before running inference. For workloads that don't need immediate responses, consider batch inference instead of keeping an online endpoint available continuously. A model that works locally isn't automatically production-ready. Track dataset versions, dependencies, training parameters, random seeds, and model artifacts. Store code in version control and preserve evaluation results for each model version. Monitor training failures, compute utilization, inference latency, and application errors using supported AWS monitoring capabilities. Also evaluate model quality over time. A service can remain operational while its predictions become less accurate because real-world data has changed. Don't rely on training loss alone. For classification tasks, consider precision, recall, F1 score, confusion matrices, and performance on representative held-out data. Select metrics based on the real business problem. Common applications include: For large-scale distributed training, additional planning is needed for GPU memory, network communication, checkpointing, and data loading. PyTorch provides the flexibility to build deep learning models, while AWS offers several ways to run training workloads, store model artifacts, and deploy predictions. Start with a small experiment, validate the results, and then introduce cloud infrastructure where it provides measurable value. The real goal is not just to train a model in the cloud. It's to create a secure, reproducible, observable machine learning workflow that can grow with your application.