使用 Distributed Data Parallel 在多 GPU 上通过 Opacus 进行训练¶

%%writefile opacus_ddp_demo.py
import os
import torch.distributed as dist
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

%%writefile -a opacus_ddp_demo.py

import torch.nn as nn
import torch.nn.functional as F

class SampleConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 8, 2, padding=3)
        self.conv2 = nn.Conv2d(16, 32, 4, 2)
        self.fc1 = nn.Linear(32 * 4 * 4, 32)
        self.fc2 = nn.Linear(32, 10)

    def forward(self, x):
        # x of shape [B, 1, 28, 28]
        x = F.relu(self.conv1(x))  # -> [B, 16, 14, 14]
        x = F.max_pool2d(x, 2, 1)  # -> [B, 16, 13, 13]
        x = F.relu(self.conv2(x))  # -> [B, 32, 5, 5]
        x = F.max_pool2d(x, 2, 1)  # -> [B, 32, 4, 4]
        x = x.view(-1, 32 * 4 * 4)  # -> [B, 512]
        x = F.relu(self.fc1(x))  # -> [B, 32]
        x = self.fc2(x)  # -> [B, 10]
        return x

%%writefile -a opacus_ddp_demo.py

from torchvision import datasets, transforms

# Precomputed characteristics of the MNIST dataset
MNIST_MEAN = 0.1307
MNIST_STD = 0.3081

DATA_ROOT = "./mnist"

mnist_train_ds = datasets.MNIST(
    DATA_ROOT,
    train=True,
    download=True,
    transform=transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((MNIST_MEAN,), (MNIST_STD,)),
        ]
    ),
)

mnist_test_ds = datasets.MNIST(
    DATA_ROOT,
    train=False,
    download=True,
    transform=transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((MNIST_MEAN,), (MNIST_STD,)),
        ]
    ),
)

%%writefile -a opacus_ddp_demo.py

import torch
import torch.optim as optim
from torch.utils.data import DataLoader, DistributedSampler

from opacus.distributed import DifferentiallyPrivateDistributedDataParallel as DPDDP
from torch.nn.parallel import DistributedDataParallel as DDP

from opacus import PrivacyEngine

LR = 0.1
BATCH_SIZE = 200
N_GPUS = torch.cuda.device_count()

def init_training(rank):
    model = SampleConvNet()
    #model = DDP(model) -- non-private
    model = DPDDP(model)

    optimizer = optim.SGD(model.parameters(), lr=LR, momentum=0)
    data_loader = DataLoader(
        mnist_train_ds,
        #batch_size=BATCH_SIZE // N_GPUS, -- non-private
        batch_size=BATCH_SIZE,
        #sampler=DistributedSampler(mnist_train_ds) -- non-private
    )

    if rank == 0:
        logger.info(
            f"(rank {rank}) Initialized model ({type(model).__name__}), "
            f"optimizer ({type(optimizer).__name__}), "
            f"data loader ({type(data_loader).__name__}, len={len(data_loader)})"
        )

    privacy_engine = PrivacyEngine()

    # PrivacyEngine looks at the model's class and enables
    # distributed processing if it's wrapped with DPDDP
    model, optimizer, data_loader = privacy_engine.make_private(
        module=model,
        optimizer=optimizer,
        data_loader=data_loader,
        noise_multiplier=1.,
        max_grad_norm=1.,
    )

    if rank == 0:
        logger.info(
            f"(rank {rank}) After privatization: model ({type(model).__name__}), "
            f"optimizer ({type(optimizer).__name__}), "
            f"data loader ({type(data_loader).__name__}, len={len(data_loader)})"
        )

    logger.info(f"(rank {rank}) Average batch size per GPU: {int(optimizer.expected_batch_size)}")

    return model, optimizer, data_loader, privacy_engine

%%writefile -a opacus_ddp_demo.py

import numpy as np

def test(model, device):
    test_loader = DataLoader(
        mnist_test_ds,
        batch_size=BATCH_SIZE,
    )

    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(
                dim=1, keepdim=True
            )
            correct += pred.eq(target.view_as(pred)).sum().item()

    model.train()
    return correct / len(mnist_test_ds)

def launch(rank, world_size, epochs):
    setup(rank, world_size)
    criterion = nn.CrossEntropyLoss()

    model, optimizer, data_loader, privacy_engine = init_training(rank)
    model.to(rank)
    model.train()


    for e in range(epochs):
        losses = []
        correct = 0
        total = 0

        for data, target in data_loader:
            data, target = data.to(rank), target.to(rank)
            optimizer.zero_grad()
            output = model(data)

            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
            total += len(data)

            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            losses.append(loss.item())

        test_accuracy = test(model, rank)
        train_accuracy = correct / total
        epsilon = privacy_engine.get_epsilon(delta=1e-5)

        if rank == 0:
            print(
                f"Epoch: {e} \t"
                f"Train Loss: {np.mean(losses):.4f} | "
                f"Train Accuracy: {train_accuracy:.2f} | "
                f"Test Accuracy: {test_accuracy:.2f} |"
                f"(ε = {epsilon:.2f})"
            )

    cleanup()

%%writefile -a opacus_ddp_demo.py

import torch.multiprocessing as mp

EPOCHS = 10
world_size = torch.cuda.device_count()

if __name__ == '__main__':
    mp.spawn(
        launch,
        args=(world_size,EPOCHS,),
        nprocs=world_size,
        join=True
    )

!python -W ignore opacus_ddp_demo.py

%%writefile opacus_distributed_data_loader_demo.py

from opacus_ddp_demo import setup, cleanup, mnist_train_ds
import logging
from torch.utils.data import DataLoader, DistributedSampler
from opacus.data_loader import DPDataLoader

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

BATCH_SIZE = 64

def init_data(rank, world_size):
    setup(rank, world_size)

    non_distributed_dl = DataLoader(
        mnist_train_ds,
        batch_size=BATCH_SIZE
    )

    distributed_non_private_dl = DataLoader(
        mnist_train_ds,
        batch_size=BATCH_SIZE // world_size,
        sampler=DistributedSampler(mnist_train_ds),
    )

    private_dl = DPDataLoader.from_data_loader(non_distributed_dl, distributed=True)

    if rank == 0:
        logger.info(
            f"(rank {rank}) Non-distributed non-private data loader. "
            f"#batches: {len(non_distributed_dl)}, "
            f"#data points: {len(non_distributed_dl.sampler)}, "
            f"batch_size: {non_distributed_dl.batch_size}"
        )

        logger.info(
            f"(rank {rank}) Distributed, non-private data loader. "
            f"#batches: {len(distributed_non_private_dl)}, "
            f"#data points: {len(distributed_non_private_dl.sampler)}, "
            f"batch_size: {distributed_non_private_dl.batch_size}"
        )

        logger.info(
            f"(rank {rank}) Distributed, private data loader. "
            f"#batches: {len(private_dl)}, "
            f"#data points: {private_dl.batch_sampler.num_samples}, "
            f"sample_rate: {private_dl.sample_rate:4f}, "
            f"avg batch_size (=sample_rate*num_data_points): {int(private_dl.sample_rate*private_dl.batch_sampler.num_samples)}"
        )

%%writefile -a opacus_distributed_data_loader_demo.py

import torch
import torch.multiprocessing as mp

world_size = torch.cuda.device_count()

if __name__ == '__main__':
    mp.spawn(
        init_data,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )

!python -W ignore opacus_distributed_data_loader_demo.py

%%writefile opacus_sync_demo.py

import sys
sys.path.append('/data/home/shilov/opacus')

from opacus_ddp_demo import setup, cleanup, mnist_train_ds, SampleConvNet
import logging
from torch.utils.data import DataLoader
import torch.optim as optim
from opacus.data_loader import DPDataLoader
from opacus import GradSampleModule
from opacus.distributed import DifferentiallyPrivateDistributedDataParallel as DPDDP
from opacus.optimizers import DistributedDPOptimizer
from torch.nn.parallel import DistributedDataParallel as DDP

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

BATCH_SIZE = 64
LR = 64

def init_training(rank, world_size):
    model = SampleConvNet()
    optimizer = optim.SGD(model.parameters(), lr=LR, momentum=0)

    model = GradSampleModule(model)
    model = DPDDP(model)

    optimizer = DistributedDPOptimizer(
        optimizer=optimizer,
        noise_multiplier=0.,
        max_grad_norm=100.,
        expected_batch_size=BATCH_SIZE//world_size,
    )

    data_loader = DPDataLoader.from_data_loader(
        data_loader=DataLoader(
            mnist_train_ds,
            batch_size=BATCH_SIZE,
        ),
        distributed=True,
    )


    return model, optimizer, data_loader

%%writefile -a opacus_sync_demo.py

import torch.nn as nn
import numpy as np

def launch(rank, world_size):
    setup(rank, world_size)
    criterion = nn.CrossEntropyLoss()

    model, optimizer, data_loader = init_training(rank, world_size)
    model.to(rank)
    model.train()

    for data, target in data_loader:
        data = data
        target = torch.tensor(target)

        data, target = data.to(rank), target.to(rank)
        optimizer.zero_grad()

        output = model(data)
        loss = criterion(output, target)
        loss.backward()

        flat_grad = torch.cat([p.grad_sample.sum(dim=0).view(-1) for p in model.parameters()]).cpu().numpy() / optimizer.expected_batch_size
        logger.info(
            f"(rank={rank}) Gradient norm before optimizer.step(): {np.linalg.norm(flat_grad):.4f}"
        )
        logger.info(
            f"(rank={rank}) Gradient sample before optimizer.step(): {flat_grad[:3]}"
        )

        optimizer.step()

        flat_grad = torch.cat([p.grad.view(-1) for p in model.parameters()]).cpu().numpy()
        logger.info(
            f"(rank={rank}) Gradient norm after optimizer.step(): {np.linalg.norm(flat_grad):.4f}"
        )
        logger.info(
            f"(rank={rank}) Gradient sample after optimizer.step(): {flat_grad[:3]}"
        )

        break

    cleanup()

%%writefile -a opacus_sync_demo.py

import torch.multiprocessing as mp
import torch

world_size = torch.cuda.device_count()

if __name__ == '__main__':
    mp.spawn(
        launch,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )

!python -W ignore opacus_sync_demo.py

%%bash
rm opacus_ddp_demo.py
rm opacus_distributed_data_loader_demo.py
rm opacus_sync_demo.py