DP-SGDにおけるバッチサイズの効果

はじめに

DP-SGDで教師データのプライバシーを保護しつつ、高いテスト精度を実現することを目指します。
学習は以前の記事をベースにしています。
Jupyter Notebookは下記にあります。

https://github.com/toshi-4886/privacy_preserving_ML/blob/main/PyTorch/3_DPSGD_batch_size.ipynb

概要

  • ResNet18によるCIFAR-10の学習に、Opacusを用いてDP-SGDを適用した際のテスト精度を改善する。

  • DP-SGDの学習におけるバッチサイズの影響を確認する。

参考資料

実装

1. ライブラリのインポート

必要なライブラリをインポートします。

import torch
import torchvision
from torch.utils.data import DataLoader
import sys
from tqdm import tqdm
import matplotlib.pyplot as plt
import numpy as np
from sklearn import metrics

!pip install 'opacus>=1.0'
import opacus
from opacus.validators import ModuleValidator
from opacus.utils.batch_memory_manager import BatchMemoryManager

import warnings
warnings.simplefilter("ignore")

2. 実行環境の確認

使用するライブラリのバージョンや、GPU環境を確認します。
Opacusは1.0.0で書き方が大きく変わっているので注意してください。
Google Colaboratoryで実行した際の例になります。

print('Python:', sys.version)
print('PyTorch:', torch.__version__)
print('Torchvision:', torchvision.__version__)
print('Opacus:', opacus.__version__)
!nvidia-smi
Python: 3.9.16 (main, Dec  7 2022, 01:11:51) 
[GCC 9.4.0]
PyTorch: 1.13.1+cu116
Torchvision: 0.14.1+cu116
Opacus: 1.4.0
Sat Mar 25 01:56:54 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   64C    P8    11W /  70W |      0MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

3. データセットの用意

データセットの取得

カラー画像データセットのCIFAR-10を、torchvisionを用いて取得します。
以前の結果に基づいて、教師データにdata augmentationは適用せずに、前処理は下記のみを適用します。

  • ToTensor:データをTensor型に変換

  • Normalize:各チャネルの平均が0、標準偏差が1となるように標準化

CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)
CIFAR10_STD = (0.2470, 0.2435, 0.2616)

training_data = torchvision.datasets.CIFAR10(
    root="data",
    train=True,
    download=True,
    transform=torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(), 
        torchvision.transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD)]),
)


test_data = torchvision.datasets.CIFAR10(
    root="data",
    train=False,
    download=True,
    transform=torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(), 
        torchvision.transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD)]),
)

4. 学習の関数化

今回は、異なるバッチサイズで複数回学習を実行するため、学習部分を関数として定義します。
各エポックのテスト精度と消費した$${\epsilon}$$を出力します。

def training(batch_size = 128):    
    # make dataloader
    train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
    test_dataloader = DataLoader(test_data, batch_size=1024, shuffle=False)

    # make model
    model = torchvision.models.resnet18()
    model = ModuleValidator.fix(model)

    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    model = model.to(device)

    # make optimizer
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

    # DP-SGD settig
    privacy_engine = opacus.PrivacyEngine()
    model, optimizer, train_dataloader = privacy_engine.make_private(
        module=model,
        optimizer=optimizer,
        data_loader=train_dataloader,
        noise_multiplier=1.0,
        max_grad_norm=1.0,
    )

    # training 
    test_acc = []
    epsilon = []

    pbar = tqdm(range(50), desc=f"[batch size: {batch_size}]")
    for epoch in pbar:
            
        with BatchMemoryManager(
            data_loader=train_dataloader, 
            max_physical_batch_size=64, 
            optimizer=optimizer
        ) as memory_safe_data_loader:
            model.train()
            for (X, y) in memory_safe_data_loader:
                X, y = X.to(device), y.to(device)

                # optimization step
                optimizer.zero_grad()
                pred = model(X)
                loss = criterion(pred, y)
                loss.backward()
                optimizer.step()

        # calculate epsilon
        epsilon.append(privacy_engine.get_epsilon(1e-5))
    
        # test
        model.eval()
        pred_list = []
        y_list = []
        with torch.no_grad():
            for X, y in test_dataloader:
                X, y = X.to(device), y.to(device)

                # predict
                pred = model(X)
                loss = criterion(pred, y)

                y_list.extend(y.to('cpu').numpy().tolist())
                pred_list.extend(pred.argmax(1).to('cpu').numpy().tolist())

        test_acc.append(metrics.accuracy_score(y_list, pred_list))
        pbar.set_postfix(epsilon=epsilon[-1],test_acc=test_acc[-1])
    return epsilon, test_acc  

5. 学習

バッチサイズを、$${2^9 (=512)}$$、$${2^{10} (=1024)}$$、$${2^{12} (=4096)}$$と変えて学習を行います。

result = {}
for batch_size in [512, 1024, 4096]:
    epsilon, test_acc = training(batch_size)
    result[batch_size] = (epsilon, test_acc)
[batch size: 512]: 100%|██████████| 50/50 [1:13:26<00:00, 88.12s/it, epsilon=4.26, test_acc=0.575]
[batch size: 1024]: 100%|██████████| 50/50 [1:12:13<00:00, 86.66s/it, epsilon=6.43, test_acc=0.601]
[batch size: 4096]: 100%|██████████| 50/50 [1:11:39<00:00, 85.98s/it, epsilon=14.2, test_acc=0.629]

6. 学習結果の表示

学習結果として、消費した$${\epsilon}$$とテストデータの精度を描画します。

for batch_size, res in result.items():
    epsilon, test_acc = res
    plt.plot(epsilon, test_acc, label=f'Batch size: {batch_size}')
plt.xlabel('Epsilon')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

おわりに

今回の結果

バッチサイズを大きくするとエポック数50で到達するテスト精度は向上しますが、消費する$${\epsilon}$$も増大することが分かりました。

  • 小さい$${\epsilon}$$で学習を行いたい場合
    バッチサイズを大きくすると、 小さい$${\epsilon}$$ではテスト精度が出ないので、バッチサイズを小さくする必要がありそうです。

  • ある程度の$${\epsilon}$$を許容できる場合
    下記2つの方法がありますが、どちらの方がよいかは、データセットやネットワークによっても変わりそうですし、どの程度学習に時間をかけられるかによっても変わってくると思います。

  1. 小さいバッチサイズでエポック数を多くしてテスト精度を少しずつ向上させる

  2. バッチサイズを大きくして、小さいエポック数で高いテスト精度に到達させる

$${\epsilon=2}$$で到達するテスト精度は、バッチサイズ128と512でほぼ同じとなっていますが、バッチサイズ512の方がエポック数を少なくできるため、今後は512を採用しようと思います。

次にやること

optimizerを変えることでDP-SGDのテスト精度がどう変わるか検証したいと思います。

この記事が気に入ったらサポートをしてみませんか?