Comparing Deep Learning on TensorFlow (Keras), PyTorch, and Apple’s MLX

Deep learning frameworks such as TensorFlow (Keras), PyTorch, and Apple’s MLX offer powerful tools to build and train machine learning models. Despite solving similar problems, these frameworks have different philosophies, APIs, and optimizations under the hood. In this post, we will examine how the same model is implemented on each platform and why the differences in code arise, especially focusing on why MLX is more similar to PyTorch than TensorFlow.

1. Model in PyTorch

PyTorch is known for giving developers granular control over model-building and training processes. The framework encourages writing custom training loops, making it highly flexible, especially for research purposes.

PyTorch Code:

</p><p>import torch</p><p>import torch.optim as optim</p><p>import torch.nn as nn</p><p>import torch.nn.functional as F</p><p># Define a simple model</p><p>class SimpleModel(nn.Module):</p><p>def __init__(self):</p><p>super(SimpleModel, self).__init__()</p><p>self.fc1 = nn.Linear(32, 64)</p><p>self.fc2 = nn.Linear(64, 10)</p><p>def forward(self, x):</p><p>x = F.relu(self.fc1(x))</p><p>return self.fc2(x)</p><p>model = SimpleModel()</p><p>optimizer = optim.SGD(model.parameters(), lr=0.01)</p><p>criterion = nn.CrossEntropyLoss()</p><p># Training loop</p><p>for epoch in range(10):</p><p>for inputs, labels in train_loader:</p><p>optimizer.zero_grad()</p><p>outputs = model(inputs)</p><p>loss = criterion(outputs, labels)</p><p>loss.backward()</p><p>torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)</p><p>optimizer.step()</p><p>

What’s Happening Behind the Scenes in PyTorch?

PyTorch gives the developer direct control over every step of the model training process. The training loop is written manually, where:

• Forward pass: Defined in the forward() method, explicitly computing the output layer by layer.
• Backward pass: After calculating the loss, the gradients are computed using loss.backward().
• Gradient updates: The optimizer manually updates the weights after each batch using optimizer.step().

This manual training loop allows researchers and developers to experiment with unconventional architectures or optimization methods. The gradient clipping function torch.nn.utils.clip_grad_norm_ prevents exploding gradients during backpropagation.

2. Model in TensorFlow (Keras)

TensorFlow with Keras abstracts many of the low-level operations you see in PyTorch. Its primary goal is to make model-building and training easy and fast, reducing the amount of boilerplate code.

TensorFlow (Keras) Code:

</p><p>import tensorflow as tf</p><p>from tensorflow.keras import layers, models, optimizers</p><p># Define a simple model</p><p>model = models.Sequential([</p><p>layers.Dense(64, activation='relu', input_shape=(32,)),</p><p>layers.Dense(10)</p><p>])</p><p># Compile the model with gradient clipping</p><p>optimizer = optimizers.SGD(learning_rate=0.01, clipnorm=1.0)</p><p>model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy')</p><p># Train the model</p><p>model.fit(X_train, y_train, epochs=10, batch_size=32)</p><p>

What’s Happening Behind the Scenes in TensorFlow (Keras)?

TensorFlow abstracts much of the model-building process, particularly:

• Automatic Training Loop: The model.fit() method handles the forward pass, backpropagation, and gradient updates. This makes the code much more concise.
• Gradient Clipping: Can be added directly into the optimizer with clipnorm, without needing explicit control of the gradients as in PyTorch.

TensorFlow’s design philosophy is to reduce complexity for users by providing higher-level APIs. However, this abstraction means that developers may have less control over the finer details of the training process. TensorFlow is often preferred for production environments or when rapid prototyping is required.

3. Model in MLX (Apple Silicon)

MLX is Apple’s machine learning framework designed specifically for Apple Silicon (M1/M2) hardware. While it shares many similarities with PyTorch, it is optimized to take full advantage of the unified memory architecture and metal-accelerated computations of Apple devices.

MLX Code:

</p><p>import mlx.core as mx</p><p>import mlx.nn as nn</p><p>import mlx.optimizers as optim</p><p># Define a simple model</p><p>class SimpleModel(nn.Module):</p><p>def __init__(self):</p><p>super().__init__()</p><p>self.fc1 = nn.Linear(32, 64)</p><p>self.fc2 = nn.Linear(64, 10)</p><p>def __call__(self, x):</p><p>x = mx.maximum(self.fc1(x), 0) # ReLU activation</p><p>return self.fc2(x)</p><p>model = SimpleModel()</p><p># Define optimizer and loss function</p><p>optimizer = optim.SGD(learning_rate=0.01)</p><p>loss_fn = nn.losses.cross_entropy</p><p># Training loop</p><p>for epoch in range(10):</p><p>for inputs, labels in train_loader:</p><p>outputs = model(inputs)</p><p>loss = loss_fn(outputs, labels)</p><p>grads = nn.value_and_grad(model, loss_fn)(model, inputs, labels)</p><p>optimizer.update(model, grads)</p><p>

What’s Happening Behind the Scenes in MLX?

MLX is conceptually similar to PyTorch, with a manual training loop and explicit model-building, but it’s optimized for Apple hardware:

• Unified Memory Architecture: Unlike TensorFlow and PyTorch, MLX is built to leverage Apple Silicon’s unified memory, meaning data does not need to be moved between CPU and GPU, improving performance.
• Dynamic Graph Construction: Like PyTorch, MLX dynamically constructs computation graphs, allowing for more flexibility in model design and debugging.
• Manual Training Loop: Just like in PyTorch, the user has to manually handle forward passes, backward passes, and weight updates.

Why MLX is More Similar to PyTorch than TensorFlow

MLX’s philosophy is closer to PyTorch’s in terms of flexibility and control over the training process:

• Manual Training Loop: Like PyTorch, MLX allows you to define every part of the training loop manually. This makes MLX more suitable for research and experimentation, similar to PyTorch.
• Dynamic Graphs: Both PyTorch and MLX use dynamic computation graphs (as opposed to TensorFlow’s static graphs), which provide more flexibility and ease of debugging.
• Hardware Optimization: While TensorFlow can also leverage hardware acceleration (like GPUs or TPUs), MLX is highly optimized for Apple’s own silicon chips, making it more efficient for Apple devices.

In contrast, TensorFlow focuses on abstraction and simplicity, using high-level APIs like Sequential and fit() to automatically handle many aspects of model training.

Why the Code is Different Across Platforms

• Philosophy: PyTorch and MLX focus on research flexibility and control. You can manually design the training loop and control every aspect of computation, making them suitable for users who need fine-tuned control or custom architectures. TensorFlow prioritizes simplicity and productivity, abstracting many complexities.
• Memory Management
  : MLX is built specifically for Apple Silicon, taking advantage of the unified memory architecture to minimize data transfers between CPU and GPU, improving performance.
• Graph Construction: PyTorch and MLX both use dynamic graph construction, whereas TensorFlow traditionally relied on static graphs, although it now in new versions it supports dynamic graphs as well.
• Hardware Optimization: MLX is optimized for Apple Silicon, allowing it to take advantage of Apple’s unified memory architecture, which means MLX can efficiently manage data across the CPU and GPU without having to copy data between them.

Conclusion

Each of these frameworks has its strengths, but which one to choose depends on the specific needs of your project:

• PyTorch: Renowned for its flexibility and intuitive design, PyTorch is ideal for research and experimentation. It allows for dynamic computation graphs, providing researchers with the ability to modify models on-the-fly. This adaptability makes it a preferred choice in academic settings.

• TensorFlow (with Keras): TensorFlow, especially when used with the Keras API, offers a user-friendly experience that abstracts many complexities of deep learning. It’s well-suited for production environments due to its robustness and scalability. TensorFlow’s comprehensive ecosystem supports deployment across various platforms, making it a strong candidate for production-level projects. Notably, with the introduction of eager execution in TensorFlow 2.0, it now supports dynamic computation graphs, enhancing its flexibility.

  TensorFlow

   

• MLX: If you’re developing machine learning models on Apple Silicon hardware, MLX is designed to leverage Apple’s unified memory architecture, maximizing performance and efficiency. This framework is optimized for Apple’s hardware, ensuring seamless integration and optimal resource utilization.

In summary:

• Choose PyTorch for research-focused projects that require flexibility and dynamic model building.

• Opt for TensorFlow (Keras) when developing production-level applications that benefit from a stable and scalable framework, now with enhanced flexibility due to its support for dynamic computation graphs.

• Select MLX to harness the full potential of Apple Silicon hardware for machine learning tasks.