Gradient Accumulation（勾配累積）

Gradient Accumulationを使えば、サイズの大きいモデルでもバッチサイズを確保しながら学習できます。
考え方としては、設定できるバッチサイズの最大が2だとしても、acc_step = 8にすると実質2×8=16のバッチサイズにできるという感じです。

Gradient Accumulationの使い方

データセット

import torchvision
train = torchvision.datasets.CIFAR10(
    root = "data",
    train = True,
    download = True,
    transform = torchvision.transforms.ToTensor()
)

valid = torchvision.datasets.CIFAR10(
    root = "data",
    train = False,
    download = True,
    transform = torchvision.transforms.ToTensor()
)

CIFAR10で画像分類をします。
trainが学習用で、validが検証用です。

データセット, モデル定義

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
# !pip install -Uqqq timm
import timm
from tqdm import tqdm
import albumentations as A
import matplotlib.pyplot as plt
from timm.utils import AverageMeter

class CustomData(Dataset):
    def __init__(self, data, transform):
        self.data = data
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img, lbl = self.data[idx]
        img = img.permute(1, 2, 0).numpy()
        if self.transform is not None:
            img = self.transform(image = img)["image"]
        img = torch.from_numpy(img).permute(2, 0, 1)
        return img, lbl

先程取得したデータを取り出すようにしています。
idx=0が最初のデータで、train[idx] = 画像, ラベルみたいに取り出せます。

model = timm.create_model("resnet34d", pretrained = True, num_classes = 10).cuda()
optimizer = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

モデルはresnet34dにしておきます。後ほど大きなモデルにしてみましょう。
ちなみに使っているGPUはGoogle ColabのT4です。

Gradient Accumulationなしで学習

batch_size = 32

train_ds = CustomData(train, A.Compose([A.Flip(), A.RandomRotate90()]))
train_dl = DataLoader(train_ds, batch_size = batch_size, shuffle = True, pin_memory = True)
valid_ds = CustomData(valid, None)
valid_dl = DataLoader(valid_ds, batch_size = batch_size, shuffle = False, pin_memory = True)

results = {"train": [], "valid": []}
for epoch in range(10):
    model.train()
    train_loss = AverageMeter()
    for x, y in tqdm(train_dl):
        x = x.cuda()
        y = y.cuda()
        logits = model(x)
        loss = criterion(logits, y)
        train_loss.update(val = loss.item(), n = len(y))
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    valid_loss = AverageMeter()
    for x, y in tqdm(valid_dl):
        x = x.cuda()
        y = y.cuda()
        with torch.no_grad():
            logits = model(x)
        valid_loss.update(val = loss.item(), n = len(y))
    print(train_loss.avg)
    print(valid_loss.avg)
    results["train"].append(train_loss.avg)
    results["valid"].append(valid_loss.avg)
plt.plot(results["train"])
plt.plot(results["valid"])
plt.show()

バッチサイズは32です。アウトプットは以下のとおり。

正解率は80%くらいになります。

Gradient Accumulationありで学習

acc_step = 4
batch_size = 8

train_ds = CustomData(train, A.Compose([A.Flip(), A.RandomRotate90()]))
train_dl = DataLoader(train_ds, batch_size = batch_size, shuffle = True, pin_memory = True)
valid_ds = CustomData(valid, None)
valid_dl = DataLoader(valid_ds, batch_size = batch_size, shuffle = False, pin_memory = True)

model = timm.create_model("resnet34d", pretrained = True, num_classes = 10).cuda()
optimizer = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

results = {"train": [], "valid": []}
for epoch in range(10):
    model.train()
    train_loss = AverageMeter()
    for i, (x, y) in enumerate(tqdm(train_dl)):
        x = x.cuda()
        y = y.cuda()
        logits = model(x)
        loss = criterion(logits, y)
        train_loss.update(val = loss.item(), n = len(y))
        loss /= acc_step
        loss.backward()
        if (i + 1) % acc_step == 0:
            optimizer.step()
            optimizer.zero_grad()
    valid_loss = AverageMeter()
    for x, y in tqdm(valid_dl):
        x = x.cuda()
        y = y.cuda()
        with torch.no_grad():
            logits = model(x)
        valid_loss.update(val = loss.item(), n = len(y))
    print(train_loss.avg)
    print(valid_loss.avg)
    results["train"].append(train_loss.avg)
    results["valid"].append(valid_loss.avg)
plt.plot(results["train"])
plt.plot(results["valid"])
plt.show()

acc_stepが勾配を溜め込むステップ数です。
batch_size x acc_step = 4 x 8 = 32が実質のバッチサイズになります。

変更点は、まずlossをacc_stepで割ること。あとは1エポック内のステップ数をiとして(i + 1)%acc_step=0のときに重みを更新することです。

結果は以下のとおり。