Distribuce trénování PyTorch pomocí torchDistributoru

6 min

PyTorch, běžně s jinými architekturami hlubokého učení, jako je TensorFlow, je navržená tak, aby škálovat napříč několika procesory (PROCESORy nebo GPU) na jednom počítači. Ve většině případů tento přístup ke vertikálnímu navýšení kapacity pomocí počítačů s více nebo rychlejšími procesory poskytuje odpovídající výkon trénování.

Pokud ale potřebujete pracovat se složitými neurálními sítěmi nebo velkými objemy trénovacích dat, můžete využít základní schopnosti Apache Sparku škálovat úlohy zpracování napříč několika pracovními uzly.

Azure Databricks používá clustery Spark, které můžou zahrnovat více pracovních uzlů. Pokud chcete zajistit optimální využití těchto clusterů, můžete použít TorchDistributor, opensourcovou knihovnu, která umožňuje distribuovat trénovací úlohy PyTorch mezi uzly v clusteru. TorchDistributor je k dispozici v Databricks Runtime ML 13.0 a vyšší.

Pokud jste model už natrénovali pomocí PyTorch, můžete převést trénování jednoho procesu na distribuované trénování pomocí torchDistributoru podle:

Přizpůsobte stávající kód: Upravte trénovací kód s jedním uzlem tak, aby byl kompatibilní s distribuovaným trénováním. Ujistěte se, že je logika trénování zapouzdřená v rámci jedné funkce.
Přesun importů v rámci trénovací funkce: Umístěte nezbytné importy, například import torch, uvnitř trénovací funkce, aby nedocházelo k běžným chybám při výběru.
Příprava trénovací funkce: Do trénovací funkce zahrňte model, optimalizátor, funkci ztráty a trénovací smyčku. Ujistěte se, že se model a data přesunou do příslušného zařízení (CPU nebo GPU).
Vytvoření instance a spuštění TorchDistributor: Vytvořte instanci TorchDistributor s požadovanými parametry a volání .run(*args) pro spuštění distribuovaného trénování.

Přizpůsobení existujícího kódu

Nejprve je potřeba upravit trénovací kód s jedním uzlem tak, aby byl kompatibilní s distribuovaným trénováním. Při úpravě kódu je potřeba zajistit zapouzdření trénovací logiky v rámci jedné funkce. Tuto funkci používá TorchDistributor k distribuci trénování mezi více uzlů.

import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)
    
    def forward(self, x):
        return self.fc(x)

Teď můžete připravit datovou sadu, která je ve formátu kompatibilním s PyTorch pomocí torch.utils.data.DataLoader.

# Sample data
inputs = torch.randn(100, 10)
targets = torch.randn(100, 1)

# Create dataset and dataloader
from torch.utils.data import DataLoader, TensorDataset
dataset = TensorDataset(inputs, targets)
dataloader = DataLoader(dataset, batch_size=10)

Přesun importů v rámci trénovací funkce

Abyste se vyhnuli běžným chybám při výběru, umístěte do trénovací funkce nezbytné importy, například import torch. Umístění všech importů do trénovací funkce zajišťuje, že všechny požadované moduly budou k dispozici, když je funkce distribuována napříč více uzly.

Příprava trénovací funkce

Do trénovací funkce zahrňte svůj model, optimalizátor, funkci ztráty a trénovací smyčku. Ujistěte se, že se model a data přesunou do příslušného zařízení (CPU nebo GPU).

def train_model(dataloader):
    import torch
    import torch.nn as nn
    from torch.optim import SGD

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SimpleModel().to(device)
    optimizer = SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()
    
    for epoch in range(10):
        for batch in dataloader:
            inputs, targets = batch
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_fn(outputs, targets)
            loss.backward()
            optimizer.step()

Vytvoření instance a spuštění torchDistributoru

Vytvořte instanci TorchDistributor s požadovanými parametry a voláním .run(*args) pro spuštění distribuovaného trénování. Spuštění TorchDistributor distribuuje trénovací úlohy napříč několika uzly.

from pyspark.ml.torch.distributor import TorchDistributor

# Distribute the training
distributor = TorchDistributor(num_workers=4)
distributor.run(train_model, dataloader)

Monitorování a vyhodnocení trénovací úlohy

Integrované nástroje můžete použít ke sledování výkonu clusteru, včetně využití procesoru nebo GPU a využití paměti. Po dokončení trénování můžete model vyhodnotit na ověřovací nebo testovací datové sadě pomocí technik vyhodnocení PyTorch k vyhodnocení výkonu modelu.

# Evaluate the model (after distributed training is complete)
model.eval()
with torch.no_grad():
    for inputs, targets in dataloader:
        outputs = model(inputs)
        # Perform evaluation logic

Tip