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

> Source: <https://dev.to/omayrashraf/aws-pytorch-build-train-and-deploy-deep-learning-models-in-the-cloud-40i5>
> Published: 2026-10-21 02:34:59+00:00

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.**