Bring Your Models to the Edge: PyTorch Mobile Deployment#

Common Challenges of Edge Deployment#

• Limited Computational Resources: Mobile CPUs/GPUs are less powerful than their server counterparts.
• Battery Constraints: Constant heavy processing can quickly drain a device’s battery.
• Model Size Limits: Large models can be impractical to store on devices with limited space, and can also be slow to run.

However, with careful planning—such as using optimized libraries, quantization, and pruning—these challenges can be mitigated.

Key Features#

1. Small Footprint: The PyTorch Mobile runtime is smaller in size compared to the full PyTorch library.
2. Cross-Platform: Official support and documentation for both Android and iOS.
3. Quantization and Special Ops: Built-in tools to reduce model size and increase inference speed.

Dependencies#

To follow along with this tutorial, you’ll need:

• Python 3.7 or later
• PyTorch (desktop installation)
• TorchVision (if you plan to work with vision models)
• A modern code editor or IDE
• Android Studio (for Android deployment)
• Xcode (for iOS deployment)
• Java Development Kit (JDK) for Android builds

Installing PyTorch Mobile Libraries#

When targeting mobile platforms, you need specific artifacts for PyTorch Mobile. For Android, you’ll include them in your Gradle files. For iOS, you’ll use CocoaPods or Swift Package Manager.

On your desktop Python environment, make sure you have the latest version of PyTorch installed:

1
pip install torch torchvision

Verifying the Installation#

You can check your PyTorch version by running:

1
import torch
2
print(torch.__version__)

Ensure you see the version you installed. If you’re using a GPU, also check:

1
print(torch.cuda.is_available())

Though the GPU check is more useful on desktop environments; on mobile, you’ll typically rely on CPU or specialized libraries.

Model Definition#

Let’s build a super-simple convolutional neural network for demonstration. This model will be small enough to run on mobile devices without too much overhead. For example:

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4

5
class SimpleCNN(nn.Module):
6
    def __init__(self, num_classes=10):
7
        super(SimpleCNN, self).__init__()
8
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3)
9
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3)
10
        self.fc1 = nn.Linear(32*6*6, 128)
11
        self.fc2 = nn.Linear(128, num_classes)
12

13
    def forward(self, x):
14
        x = F.relu(self.conv1(x))
15
        x = F.max_pool2d(x, 2)
16
        x = F.relu(self.conv2(x))
17
        x = F.max_pool2d(x, 2)
18
        x = x.view(x.size(0), -1)
19
        x = F.relu(self.fc1(x))
20
        x = self.fc2(x)
21
        return x

Training the Model#

For the sake of brevity, let’s assume you have a dataset ready. Typically, you’d use something like the CIFAR-10 dataset for a classification task. Below is a simplified training loop:

1
from torch.utils.data import DataLoader
2
import torch.optim as optim
3

4
# Suppose 'train_dataset' is your training dataset
5
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
6
model = SimpleCNN(num_classes=10)
7
optimizer = optim.Adam(model.parameters(), lr=0.001)
8
criterion = nn.CrossEntropyLoss()
9

10
for epoch in range(10):
11
    model.train()
12
    running_loss = 0.0
13

14
    for images, labels in train_loader:
15
        optimizer.zero_grad()
16
        outputs = model(images)
17
        loss = criterion(outputs, labels)
18
        loss.backward()
19
        optimizer.step()
20
        running_loss += loss.item()
21

22
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}")

Saving the Model#

After training is complete, you can save your model’s state dictionary or the entire model:

1
MODEL_PATH = "simple_cnn.pth"
2
torch.save(model.state_dict(), MODEL_PATH)

Scripting vs Tracing#

PyTorch provides two main methods to convert a model to TorchScript:

1. Tracing: Traces the execution of the model with a dummy input, suitable for models with static control flow.
2. Scripting: Analyzes the source code of your model to create an executable graph, suitable for dynamic control flow.

For simpler models without complex branching logic, tracing can be sufficient. Otherwise, use scripting.

1
# Example of tracing
2
model = SimpleCNN()
3
model.load_state_dict(torch.load(MODEL_PATH))
4
model.eval()
5

6
example_input = torch.rand(1, 3, 28, 28)
7
traced_script_module = torch.jit.trace(model, example_input)

Converting to TorchScript#

Once you have either a traced script module or a scripted module, you’ll have a torch.jit.ScriptModule. You can save it:

1
mobile_model_path = "simple_cnn_mobile.pt"
2
traced_script_module.save(mobile_model_path)

Optimizing the Model for Mobile#

PyTorch Mobile provides an additional optimization step you can use:

1
from torch.utils.mobile_optimizer import optimize_for_mobile
2

3
optimized_scripted_module = optimize_for_mobile(traced_script_module)
4
optimized_scripted_module._save_for_lite_interpreter("simple_cnn_mobile_optimized.ptl")

This produces an optimized model (.ptl) specifically for mobile inference.

Android Project Setup#

1. Create a New Android Studio Project: Begin by creating a new Android project in Android Studio.
2. Gradle Configuration: Open the build.gradle for your app module and add dependencies for PyTorch Mobile.

Integrating the PyTorch Android Libraries#

In your module-level build.gradle, include:

1
dependencies {
2
    implementation 'org.pytorch:pytorch_android_lite:1.12.0'
3
    implementation 'org.pytorch:pytorch_android_torchvision_lite:1.12.0'
4
}

Make sure the versions match the PyTorch Mobile release you intend to use. Note that the “lite” versions refer to optimized libraries for mobile.

Loading and Running the Model#

In your Android code (e.g., an Activity), load the model file that you’ve placed in the assets folder.

1. Put the .ptl file in your project’s app/src/main/assets directory.
2. Use PyTorch Android APIs to load the model during runtime.

Example Android Code Snippet#

1
import org.pytorch.IValue;
2
import org.pytorch.Module;
3
import org.pytorch.Tensor;
4

5
public class MainActivity extends AppCompatActivity {
6

7
    private Module module;
8

9
    @Override
10
    protected void onCreate(Bundle savedInstanceState) {
11
        super.onCreate(savedInstanceState);
12
        setContentView(R.layout.activity_main);
13

14
        // Load the PyTorch model
15
        module = Module.load(assetFilePath("simple_cnn_mobile_optimized.ptl"));
16

17
        // Prepare input tensor
18
        float[] inputData = new float[3 * 28 * 28]; // Example
19
        Tensor inputTensor = Tensor.fromBlob(inputData, new long[]{1, 3, 28, 28});
20

21
        // Run inference
22
        IValue outputs = module.forward(IValue.from(inputTensor));
23
        float[] scores = outputs.toTensor().getDataAsFloatArray();
24

25
        // Process results (for example, find the class with the highest score)
26
        int predictedClass = argMax(scores);
27
        Log.d("PyTorchMobile", "Predicted class: " + predictedClass);
28
    }
29

30
    private String assetFilePath(String assetName) {
31
        File file = new File(getFilesDir(), assetName);
32
        // Copy the file from assets if not already present
33
        // Implementation omitted for brevity
34
        return file.getAbsolutePath();
35
    }
36

37
    private int argMax(float[] array) {
38
        int maxIndex = 0;
39
        for (int i = 1; i < array.length; i++) {
40
            if (array[i] > array[maxIndex]) {
41
                maxIndex = i;
42
            }
43
        }
44
        return maxIndex;
45
    }
46
}

iOS Project Setup#

1. Create a New Xcode Project: iOS Single View App or SwiftUI App.
2. Add PyTorch Mobile: You can use CocoaPods or Swift Package Manager to integrate PyTorch Mobile.

Integrating the PyTorch iOS Libraries#

If you’re using CocoaPods, add the following to your Podfile:

1
target 'YourApp' do
2
  pod 'LibTorch-Lite', '~> 1.12.0'
3
  pod 'LibTorch-Lite-vision', '~> 1.12.0'
4
end

Run pod install in your terminal. Then open the .xcworkspace project file in Xcode.

Loading and Running the Model#

Add your .ptl model file to your Xcode project (e.g., in the “Resources” folder). Make sure the file is included in the app’s build resources.

Example iOS Code Snippet#

1
import UIKit
2
import AVFoundation
3

4
class ViewController: UIViewController {
5
    var module: TorchModule?
6

7
    override func viewDidLoad() {
8
        super.viewDidLoad()
9

10
        if let filePath = Bundle.main.path(forResource: "simple_cnn_mobile_optimized", ofType: "ptl") {
11
            self.module = TorchModule(fileAtPath: filePath)
12
        } else {
13
            fatalError("Model file not found.")
14
        }
15

16
        // Prepare input tensor
17
        let dataCount = 3 * 28 * 28
18
        var inputData = [Float](repeating: 0.0, count: dataCount)
19
        let tensor = TorchTensor(fromArray: &inputData, shape: [1, 3, 28, 28])
20

21
        // Run inference
22
        guard let outputTensor = module?.predict(imageTensor: tensor) else {
23
            fatalError("Model inference failed.")
24
        }
25

26
        let scores = outputTensor.toArray()
27
        let predictedClass = argMax(scores)
28
        print("Predicted class: \(predictedClass)")
29
    }
30

31
    func argMax(_ array: [Float]) -> Int {
32
        var maxIndex = 0
33
        for i in 1..<array.count {
34
            if array[i] > array[maxIndex] {
35
                maxIndex = i
36
            }
37
        }
38
        return maxIndex
39
    }
40
}

The key differences from Android revolve around the way you include libraries and resources, but the conceptual flow (load model, create tensor, run inference, process output) remains the same.

Quantization#

Quantization reduces the precision of weights and/or activations from 32-bit floating-point to lower-bit representations (e.g., 8-bit). This can significantly reduce model size and speed up inference on certain hardware. PyTorch supports:

1. Static Quantization: Quantizing weights and activations for the entire model.
2. Dynamic Quantization: Quantizing certain layers (like fully connected layers) on the fly.

Here’s a quick example of static quantization:

1
import torch.quantization as quant
2

3
model_fp32 = SimpleCNN()
4
model_fp32.load_state_dict(torch.load("simple_cnn.pth"))
5
model_fp32.eval()
6

7
# Fuse modules where applicable. E.g., fuse conv + bn + relu
8
# model_fp32_fused = torch.quantization.fuse_modules(model_fp32, ...)
9

10
model_qconfig = quant.get_default_qconfig("fbgemm")
11
model_fp32.qconfig = model_qconfig
12

13
torch.quantization.prepare(model_fp32, inplace=True)
14
# Calibrate with a representative dataset
15
for images, _ in calibration_loader:
16
    _ = model_fp32(images)
17

18
torch.quantization.convert(model_fp32, inplace=True)
19

20
# Save quantized model
21
torch.jit.script(model_fp32).save("simple_cnn_quantized.pt")

Pruning#

Pruning removes weights that have minimal impact on the output. PyTorch’s pruning modules allow channel pruning or structured pruning, potentially reducing the model size further.

Hardware Acceleration#

In some cases, you might leverage proprietary libraries or specialized accelerators (e.g., Apple’s Neural Engine via Core ML, or Android’s NNAPI). PyTorch Mobile supports certain integrations, but usage can vary depending on hardware constraints.

Best Practices#

1. Benchmark Early: Measure inference speed and memory usage on actual devices.
2. Profile Memory: Keep an eye on how your model loads and if device memory is sufficient.
3. Optimize for Battery: Use smaller, quantized models and minimize background inference tasks.
4. Iterative Testing: Integrate with your CI/CD pipeline to test model performance on various device configurations.

Dynamic Quantization#

While static quantization quantizes both weights and activations, dynamic quantization typically only quantizes the weights by converting activations to lower precision “on the fly.” This is faster to implement and often sufficient for many linear-based layers like LSTM or fully connected.

1
model = SimpleCNN()
2
model.load_state_dict(torch.load("simple_cnn.pth"))
3
model.eval()
4

5
quantized_model = torch.quantization.quantize_dynamic(
6
    model, {nn.Linear}, dtype=torch.qint8
7
)
8

9
# Save dynamic quantized model
10
quantized_model_scripted = torch.jit.script(quantized_model)
11
quantized_model_scripted.save("simple_cnn_dynamic_quantized.pt")

Domain-Specific Optimizations#

• Vision: Use specialized operations like torchvision.ops or hardware-accelerated libraries.
• NLP: Consider subword tokenization and quantization for word embedding layers.

Versioning and A/B Testing#

Maintaining multiple versions of your model is often necessary for continuous improvement and canary releases. Consider storing versioned models and using environment-based toggles. You can serve different model versions to different user segments for live performance monitoring.

Offline Updates and Rollbacks#

Mobile apps can fetch updated model files from a remote server when connected, caching them locally. If performance regresses or issues arise, store the previous model to roll back. This ensures minimal downtime and quick recovery.