📂機械学習

PyTorchでNumpy配列からカスタムデータセットを作成して使用する方法

説明

>>> import numpy as np
>>> import torch
>>> from torch.utils.data import TensorDataset, DataLoader

$32\times 32$のサイズの「モノクロ」写真100枚を積み重ねたnumpy配列$X$と、それに対するラベル$Y$が用意されているとしよう。以下のコードで読み込んだとする。

>>> X = np.load("X.npy")
>>> X.shape
(100, 32, 32)
>>> Y = np.load("Y.npy")
>>> Y.shape
(100)

学習用のデータセットを作るために、次のようにnumpy配列をテンソルに変換する。

>>> X = torch.from_numpy(X)
>>> Y = torch.from_numpy(Y)

TensorDatasetで学習データをまとめる。もちろん、3種類以上のデータをまとめてデータセットを作ることも可能だ。

>>> train_data = TensorDataset(X, Y)

DataLoaderを使用して、データの順序をランダムにシャッフルし、望むバッチサイズで作成できる。例えば、バッチサイズを25に設定した場合、合計100であるため、$25 \times 32 \times 32$のサイズのデータが4つできることになる。

>> train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)

次のように学習関数を定義して使用できる。

def train(model):
    model.train()
            
    for batch_idx, (X, Y) in enumerate(Data_set):
        
        optimizer.zero_grad()

        output = model(X)
        
        loss = criterion(output, Y)
        loss.backward()
        optimizer.step()

    return loss

全コード

適切なモデルが定義されているとする。

#필요한 라이브러리 import
import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader

#numpy 배열 불러오기
X = np.load("X.npy")
Y = np.load("Y.npy")

#tensor로 데이터 타입 변환
X = torch.from_numpy(X)
Y = torch.from_numpy(Y)

#데이터셋으로 변환
train_data = TensorDataset(X, Y)

#셔플, 배치사이즈 적용
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)

#학습 함수 정의
def train():
    model.train()
            
    for batch_idx, (X, Y) in enumerate(Data_set):
        
        optimizer.zero_grad()

        output = model(X)
        
        loss = criterion(output, Y)
        loss.backward()
        optimizer.step()

    return loss

#학습
for epoch in range(1, EPOCHS):
    loss = train()
    
    print(loss)

環境

• OS: Windows10
• バージョン: Python 3.9.2, torch 1.8.1+cu111