n4okins / 今まで見つけたPyTorchとかの知見集

Created Sun, 31 Mar 2024 15:11:33 +0900 Modified Sun, 31 Mar 2024 15:11:33 +0900

多くなってきたため後々分割するかも

Error Tips

import torchでエラー?

    ...
    if any(filename.endswith(s) for s in all_bytecode_suffixes):
AttributeError: 'PosixPath' object has no attribute 'endswith'

まずこんなエラーに遭遇する人は稀だろうがimport torchでエラーが出ることがある。

自分の場合は、torchをimportする前に__file__をstr型からPath型に変更していたために発生していた。

from pathlib import Path
__file__ = Path(__file__)
import torch

print(torch)

上記ではエラーが出るので、下記のように__file__をstr型のままにしておく。

from pathlib import Path
import torch
__file__ = Path(__file__)

print(torch)

なお、そもそも__file__は予約語なので変更すべきではない。

学習Tips

ConvTranspose2DとUpsample + Conv2Dの違い

どちらもアップサンプリングだが、違いがある。

PixelShuffleという選択肢

下記によるとDeconvolution系の処理よりもPixelShuffleの方が学習速度が速いらしい。(注: 2017年時点)

BatchNormalizationはどの層に入れるべきか

上の記事でも触れられているため一応メモ。 例えば活性化関数ReLUのあとに正規化すると負の値が0になってしまうため、意味がない?

全結合のニューラルネットワークの場合、Affineの後、活性化関数の前にいれると良いらしい。

学習中のMatplotlibでメモリリーク

from multiprocessing import Pool
import matplotlib.pyplot as plt

def plot(args):
    x, y = args
    plt.plot(x, y)
    plt.show()

N = 4
for i in range(100):
    with Pool(N) as p:
        p.map(plot, [(range(10), range(10))])

Poolを使って書くと良い。

高速化Tips

WSL2での高速化

データのあるディレクトリをWSL2のディレクトリにマウントすると高速化できる。

WSL2から/mnt/c/などのWSLにマウントされているWindowsフォルダ(NFTS)へのアクセスは、ext4でマウントされたフォルダにアクセスするよりも遥かに遅い。

可能であればext4専用のSSDを使用するのが良い。

torch.compileの活用

Python3.9以下 かつ PyTorch2.0以上 で使える機能。

model = Model()
compiled_model = torch.compile(model) 
# compiled_model = torch.compile(model, mode="reduce-overhead")
# compiled_model = torch.compile(model, mode="max-autotune")
# compiled_model = torch.compile(model, mode="max-autotune-no-cudagraphs")

とすることで、モデルの高速化が可能。 ただしbf16での学習時にエラーが出ることがある?

torch.ampの活用

多くの場合はfloat32で学習/推論するのは冗長であるらしい。

torch.ampscalerを使うことで、float16やbfloat16で学習/推論することが可能。 必要な時間計算量・空間計算量を減らすことができる。

class Model(nn.Module):
    ...

model = Model()

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = torch.cuda.amp.grad_scaler.GradScaler()

for images, labels in dataloader:
    with torch.amp.autocast_mode.autocast(device_type="cuda", dtype=torch.float16):
    # with torch.amp.autocast_mode.autocast(device_type="cuda", dtype=torch.bfloat16):
        out = model(images)
        loss = criterion(out, labels)

    optimizer.zero_grad()
    scaler.scale(loss).backward()  # type: ignore
    scaler.unscale_(optimizer)
    nn.utils.clip_grad.clip_grad_norm_(model.parameters(), 1)
    scaler.step(optimizer)
    scaler.update()

nvidia-Daliの活用

画像データを読み込む際、PyTorchのDataLoaderはめちゃくちゃ遅い。 nvidia-Daliを使うことで、高速化が可能。

import ctypes
from functools import partial
from logging import Logger, getLogger
from typing import Any, Callable

import torch
from torchvision.datasets import ImageFolder
from nvidia.dali import fn as dali_fn
from nvidia.dali.backend import TensorCPU, TensorGPU
from nvidia.dali.pipeline import Pipeline as DALIPipeline
from nvidia.dali.types import DALIDataType, DALIImageType

DALIDataType2TorchDataType = {
    DALIDataType.FLOAT: torch.float32,
    DALIDataType.FLOAT64: torch.float64,
    DALIDataType.FLOAT16: torch.float16,
    DALIDataType.UINT8: torch.uint8,
    DALIDataType.INT8: torch.int8,
    DALIDataType.BOOL: torch.bool,
    DALIDataType.INT16: torch.int16,
    DALIDataType.INT32: torch.int32,
    DALIDataType.INT64: torch.int64,
}


def dali_to_torch_tensor(
    dali_tensor: TensorCPU | TensorGPU, device_id=None
) -> torch.Tensor:
    stream = None
    device = torch.device("cpu")
    if isinstance(dali_tensor, TensorGPU):
        device = torch.device("cuda", index=device_id or 0)
        stream = torch.cuda.current_stream(device=device)
    torch_tensor = torch.empty(
        dali_tensor.shape(),
        dtype=DALIDataType2TorchDataType[dali_tensor.dtype],
        device=device,
    )
    pointer = ctypes.c_void_p(torch_tensor.data_ptr())
    if device.type == "cuda":
        stream = stream if stream is None else ctypes.c_void_p(stream.cuda_stream)
        dali_tensor.copy_to_external(pointer, stream, non_blocking=True)
    else:
        dali_tensor.copy_to_external(pointer)

    return torch_tensor


class DALIFromImageFolder(DALIPipeline):
    def default_collate_fn(self, batch):
        images, labels = zip(*batch)
        images = torch.stack(images)
        labels = torch.stack(labels)
        return images, labels

    def __init__(
        self,
        dataset: ImageFolder,
        transform: Callable = lambda x: x,
        target_transform: Callable = lambda x: x,
        batch_size: int = 1,
        num_workers: int = 12,
        shuffle: bool = True,
        device_id: int = 0,
        seed: int | None = None,
        logger: Logger = getLogger(__name__),
        collate_fn: Callable[[Any], Any] | None = None,
        *args,
        **kwargs,
    ):
        super(DALIFromImageFolder, self).__init__(batch_size, num_workers, device_id)
        self.dataset = dataset

        self.transform = transform
        self.target_transform = target_transform

        self.file_pathes, self.labels = zip(*self.dataset.samples)
        self.collate_fn = collate_fn or self.default_collate_fn

        self.input = dali_fn.readers.file(
            files=self.file_pathes,
            labels=self.labels,
            name="Reader",
            random_shuffle=shuffle,
            seed=seed,
        )
        self.decode = partial(
            dali_fn.decoders.image, device="cpu", output_type=DALIImageType.RGB
        )
        self.build()

        self._index = 0

    def define_graph(self):
        images, labels = self.input
        images = self.decode(images)
        return (images, labels)

    @property
    def index(self):
        return self._index

    def __len__(self):
        return len(self.dataset) // self.max_batch_size + 1

    def __iter__(self):
        for i in range(len(self)):
            images, labels = self.run()
            images = list(map(dali_to_torch_tensor, images))
            labels = list(map(dali_to_torch_tensor, labels))
            if self.transform:
                images = list(map(self.transform, images))
            if self.target_transform:
                labels = list(map(self.target_transform, labels))
            images, labels = self.collate_fn(list(zip(images, labels)))
            yield images, labels

改善の余地が多そう&おそらくは公式ドキュメント通りにDALIGenericIteratorなどを使うべきではあるが、こんな感じだとImageFolderの感覚で使える(と思う)。

手元の実験としてImageNetのtrainデータを224x224にリサイズし、batch_size=256で100回取り出す処理をしたところ、PyTorchのDataLoaderよりもDaliのPipelineが約4倍速かった。