PyTorch

参考文献

1. torch.utils.data — PyTorch 2.1 documentation. (2023年10月8日参照).

CPU 版と GPU 版の違い

CPU 版と GPU 版で以下の実行結果が異なるのがわかる。これは CPU 版と GPU 版ではより低レイヤーでソートアルゴリズムが異なることに起因していると自分は推測しているが不明である。

import torch
x = torch.tensor([0.3, 0.2, 0.1, 0.4, 0.5])
print(x.topk(3, sorted=False))

# GPU
torch.return_types.topk(
values=tensor([0.5000, 0.4000, 0.3000]),
indices=tensor([4, 3, 0]))

# CPU
torch.return_types.topk(
values=tensor([0.4000, 0.5000, 0.3000]),
indices=tensor([3, 4, 0]))

関数・クラス

torch.utils.data.DataLoader

DataLoader

は以下のように使用できる。明示的に

suffle=True

としない限りは単に前から順にバッチサイズずつバッチを切り出す。

suffle=True

とした場合はエポックごとにリシャッフルされる。

from torch.utils.data import Dataset, DataLoader
import pandas as pd

class MyDataset(Dataset):
    def __init__(self, df):
        self.data = df.loc[:, ['x0', 'x1']]
        self.y = df['y']
        self.n_sample = len(df)
    def __len__(self):
        return self.n_sample
    def __getitem__(self, idx):
        return self.data.loc[idx, :].values, self.y[idx]

df = pd.DataFrame({
    'x0': [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
    'x1': [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
    'y' : [20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
})
dataset = MyDataset(df)

batch_size = 4
dataloader = DataLoader(dataset=dataset, batch_size=batch_size)
for i_epoch in range(2):
    print(f'===== epoch {i_epoch} =====')
    for i_batch, (data, y) in enumerate(dataloader):
        print(f'----- batch {i_batch} -----')
        print(data)
        print(y)

===== epoch 0 =====
----- batch 0 -----
tensor([[ 0, 10],
        [ 1, 11],
        [ 2, 12],
        [ 3, 13]])
tensor([20, 21, 22, 23])
----- batch 1 -----
tensor([[ 4, 14],
        [ 5, 15],
        [ 6, 16],
        [ 7, 17]])
tensor([24, 25, 26, 27])
----- batch 2 -----
tensor([[ 8, 18],
        [ 9, 19]])
tensor([28, 29])
===== epoch 1 =====
----- batch 0 -----
tensor([[ 0, 10],
        [ 1, 11],
        [ 2, 12],
        [ 3, 13]])
tensor([20, 21, 22, 23])
----- batch 1 -----
tensor([[ 4, 14],
        [ 5, 15],
        [ 6, 16],
        [ 7, 17]])
tensor([24, 25, 26, 27])
----- batch 2 -----
tensor([[ 8, 18],
        [ 9, 19]])
tensor([28, 29])

バッチの切り出し方を制御するには、

batch_sampler

引数に各バッチのインデックスのリストを返すイテラブルなオブジェクトを指定すればよい。例えば以下のようにリストで渡してもよい。ただし、これだとすべてのエポックでバッチの括り方もバッチの順序も固定になる。

dataloader = DataLoader(
    dataset=dataset,
    batch_sampler=[[0, 1, 4, 8], [2, 3, 6, 7], [5, 9]],
)
for i_epoch in range(2):
    print(f'===== epoch {i_epoch} =====')
    for i_batch, (data, y) in enumerate(dataloader):
        print(f'----- batch {i_batch} -----')
        print(data)
        print(y)

===== epoch 0 =====
----- batch 0 -----
tensor([[ 0, 10],
        [ 1, 11],
        [ 4, 14],
        [ 8, 18]])
tensor([20, 21, 24, 28])
----- batch 1 -----
tensor([[ 2, 12],
        [ 3, 13],
        [ 6, 16],
        [ 7, 17]])
tensor([22, 23, 26, 27])
----- batch 2 -----
tensor([[ 5, 15],
        [ 9, 19]])
tensor([25, 29])
===== epoch 1 =====
----- batch 0 -----
tensor([[ 0, 10],
        [ 1, 11],
        [ 4, 14],
        [ 8, 18]])
tensor([20, 21, 24, 28])
----- batch 1 -----
tensor([[ 2, 12],
        [ 3, 13],
        [ 6, 16],
        [ 7, 17]])
tensor([22, 23, 26, 27])
----- batch 2 -----
tensor([[ 5, 15],
        [ 9, 19]])
tensor([25, 29])