{"id":3443,"date":"2025-06-26T08:00:12","date_gmt":"2025-06-26T15:00:12","guid":{"rendered":"https:\/\/devblogs.microsoft.com\/azure-sdk\/?p=3443"},"modified":"2025-06-25T14:31:54","modified_gmt":"2025-06-25T21:31:54","slug":"introducing-the-azure-storage-connector-for-pytorch","status":"publish","type":"post","link":"https:\/\/devblogs.microsoft.com\/azure-sdk\/introducing-the-azure-storage-connector-for-pytorch\/","title":{"rendered":"Introducing the Azure Storage Connector for PyTorch"},"content":{"rendered":"

We’re excited to introduce the Azure Storage Connector for PyTorch (azstoragetorch<\/code>), a new library that brings seamless, performance-optimized integration between Azure Storage and PyTorch. The library makes it easy to access and store data in Azure Blob Storage directly within your training workflows.<\/p>\n

What is PyTorch?<\/h2>\n

PyTorch<\/a> is a widely used open-source machine learning framework, known for its flexibility and strong support for both research and production deployments. Training models with PyTorch often involves handling large volumes of data. This can include loading massive datasets, saving and restoring model checkpoints, and managing data pipelines. These pipelines can live across environments like local machines, cloud virtual machines, and distributed compute clusters, adding to complexity.<\/p>\n

How can the Azure Storage Connector for PyTorch help?<\/h2>\n

The Azure Storage Connector for PyTorch (azstoragetorch<\/code>) bridges the powerful storage capabilities of Azure Storage with PyTorch, enabling seamless integrations with key PyTorch APIs for your model training workflows.<\/p>\n

The package supports model checkpointing directly with Azure Storage with torch.save()<\/code> and torch.load()<\/code> and directly loading data from Azure Storage to PyTorch Dataset<\/code> classes.<\/p>\n

Use the Azure Storage Connector for PyTorch<\/h2>\n

Installation<\/h3>\n

The Azure Storage Connector for PyTorch is listed on PyPI, and you can install it using your favorite package manager. This example utilizes pip<\/code>:<\/p>\n

pip install azstoragetorch<\/code><\/pre>\n
This installs the azstoragetorch<\/code> library and other dependencies such as torch<\/code> and azure-storage-blob<\/code>.<\/p>\n
Authentication<\/h3>\n
The Azure Storage Connector for PyTorch package’s interfaces default to using the Azure Identity library’s DefaultAzureCredential<\/code><\/a>, which automatically retrieves Microsoft Entra ID tokens based on your current environment. For more information, see DefaultAzureCredential overview<\/a>.<\/p>\n
This means that as long as there exists a credential on your machine in the chained credential list for DefaultAzureCredential<\/code>, your credentials are securely handled automatically and you’re ready to start using the library.<\/p>\n
Save and Load a Model Checkpoint<\/h3>\n
The Azure Storage Connector for PyTorch provides the azstoragetorch.io.BlobIO<\/code><\/a> class to save and load models directly to and from Azure Blob Storage. This class adheres to the file-like object accepted when using torch.save()<\/code> and torch.load()<\/code> for model checkpointing in PyTorch.<\/p>\n
The BlobIO<\/code> class takes an Azure Storage Blob URL and the mode you’d like to operate in – either write (\"wb\"<\/code>) for saving or read (\"rb\"<\/code>) for loading. Because the BlobIO<\/code> class is file-like in Python, it can be safely handled when using the with<\/code> statement.<\/p>\n
import torch\r\nimport torchvision.models # Install this separately, e.g. ``pip install torchvision``\r\nfrom azstoragetorch.io import BlobIO\r\n\r\n# The URL to your container\r\nCONTAINER_URL = \"https:\/\/<my-storage-account-name>.blob.core.windows.net\/<my-container-name>\"\r\n\r\n# Your model of choice, in this case ResNet18 for image recognition   \r\nmodel = torchvision.models.resnet18(weights=\"DEFAULT\")\r\n\r\n# Save model weights to an Azure Storage Blob named \"model_weights.pth\" in the container\r\nwith BlobIO(f\"{CONTAINER_URL}\/model_weights.pth\", \"wb\") as f:\r\n    torch.save(model.state_dict(), f)\r\n\r\n# Load model weights from model_weights.pth in the Azure Storage Container\r\nwith BlobIO(f\"{CONTAINER_URL}\/model_weights.pth\", \"rb\") as f:\r\n    model.load_state_dict(torch.load(f))<\/code><\/pre>\n
Sample: Use PyTorch Datasets on Azure Storage<\/h3>\n
PyTorch has two primitives for loading samples, DataSet<\/code> and DataLoader<\/code><\/a>. The Azure Storage Connector for PyTorch has implementations for both of PyTorch datasets, map-style<\/a> and iterable-style<\/a>.<\/p>\n
The azstoragetorch.datasets.BlobDataset<\/code><\/a> class is a map-style dataset, enabling random access to data samples. The azstoragetorch.datasets.IterableBlobDataset<\/code><\/a> class is an iterable dataset, which should be used when working on large datasets that may not fit in memory.<\/p>\n
Both classes support two methods: from_container_url()<\/code> and from_blob_urls()<\/code>. The from_container_url()<\/code> method instantiates a dataset by listing blobs in a container, and the from_blob_urls()<\/code> method takes a list of blob URLs when creating a dataset.<\/p>\n
These Dataset<\/code> integrations fit naturally into a PyTorch workflow. Let’s dive into an image example using the caltech101<\/a> dataset and the resnet18<\/code> model.<\/p>\n
The caltech101 dataset contains about 9,000 images across various categories and the resnet18<\/code> model is a residual neural network introduced in the paper Deep Residual Learning for Image Recognition<\/a>.<\/p>\n
Prerequisites and Setup<\/h4>\n
First, install prerequisite packages azstoragetorch<\/code>, pillow<\/code>, and torchvision<\/code>.<\/p>\n
pip install azstoragetorch pillow torchvision<\/code><\/pre>\n
Once you have the packages installed, ensure you have the caltech101 dataset in your Azure Storage Account. To ease this setup process, clone the GitHub repository<\/a> and run the provided setup file<\/a>.<\/p>\n
Lastly, it’s helpful to understand the importance of transforming data for PyTorch operations. The default output of dataset samples in the azstoragetorch<\/code> package are dictionaries representing a blob in the dataset. It’s often necessary to transform this data into the shape needed for your PyTorch workflows.<\/p>\n
To override the default output, we can provide a transform<\/code> callable in the from_blob_urls<\/code> or from_container_url<\/code> methods that accept an argument of type azstoragetorch.datasets.Blob<\/code><\/a>. The transform callable in the following code sample is based on the PyTorch documentation<\/a>. To learn more about how to use the transform<\/code> callable in the azstoragetorch<\/code> library, visit the documentation<\/a>.<\/p>\n
from torch.utils.data import DataLoader\r\nimport torchvision.transforms\r\nfrom PIL import Image\r\n\r\nfrom azstoragetorch.datasets import BlobDataset\r\n\r\n# Method to transform blob data to a torch.Tensor\r\n# To learn more about why these particular transforms were chosen,\r\n# Visit the documentation site: https:\/\/pytorch.org\/hub\/pytorch_vision_resnet\/\r\ndef blob_to_category_and_tensor(blob):\r\n    with blob.reader() as f:\r\n        img = Image.open(f).convert(\"RGB\")\r\n        img_transform = torchvision.transforms.Compose([\r\n            torchvision.transforms.Resize(256),\r\n            torchvision.transforms.CenterCrop(224),\r\n            torchvision.transforms.ToTensor(),\r\n            torchvision.transforms.Normalize(\r\n                mean=[0.485, 0.456, 0.406],\r\n                std=[0.229, 0.224, 0.225]\r\n            ),\r\n        ])\r\n        img_tensor = img_transform(img)\r\n\r\n    # Get second to last component of blob name which will be the image category\r\n    # Example: blob.blob_name -> datasets\/caltech101\/dalmatian\/image_0001.jpg\r\n    # category -> dalmatian\r\n    category = blob.blob_name.split(\"\/\")[-2]\r\n    return category, img_tensor\r\n\r\n# The URL to your container\r\nCONTAINER_URL = \"https:\/\/<my-storage-account-name>.blob.core.windows.net\/<my-container-name>\"\r\n\r\n# Initialize dataset with the azstoragetorch library\r\ndataset = BlobDataset.from_container_url(\r\n    CONTAINER_URL,\r\n    prefix=\"datasets\/caltech101\",\r\n    transform=blob_to_category_and_tensor,\r\n)\r\n\r\n# Set up data loader\r\nloader = DataLoader(dataset, batch_size=32)\r\n\r\nfor categories, img_tensors in loader:\r\n    print(categories, img_tensors.size())<\/code><\/pre>\n
Conclusion<\/h2>\n
The Azure Storage Connector for PyTorch is designed around the principle that cloud storage integration for your ML workflows shouldn’t require learning new paradigms. Key features include:<\/p>\n