🐏

PyTorchで使われる擬似乱数生成器

2025/02/07に公開

PyTorch

乱数

tech

機械学習の結果を比較検証する際、学習時のサンプリングの偏りや再現性は重要な要素の一つです。
その中でも、Pythonの機械学習用ライブラリであるPyTorchで使用されている擬似乱数生成器についてまとめます。

公式サイトによると、デフォルトではCPU用にMersenne Twister、CUDA用にNvidiaのcuRAND Philoxが使用されており、オプションで暗号論的にセキュアなTORCHCSPRINGも使用できるようです。

Mersenne Twisterの特徴

非常に長い $2^{19937}-1$ の周期をもつ
$623$ 次元までは均等分布する
- 例えば、多変量の確率分布をサンプリングしても $623$ 次元までは偏りが起こらないことが保証されている
内部に $623$ 個の $32bit$ 長の状態ベクトルを持つため、ワーキングメモリが大きい

Philoxの特徴

$2^{130}$ の周期をもつ
カウンタベースの擬似乱数生成器
- Mersenne Twisterのように大きな内部状態を持たないため、ワーキングメモリが小さい
- 各スレッドにカウンタを割り当てることで、衝突や相関のない乱数列を高速に生成できるので、大規模並列化が容易

触ってみる

import torch

def get_tensor_size_in_bytes(tensor: torch.Tensor) -> int:
    """
    PyTorch Tensor の要素数 × 1要素あたりバイト数 から
    実際に保持しているデータサイズを計算して返す。
    """
    return tensor.numel() * tensor.element_size()

# CPU 側 RNG state (Mersenne Twister)
cpu_state = torch.random.get_rng_state()
cpu_bytes = get_tensor_size_in_bytes(cpu_state)

print("=== CPU RNG State (Mersenne Twister) ===")
print(f"Shape   : {cpu_state.shape}")
print(f"Size    : {cpu_bytes} bytes")

# GPU 側 RNG state (Philox) 
gpu_state = torch.cuda.get_rng_state()
gpu_bytes = get_tensor_size_in_bytes(gpu_state)

print("\n=== GPU RNG State (Philox) ===")
print(f"Shape   : {gpu_state.shape}")
print(f"Size    : {gpu_bytes} bytes")

=== CPU RNG State (Mersenne Twister) ===
Shape   : torch.Size([5056])
Size    : 5056 bytes

=== GPU RNG State (Philox) ===
Shape   : torch.Size([16])
Size    : 16 bytes

手元の環境でCPU、GPU上で乱数生成器を呼び出し、その内部状態をprintした結果です。
GPU上の乱数生成器はかなり小さくなっており、独立した乱数列を大量に用意してもメモリ消費を抑えられることがわかります。

import torch
import time

def measure_time_and_stats(size, device='cpu'):
    """
    指定デバイスで0/ランダムテンソルを生成し、実行時間を計測する。
    """

    # 乱数生成の時間計測
    times = []
    for _ in range(10):
        start = time.time()
        torch.randn(size, device=device) or torch.zeros(size, device=device)
        torch.cuda.synchronize()
        end = time.time()
        times.append(end - start)

    return sum(times) / len(times)

# 比較するサイズ
tensor_size = (20000, 20000)

# CPU (Mersenne Twister)
cpu_time = measure_time_and_stats(
    size=tensor_size, device='cpu'
)

# GPU (Philox)
gpu_time = measure_time_and_stats(
    size=tensor_size, device='cuda'
)

# 結果表示
print("=== CPU (Mersenne Twister) ===")
print(f"Time : {cpu_time:.6f} sec")

print("\n=== GPU (Philox) ===")
print(f"Time : {gpu_time:.6f} sec")

# 0テンソルの場合
=== CPU (Mersenne Twister) ===
Time : 0.252750 sec

=== GPU (Philox) ===
Time : 0.008937 sec

# ランダムテンソルの場合
=== CPU (Mersenne Twister) ===
Time : 2.065133 sec

=== GPU (Philox) ===
Time : 0.031702 sec

CPU、GPU上で指定したサイズの0テンソルとランダムテンソルを生成し、実行時間を計測した結果です。

ランダムテンソルを生成する場合、CPU上では実行時間が8倍になっているのに対し、GPU上では実行時間が3~4倍程度になっています。このことから、PyTorchではGPU上で計算を行うことで乱数生成の計算コストを抑えられることが確認できました。

PyTorchで使われる擬似乱数生成器

Mersenne Twisterの特徴

Philoxの特徴

触ってみる

参考

Discussion