trainer/examples/mlx/image-classification/mnist.ipynb

Distributed MNIST with MLX and Kubeflow Trainer

This Notebook will show how to run distributed MLX on Kubernetes with Kubeflow Trainer.

The MLX runtime uses mlx[cuda] package to run distributed training on GPUs.

MLX Distributed: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Install the Kubeflow SDK

You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:

In [ ]:

# !pip install -U kubeflow

Install the MLX Dependencies

You also need to install MLX packages to evaluate the trained model.

In [ ]:

!pip install "mlx>=0.28.0"
!pip install "mlx-data>=0.1.0"
!pip install "matplotlib>=3.10.5"

Create MLX training script

We need to wrap our training script into a function to create Kubeflow TrainJob.

This is the simple MLP model to recognize digits from the MNIST dataset.

In [1]:

def mlx_train_mnist(bucket: str, model_name: str):
    import time
    from functools import partial
    import mlx.core as mx
    import mlx.nn as nn
    import mlx.optimizers as optim
    from mlx.data.datasets import load_mnist
    import boto3

    # Define a simple MLP model with MLX.
    class MLP(nn.Module):
        def __init__(
            self, in_dims: int, hidden_dims: int, num_layers: int, out_dims: int
        ):
            super().__init__()
            layer_sizes = [in_dims] + [hidden_dims] * num_layers + [out_dims]
            self.layers = [
                nn.Linear(idim, odim)
                for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
            ]

        def __call__(self, x):
            for layer in self.layers[:-1]:
                x = nn.relu(layer(x))
            return self.layers[-1](x)

    # Initializer MLX distributed backend.
    dist = mx.distributed.init(strict=True, backend="mpi")
    world_size = dist.size()
    rank = dist.rank()
    print(f"Start Distributed Training, WORLD_SIZE: {world_size}, RANK: {rank}")

    if rank == 0:
        print("Download MNIST dataset")

    # Load MNIST dataset and partition it.
    BATCH_SIZE = 128
    train_dataset = load_mnist()

    distributed_ds = (
        train_dataset.shuffle()
        .partition_if(world_size > 1, world_size, rank)
        .key_transform("image", lambda x: (x.astype("float32") / 255.0).ravel())
    )

    # Create the MLP model and SGD optimizer
    model = MLP(
        in_dims=distributed_ds[0]["image"].shape[-1],
        hidden_dims=32,
        num_layers=2,
        out_dims=10,
    )
    optimizer = optim.SGD(learning_rate=0.01)

    # Define function to calculate loss and accuracy.
    def loss_fn(model, x, y):
        output = model(x)
        loss = mx.mean(nn.losses.cross_entropy(output, y))
        acc = mx.mean(mx.argmax(output, axis=1) == y)
        return loss, acc

    # Define single training step.
    @partial(mx.compile, inputs=model.state, outputs=model.state)
    def step(x, y):
        loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
        (loss, acc), grads = loss_and_grad_fn(model, x, y)
        # Average grads to aggregate them across distributed nodes.
        grads = nn.utils.average_gradients(grads)
        optimizer.update(model, grads)
        return loss, acc

    # Average statistic across distributed nodes.
    def average_stats(stats, count):
        with mx.stream(mx.cpu):
            stats = mx.distributed.all_sum(mx.array(stats))
            count = mx.distributed.all_sum(count)
            return (stats / count).tolist()

    # Start distributed training.
    if rank == 0:
        print("Start MLX training")
    for epoch in range(15):
        epoch_start = time.perf_counter()
        losses = accuracies = count = 0

        for batch_idx, batch_sample in enumerate(distributed_ds.batch(BATCH_SIZE)):
            x = mx.array(batch_sample["image"])
            y = mx.array(batch_sample["label"])
            loss, acc = step(x, y)
            mx.eval(loss, acc, model.state)

            losses += loss.item()
            accuracies += acc.item()
            count += 1

            # Print the results.
            if batch_idx % 10 == 0:
                loss, acc = average_stats([losses, accuracies], count)
                if rank == 0:
                    print(
                        "Epoch: {} [{}/{} ({:.0f}%)] \tTrain loss: {:.3f}, acc: {:.3f}".format(
                            epoch,
                            batch_idx * len(x),
                            len(train_dataset),
                            100.0 * batch_idx * len(x) / len(train_dataset),
                            loss,
                            acc,
                        )
                    )
        if rank == 0:
            print(
                "Epoch: {}, time: {:.2f} seconds\n\n".format(
                    epoch, time.perf_counter() - epoch_start
                )
            )

    # Finally, export model to S3.
    if rank == 0:
        print("Exporting MNIST model to S3")
        model.save_weights(model_name)
        boto3.client("s3").upload_file(model_name, bucket, model_name)

Get the MLX Runtime

You can list the available Kubeflow Trainer runtimes with the list_runtimes() API.

The name of MLX runtime is mlx-distributed.

In [2]:

from kubeflow.trainer import TrainerClient, CustomTrainer

for r in TrainerClient().list_runtimes():
    if r.name == "mlx-distributed":
        print(f"Name: {r.name}, Framework: {r.trainer.framework}, Trainer Type: {r.trainer.trainer_type.value}\n")
        mlx_runtime = r

Name: mlx-distributed, Framework: mlx, Trainer Type: CustomTrainer

Create TrainJob for Distributed Training

Use the train() API to create distributed TrainJob on 3 GPUs. Every MPI training node uses 1 GPU.

Note: Update the S3 bucket name.

In [4]:

# BUCKET = "Add your S3 bucket"
MODEL = "mnist.npz"

In [5]:

args = {"bucket": BUCKET, "model_name": MODEL}

job_id = TrainerClient().train(
    trainer=CustomTrainer(
        func=mlx_train_mnist,
        func_args=args,
        num_nodes=3,
        packages_to_install=["boto3"],
        resources_per_node={
            "gpu": 1
        },
    ),
    runtime=mlx_runtime,
)

In [6]:

# Train API generates a random TrainJob id.
job_id

Out [6]:

'vd59b4f8583c'

Check the TrainJob Info

Use the list_jobs() and get_job() APIs to get information about created TrainJob and its steps.

In [7]:

for job in TrainerClient().list_jobs():
    print(f"TrainJob: {job.name}, Status: {job.status}, Created at: {job.creation_timestamp}")

TrainJob: vd59b4f8583c, Status: Created, Created at: 2025-10-18 00:21:13+00:00

In [8]:

# We execute mpirun command on node-0, which functions as the MPI Launcher node.
for c in TrainerClient().get_job(name=job_id).steps:
    print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\n")

Step: node-0, Status: Running, Devices: gpu x 1

Step: node-1, Status: Running, Devices: gpu x 1

Step: node-2, Status: Running, Devices: gpu x 1

Get the TrainJob Logs

Use the get_job_logs() API to retrieve the TrainJob logs.

Since we distribute the dataset across 3 GPUs, each rank processes round(60,000 / 3) = 20,000 samples.

In [9]:

for logline in TrainerClient().get_job_logs(job_id, follow=True):
    print(logline)

Warning: Permanently added '[vd59b4f8583c-node-0-1.vd59b4f8583c]:2222' (ECDSA) to the list of known hosts.
Warning: Permanently added '[vd59b4f8583c-node-0-0.vd59b4f8583c]:2222' (ECDSA) to the list of known hosts.
Start Distributed Training, WORLD_SIZE: 3, RANK: 0
Download MNIST dataset
Start Distributed Training, WORLD_SIZE: 3, RANK: 1
Start Distributed Training, WORLD_SIZE: 3, RANK: 2
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/train-images-idx3-ubyte.gz |███████████████████████████████████     | 8.2MiB / 9.5MiB (134.8MiB/s)  
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/train-images-idx3-ubyte.gz 9.5MiB (147.3MiB/s)                                                      
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/train-images-idx3-ubyte.gz 9.5MiB (119.3MiB/s)                                                      
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/t10k-images-idx3-ubyte.gz |█                                       | 32.0KiB / 1.6MiB (3.9MiB/s)   
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/t10k-images-idx3-ubyte.gz 1.6MiB (35.0MiB/s)                                                       
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/t10k-images-idx3-ubyte.gz 1.6MiB (35.3MiB/s)                                                       
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/train-labels-idx1-ubyte.gz 32.0KiB (17.1MiB/s)                                                     
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/train-labels-idx1-ubyte.gz 32.0KiB (18.4MiB/s)                                                     
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/t10k-labels-idx1-ubyte.gz 8.0KiB (54.3MiB/s)                                               
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/train-labels-idx1-ubyte.gz 32.0KiB (18.5MiB/s)                                                     
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/t10k-labels-idx1-ubyte.gz 8.0KiB (53.5MiB/s)                                               
Downloading https://raw.githubusercontent.com/fgnt/mnist/master/t10k-labels-idx1-ubyte.gz 8.0KiB (56.3MiB/s)                                               
Start MLX training
Epoch: 0 [0/60000 (0%)] 	Train loss: 2.310, acc: 0.086
Epoch: 0 [1280/60000 (2%)] 	Train loss: 2.309, acc: 0.100
Epoch: 0 [2560/60000 (4%)] 	Train loss: 2.308, acc: 0.100
Epoch: 0 [3840/60000 (6%)] 	Train loss: 2.310, acc: 0.097
Epoch: 0 [5120/60000 (9%)] 	Train loss: 2.310, acc: 0.098
Epoch: 0 [6400/60000 (11%)] 	Train loss: 2.309, acc: 0.098
Epoch: 0 [7680/60000 (13%)] 	Train loss: 2.309, acc: 0.100
Epoch: 0 [8960/60000 (15%)] 	Train loss: 2.309, acc: 0.100
Epoch: 0 [10240/60000 (17%)] 	Train loss: 2.308, acc: 0.102
Epoch: 0 [11520/60000 (19%)] 	Train loss: 2.308, acc: 0.103
Epoch: 0 [12800/60000 (21%)] 	Train loss: 2.307, acc: 0.104
Epoch: 0 [14080/60000 (23%)] 	Train loss: 2.307, acc: 0.105
Epoch: 0 [15360/60000 (26%)] 	Train loss: 2.307, acc: 0.105
Epoch: 0 [16640/60000 (28%)] 	Train loss: 2.307, acc: 0.105
Epoch: 0 [17920/60000 (30%)] 	Train loss: 2.306, acc: 0.106
Epoch: 0 [19200/60000 (32%)] 	Train loss: 2.306, acc: 0.106
Epoch: 0, time: 29.13 seconds


Epoch: 1 [0/60000 (0%)] 	Train loss: 2.299, acc: 0.115
Epoch: 1 [1280/60000 (2%)] 	Train loss: 2.298, acc: 0.126
Epoch: 1 [2560/60000 (4%)] 	Train loss: 2.297, acc: 0.126
Epoch: 1 [3840/60000 (6%)] 	Train loss: 2.299, acc: 0.122
Epoch: 1 [5120/60000 (9%)] 	Train loss: 2.299, acc: 0.123
Epoch: 1 [6400/60000 (11%)] 	Train loss: 2.298, acc: 0.122
Epoch: 1 [7680/60000 (13%)] 	Train loss: 2.298, acc: 0.124
Epoch: 1 [8960/60000 (15%)] 	Train loss: 2.298, acc: 0.124
Epoch: 1 [10240/60000 (17%)] 	Train loss: 2.297, acc: 0.126
Epoch: 1 [11520/60000 (19%)] 	Train loss: 2.296, acc: 0.127
Epoch: 1 [12800/60000 (21%)] 	Train loss: 2.296, acc: 0.128
Epoch: 1 [14080/60000 (23%)] 	Train loss: 2.296, acc: 0.129
Epoch: 1 [15360/60000 (26%)] 	Train loss: 2.296, acc: 0.128
Epoch: 1 [16640/60000 (28%)] 	Train loss: 2.295, acc: 0.128
Epoch: 1 [17920/60000 (30%)] 	Train loss: 2.295, acc: 0.129
Epoch: 1 [19200/60000 (32%)] 	Train loss: 2.295, acc: 0.130
Epoch: 1, time: 0.48 seconds


Epoch: 2 [0/60000 (0%)] 	Train loss: 2.288, acc: 0.133
Epoch: 2 [1280/60000 (2%)] 	Train loss: 2.286, acc: 0.144
Epoch: 2 [2560/60000 (4%)] 	Train loss: 2.286, acc: 0.147
Epoch: 2 [3840/60000 (6%)] 	Train loss: 2.287, acc: 0.146
Epoch: 2 [5120/60000 (9%)] 	Train loss: 2.287, acc: 0.148
Epoch: 2 [6400/60000 (11%)] 	Train loss: 2.286, acc: 0.149
Epoch: 2 [7680/60000 (13%)] 	Train loss: 2.286, acc: 0.151
Epoch: 2 [8960/60000 (15%)] 	Train loss: 2.286, acc: 0.153
Epoch: 2 [10240/60000 (17%)] 	Train loss: 2.285, acc: 0.156
Epoch: 2 [11520/60000 (19%)] 	Train loss: 2.284, acc: 0.159
Epoch: 2 [12800/60000 (21%)] 	Train loss: 2.284, acc: 0.161
Epoch: 2 [14080/60000 (23%)] 	Train loss: 2.284, acc: 0.163
Epoch: 2 [15360/60000 (26%)] 	Train loss: 2.283, acc: 0.164
Epoch: 2 [16640/60000 (28%)] 	Train loss: 2.283, acc: 0.165
Epoch: 2 [17920/60000 (30%)] 	Train loss: 2.282, acc: 0.167
Epoch: 2 [19200/60000 (32%)] 	Train loss: 2.282, acc: 0.169
Epoch: 2, time: 0.50 seconds


Epoch: 3 [0/60000 (0%)] 	Train loss: 2.274, acc: 0.172
Epoch: 3 [1280/60000 (2%)] 	Train loss: 2.272, acc: 0.197
Epoch: 3 [2560/60000 (4%)] 	Train loss: 2.271, acc: 0.199
Epoch: 3 [3840/60000 (6%)] 	Train loss: 2.272, acc: 0.197
Epoch: 3 [5120/60000 (9%)] 	Train loss: 2.272, acc: 0.197
Epoch: 3 [6400/60000 (11%)] 	Train loss: 2.271, acc: 0.197
Epoch: 3 [7680/60000 (13%)] 	Train loss: 2.270, acc: 0.199
Epoch: 3 [8960/60000 (15%)] 	Train loss: 2.270, acc: 0.201
Epoch: 3 [10240/60000 (17%)] 	Train loss: 2.269, acc: 0.204
Epoch: 3 [11520/60000 (19%)] 	Train loss: 2.268, acc: 0.206
Epoch: 3 [12800/60000 (21%)] 	Train loss: 2.267, acc: 0.208
Epoch: 3 [14080/60000 (23%)] 	Train loss: 2.266, acc: 0.210
Epoch: 3 [15360/60000 (26%)] 	Train loss: 2.266, acc: 0.211
Epoch: 3 [16640/60000 (28%)] 	Train loss: 2.265, acc: 0.211
Epoch: 3 [17920/60000 (30%)] 	Train loss: 2.264, acc: 0.214
Epoch: 3 [19200/60000 (32%)] 	Train loss: 2.264, acc: 0.216
Epoch: 3, time: 0.49 seconds


Epoch: 4 [0/60000 (0%)] 	Train loss: 2.250, acc: 0.219
Epoch: 4 [1280/60000 (2%)] 	Train loss: 2.247, acc: 0.236
Epoch: 4 [2560/60000 (4%)] 	Train loss: 2.247, acc: 0.236
Epoch: 4 [3840/60000 (6%)] 	Train loss: 2.247, acc: 0.238
Epoch: 4 [5120/60000 (9%)] 	Train loss: 2.247, acc: 0.236
Epoch: 4 [6400/60000 (11%)] 	Train loss: 2.245, acc: 0.237
Epoch: 4 [7680/60000 (13%)] 	Train loss: 2.244, acc: 0.238
Epoch: 4 [8960/60000 (15%)] 	Train loss: 2.244, acc: 0.239
Epoch: 4 [10240/60000 (17%)] 	Train loss: 2.242, acc: 0.242
Epoch: 4 [11520/60000 (19%)] 	Train loss: 2.240, acc: 0.244
Epoch: 4 [12800/60000 (21%)] 	Train loss: 2.239, acc: 0.245
Epoch: 4 [14080/60000 (23%)] 	Train loss: 2.238, acc: 0.247
Epoch: 4 [15360/60000 (26%)] 	Train loss: 2.237, acc: 0.246
Epoch: 4 [16640/60000 (28%)] 	Train loss: 2.236, acc: 0.247
Epoch: 4 [17920/60000 (30%)] 	Train loss: 2.234, acc: 0.248
Epoch: 4 [19200/60000 (32%)] 	Train loss: 2.233, acc: 0.249
Epoch: 4, time: 0.49 seconds


Epoch: 5 [0/60000 (0%)] 	Train loss: 2.209, acc: 0.279
Epoch: 5 [1280/60000 (2%)] 	Train loss: 2.207, acc: 0.270
Epoch: 5 [2560/60000 (4%)] 	Train loss: 2.206, acc: 0.267
Epoch: 5 [3840/60000 (6%)] 	Train loss: 2.205, acc: 0.269
Epoch: 5 [5120/60000 (9%)] 	Train loss: 2.204, acc: 0.268
Epoch: 5 [6400/60000 (11%)] 	Train loss: 2.202, acc: 0.271
Epoch: 5 [7680/60000 (13%)] 	Train loss: 2.200, acc: 0.272
Epoch: 5 [8960/60000 (15%)] 	Train loss: 2.199, acc: 0.275
Epoch: 5 [10240/60000 (17%)] 	Train loss: 2.196, acc: 0.279
Epoch: 5 [11520/60000 (19%)] 	Train loss: 2.194, acc: 0.282
Epoch: 5 [12800/60000 (21%)] 	Train loss: 2.192, acc: 0.284
Epoch: 5 [14080/60000 (23%)] 	Train loss: 2.189, acc: 0.286
Epoch: 5 [15360/60000 (26%)] 	Train loss: 2.188, acc: 0.287
Epoch: 5 [16640/60000 (28%)] 	Train loss: 2.186, acc: 0.288
Epoch: 5 [17920/60000 (30%)] 	Train loss: 2.183, acc: 0.290
Epoch: 5 [19200/60000 (32%)] 	Train loss: 2.181, acc: 0.292
Epoch: 5, time: 0.49 seconds


Epoch: 6 [0/60000 (0%)] 	Train loss: 2.138, acc: 0.341
Epoch: 6 [1280/60000 (2%)] 	Train loss: 2.136, acc: 0.326
Epoch: 6 [2560/60000 (4%)] 	Train loss: 2.135, acc: 0.323
Epoch: 6 [3840/60000 (6%)] 	Train loss: 2.131, acc: 0.325
Epoch: 6 [5120/60000 (9%)] 	Train loss: 2.130, acc: 0.324
Epoch: 6 [6400/60000 (11%)] 	Train loss: 2.125, acc: 0.327
Epoch: 6 [7680/60000 (13%)] 	Train loss: 2.123, acc: 0.327
Epoch: 6 [8960/60000 (15%)] 	Train loss: 2.119, acc: 0.329
Epoch: 6 [10240/60000 (17%)] 	Train loss: 2.114, acc: 0.332
Epoch: 6 [11520/60000 (19%)] 	Train loss: 2.111, acc: 0.334
Epoch: 6 [12800/60000 (21%)] 	Train loss: 2.107, acc: 0.334
Epoch: 6 [14080/60000 (23%)] 	Train loss: 2.102, acc: 0.334
Epoch: 6 [15360/60000 (26%)] 	Train loss: 2.099, acc: 0.334
Epoch: 6 [16640/60000 (28%)] 	Train loss: 2.095, acc: 0.335
Epoch: 6 [17920/60000 (30%)] 	Train loss: 2.091, acc: 0.336
Epoch: 6 [19200/60000 (32%)] 	Train loss: 2.086, acc: 0.337
Epoch: 6, time: 0.49 seconds


Epoch: 7 [0/60000 (0%)] 	Train loss: 2.008, acc: 0.372
Epoch: 7 [1280/60000 (2%)] 	Train loss: 2.004, acc: 0.359
Epoch: 7 [2560/60000 (4%)] 	Train loss: 2.000, acc: 0.357
Epoch: 7 [3840/60000 (6%)] 	Train loss: 1.993, acc: 0.362
Epoch: 7 [5120/60000 (9%)] 	Train loss: 1.989, acc: 0.362
Epoch: 7 [6400/60000 (11%)] 	Train loss: 1.981, acc: 0.368
Epoch: 7 [7680/60000 (13%)] 	Train loss: 1.976, acc: 0.368
Epoch: 7 [8960/60000 (15%)] 	Train loss: 1.970, acc: 0.372
Epoch: 7 [10240/60000 (17%)] 	Train loss: 1.961, acc: 0.376
Epoch: 7 [11520/60000 (19%)] 	Train loss: 1.955, acc: 0.380
Epoch: 7 [12800/60000 (21%)] 	Train loss: 1.948, acc: 0.383
Epoch: 7 [14080/60000 (23%)] 	Train loss: 1.940, acc: 0.387
Epoch: 7 [15360/60000 (26%)] 	Train loss: 1.933, acc: 0.390
Epoch: 7 [16640/60000 (28%)] 	Train loss: 1.927, acc: 0.394
Epoch: 7 [17920/60000 (30%)] 	Train loss: 1.919, acc: 0.398
Epoch: 7 [19200/60000 (32%)] 	Train loss: 1.912, acc: 0.402
Epoch: 7, time: 0.49 seconds


Epoch: 8 [0/60000 (0%)] 	Train loss: 1.782, acc: 0.510
Epoch: 8 [1280/60000 (2%)] 	Train loss: 1.775, acc: 0.464
Epoch: 8 [2560/60000 (4%)] 	Train loss: 1.768, acc: 0.469
Epoch: 8 [3840/60000 (6%)] 	Train loss: 1.757, acc: 0.477
Epoch: 8 [5120/60000 (9%)] 	Train loss: 1.751, acc: 0.479
Epoch: 8 [6400/60000 (11%)] 	Train loss: 1.739, acc: 0.486
Epoch: 8 [7680/60000 (13%)] 	Train loss: 1.732, acc: 0.489
Epoch: 8 [8960/60000 (15%)] 	Train loss: 1.723, acc: 0.494
Epoch: 8 [10240/60000 (17%)] 	Train loss: 1.712, acc: 0.499
Epoch: 8 [11520/60000 (19%)] 	Train loss: 1.703, acc: 0.503
Epoch: 8 [12800/60000 (21%)] 	Train loss: 1.693, acc: 0.508
Epoch: 8 [14080/60000 (23%)] 	Train loss: 1.682, acc: 0.512
Epoch: 8 [15360/60000 (26%)] 	Train loss: 1.674, acc: 0.516
Epoch: 8 [16640/60000 (28%)] 	Train loss: 1.665, acc: 0.520
Epoch: 8 [17920/60000 (30%)] 	Train loss: 1.655, acc: 0.524
Epoch: 8 [19200/60000 (32%)] 	Train loss: 1.646, acc: 0.528
Epoch: 8, time: 0.49 seconds


Epoch: 9 [0/60000 (0%)] 	Train loss: 1.483, acc: 0.604
Epoch: 9 [1280/60000 (2%)] 	Train loss: 1.478, acc: 0.576
Epoch: 9 [2560/60000 (4%)] 	Train loss: 1.471, acc: 0.584
Epoch: 9 [3840/60000 (6%)] 	Train loss: 1.458, acc: 0.595
Epoch: 9 [5120/60000 (9%)] 	Train loss: 1.452, acc: 0.598
Epoch: 9 [6400/60000 (11%)] 	Train loss: 1.438, acc: 0.605
Epoch: 9 [7680/60000 (13%)] 	Train loss: 1.431, acc: 0.608
Epoch: 9 [8960/60000 (15%)] 	Train loss: 1.422, acc: 0.610
Epoch: 9 [10240/60000 (17%)] 	Train loss: 1.411, acc: 0.614
Epoch: 9 [11520/60000 (19%)] 	Train loss: 1.402, acc: 0.618
Epoch: 9 [12800/60000 (21%)] 	Train loss: 1.392, acc: 0.621
Epoch: 9 [14080/60000 (23%)] 	Train loss: 1.381, acc: 0.624
Epoch: 9 [15360/60000 (26%)] 	Train loss: 1.373, acc: 0.626
Epoch: 9 [16640/60000 (28%)] 	Train loss: 1.364, acc: 0.628
Epoch: 9 [17920/60000 (30%)] 	Train loss: 1.355, acc: 0.631
Epoch: 9 [19200/60000 (32%)] 	Train loss: 1.347, acc: 0.634
Epoch: 9, time: 0.49 seconds


Epoch: 10 [0/60000 (0%)] 	Train loss: 1.200, acc: 0.677
Epoch: 10 [1280/60000 (2%)] 	Train loss: 1.196, acc: 0.661
Epoch: 10 [2560/60000 (4%)] 	Train loss: 1.191, acc: 0.668
Epoch: 10 [3840/60000 (6%)] 	Train loss: 1.177, acc: 0.679
Epoch: 10 [5120/60000 (9%)] 	Train loss: 1.173, acc: 0.679
Epoch: 10 [6400/60000 (11%)] 	Train loss: 1.160, acc: 0.685
Epoch: 10 [7680/60000 (13%)] 	Train loss: 1.155, acc: 0.686
Epoch: 10 [8960/60000 (15%)] 	Train loss: 1.148, acc: 0.688
Epoch: 10 [10240/60000 (17%)] 	Train loss: 1.140, acc: 0.690
Epoch: 10 [11520/60000 (19%)] 	Train loss: 1.132, acc: 0.693
Epoch: 10 [12800/60000 (21%)] 	Train loss: 1.125, acc: 0.696
Epoch: 10 [14080/60000 (23%)] 	Train loss: 1.115, acc: 0.698
Epoch: 10 [15360/60000 (26%)] 	Train loss: 1.109, acc: 0.699
Epoch: 10 [16640/60000 (28%)] 	Train loss: 1.103, acc: 0.701
Epoch: 10 [17920/60000 (30%)] 	Train loss: 1.096, acc: 0.703
Epoch: 10 [19200/60000 (32%)] 	Train loss: 1.090, acc: 0.704
Epoch: 10, time: 0.49 seconds


Epoch: 11 [0/60000 (0%)] 	Train loss: 0.986, acc: 0.745
Epoch: 11 [1280/60000 (2%)] 	Train loss: 0.980, acc: 0.727
Epoch: 11 [2560/60000 (4%)] 	Train loss: 0.980, acc: 0.728
Epoch: 11 [3840/60000 (6%)] 	Train loss: 0.966, acc: 0.736
Epoch: 11 [5120/60000 (9%)] 	Train loss: 0.966, acc: 0.734
Epoch: 11 [6400/60000 (11%)] 	Train loss: 0.955, acc: 0.739
Epoch: 11 [7680/60000 (13%)] 	Train loss: 0.952, acc: 0.739
Epoch: 11 [8960/60000 (15%)] 	Train loss: 0.948, acc: 0.739
Epoch: 11 [10240/60000 (17%)] 	Train loss: 0.943, acc: 0.739
Epoch: 11 [11520/60000 (19%)] 	Train loss: 0.937, acc: 0.740
Epoch: 11 [12800/60000 (21%)] 	Train loss: 0.933, acc: 0.742
Epoch: 11 [14080/60000 (23%)] 	Train loss: 0.925, acc: 0.744
Epoch: 11 [15360/60000 (26%)] 	Train loss: 0.922, acc: 0.744
Epoch: 11 [16640/60000 (28%)] 	Train loss: 0.917, acc: 0.745
Epoch: 11 [17920/60000 (30%)] 	Train loss: 0.913, acc: 0.746
Epoch: 11 [19200/60000 (32%)] 	Train loss: 0.909, acc: 0.746
Epoch: 11, time: 0.49 seconds


Epoch: 12 [0/60000 (0%)] 	Train loss: 0.847, acc: 0.760
Epoch: 12 [1280/60000 (2%)] 	Train loss: 0.839, acc: 0.764
Epoch: 12 [2560/60000 (4%)] 	Train loss: 0.842, acc: 0.763
Epoch: 12 [3840/60000 (6%)] 	Train loss: 0.830, acc: 0.769
Epoch: 12 [5120/60000 (9%)] 	Train loss: 0.832, acc: 0.766
Epoch: 12 [6400/60000 (11%)] 	Train loss: 0.823, acc: 0.769
Epoch: 12 [7680/60000 (13%)] 	Train loss: 0.821, acc: 0.769
Epoch: 12 [8960/60000 (15%)] 	Train loss: 0.820, acc: 0.769
Epoch: 12 [10240/60000 (17%)] 	Train loss: 0.816, acc: 0.769
Epoch: 12 [11520/60000 (19%)] 	Train loss: 0.812, acc: 0.770
Epoch: 12 [12800/60000 (21%)] 	Train loss: 0.809, acc: 0.771
Epoch: 12 [14080/60000 (23%)] 	Train loss: 0.803, acc: 0.772
Epoch: 12 [15360/60000 (26%)] 	Train loss: 0.801, acc: 0.772
Epoch: 12 [16640/60000 (28%)] 	Train loss: 0.798, acc: 0.772
Epoch: 12 [17920/60000 (30%)] 	Train loss: 0.795, acc: 0.773
Epoch: 12 [19200/60000 (32%)] 	Train loss: 0.793, acc: 0.774
Epoch: 12, time: 0.49 seconds


Epoch: 13 [0/60000 (0%)] 	Train loss: 0.757, acc: 0.786
Epoch: 13 [1280/60000 (2%)] 	Train loss: 0.746, acc: 0.791
Epoch: 13 [2560/60000 (4%)] 	Train loss: 0.751, acc: 0.788
Epoch: 13 [3840/60000 (6%)] 	Train loss: 0.740, acc: 0.793
Epoch: 13 [5120/60000 (9%)] 	Train loss: 0.743, acc: 0.789
Epoch: 13 [6400/60000 (11%)] 	Train loss: 0.735, acc: 0.791
Epoch: 13 [7680/60000 (13%)] 	Train loss: 0.735, acc: 0.791
Epoch: 13 [8960/60000 (15%)] 	Train loss: 0.734, acc: 0.790
Epoch: 13 [10240/60000 (17%)] 	Train loss: 0.732, acc: 0.789
Epoch: 13 [11520/60000 (19%)] 	Train loss: 0.728, acc: 0.790
Epoch: 13 [12800/60000 (21%)] 	Train loss: 0.726, acc: 0.791
Epoch: 13 [14080/60000 (23%)] 	Train loss: 0.721, acc: 0.792
Epoch: 13 [15360/60000 (26%)] 	Train loss: 0.720, acc: 0.792
Epoch: 13 [16640/60000 (28%)] 	Train loss: 0.718, acc: 0.792
Epoch: 13 [17920/60000 (30%)] 	Train loss: 0.715, acc: 0.793
Epoch: 13 [19200/60000 (32%)] 	Train loss: 0.714, acc: 0.793
Epoch: 13, time: 0.49 seconds


Epoch: 14 [0/60000 (0%)] 	Train loss: 0.695, acc: 0.797
Epoch: 14 [1280/60000 (2%)] 	Train loss: 0.680, acc: 0.808
Epoch: 14 [2560/60000 (4%)] 	Train loss: 0.685, acc: 0.803
Epoch: 14 [3840/60000 (6%)] 	Train loss: 0.675, acc: 0.808
Epoch: 14 [5120/60000 (9%)] 	Train loss: 0.679, acc: 0.805
Epoch: 14 [6400/60000 (11%)] 	Train loss: 0.671, acc: 0.807
Epoch: 14 [7680/60000 (13%)] 	Train loss: 0.672, acc: 0.807
Epoch: 14 [8960/60000 (15%)] 	Train loss: 0.672, acc: 0.806
Epoch: 14 [10240/60000 (17%)] 	Train loss: 0.670, acc: 0.806
Epoch: 14 [11520/60000 (19%)] 	Train loss: 0.667, acc: 0.808
Epoch: 14 [12800/60000 (21%)] 	Train loss: 0.666, acc: 0.808
Epoch: 14 [14080/60000 (23%)] 	Train loss: 0.661, acc: 0.809
Epoch: 14 [15360/60000 (26%)] 	Train loss: 0.660, acc: 0.809
Epoch: 14 [16640/60000 (28%)] 	Train loss: 0.658, acc: 0.809
Epoch: 14 [17920/60000 (30%)] 	Train loss: 0.657, acc: 0.809
Epoch: 14 [19200/60000 (32%)] 	Train loss: 0.656, acc: 0.809
Epoch: 14, time: 0.49 seconds


Exporting MNIST model to S3

Verify TrainJob Completion

Check the TrainJob status to ensure it completed successfully.

In [ ]:

TrainerClient().wait_for_job_status(name=job_id, timeout=20)

Evaluate the Trained Model

Download the trained model from S3 to evaluate it.

We will use test images from the MNIST dataset for prediction.

• Green label indicate correct predictions.
• Red label indicate incorrect predictions, with the correct value shown in parentheses.

In [10]:

import boto3

boto3.client("s3").download_file(BUCKET, MODEL, MODEL)

In [12]:

import mlx.core as mx
import mlx.nn as nn
import matplotlib.pyplot as plt
from mlx.data.datasets import load_mnist


# Load test dataset and take the random batch from it.
test_batch = (
    load_mnist(train=False)
    .key_transform("image", lambda x: (x.astype("float32") / 255.0).ravel())
    .batch(20)[mx.random.randint(10, 500)]
)

class MLP(nn.Module):
    def __init__(self, in_dims: int, hidden_dims: int, num_layers: int, out_dims: int):
        super().__init__()
        layer_sizes = [in_dims] + [hidden_dims] * num_layers + [out_dims]
        self.layers = [
            nn.Linear(idim, odim)
            for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
        ]

    def __call__(self, x):
        for layer in self.layers[:-1]:
            x = nn.relu(layer(x))
        return self.layers[-1](x)

# Load weights from the trained model.
model = MLP(
    in_dims=test_batch["image"][0].shape[-1],
    hidden_dims=32,
    num_layers=2,
    out_dims=10,
).load_weights(MODEL)

# Send test batch to the pre-trained MLP model.
x = mx.array(test_batch["image"])
output = model(x)
fig = plt.figure(figsize=(15, 5))
for i in range(20):
    # Format the input image and the model output.
    image = test_batch["image"][i].reshape((28, 28))
    pred_label = mx.argmax(output[i])

    # Add data to the plot.
    ax = fig.add_subplot(4, 5, i + 1, xticks=[], yticks=[])
    ax.imshow(image, cmap="gray")
    if test_batch["label"][i] == pred_label:
        ax.set_title(test_batch["label"][i], color="green")
    else:
        ax.set_title("{} ({})".format(pred_label, test_batch["label"][i]), color="red")

    ax.title.set_fontsize(20)
    fig.tight_layout()

Delete the TrainJob

When TrainJob is finished, you can delete the resource.

In [ ]:

TrainerClient().delete_job(job_id)