【PyTorch】データをテスト用・学習用に分割する

大昔に scikit-learn を触ったとき、データの分割が引数だけで完結しました。
PyTorchにもあると持っていたけど、ないっぽいので分割するコードを書いてみました。

※ディープラーニングやPyTorchの初学者です。間違い・勘違いが含まれているかもしれません。

データ

データ数は89行, 13列のものです。

# 0行目はid
# 1行目がターゲット
# 2行目〜は特徴量
array([[0.000e+00, 3.000e+00, 1.236e+01, ..., 5.600e-01, 1.580e+00,
        5.200e+02],
       [1.000e+00, 2.000e+00, 1.242e+01, ..., 9.200e-01, 3.120e+00,
        3.650e+02],
       [3.000e+00, 1.000e+00, 1.305e+01, ..., 8.800e-01, 3.350e+00,
        8.850e+02],

分割

以下のように分割しています。
学習用 -> 70
テスト用 -> 19

引数test_modeの値で「学習用データを取得する」 or 「テスト用データを取得する」を分岐させています。

#　
class WineDataset(Dataset):
    def __init__(self, test_mode=None): ### -> 引数
        xy = np.loadtxt('train.tsv', skiprows=1)
        self.train_x = torch.tensor(xy[:70, 2:], dtype=torch.float32)
        self.train_y = torch.tensor(xy[:70, 1], dtype=torch.float32).reshape(70, 1)
        self.test_x = torch.tensor(xy[70:, 2:], dtype=torch.float32)
        self.test_y = torch.tensor(xy[70:, 1], dtype=torch.float32).reshape(19, 1)
        self.sample_size = self.train_x.shape[0]
        self.test_mode = test_mode ### -> 格納

    def __getitem__(self, index):
        if self.test_mode: ### -> 分岐
            return self.test_x[index], self.test_y[index]
        else: ### -> 分岐
            return self.train_x[index], self.train_y[index]

    def __len__(self):
        return self.sample_size

以下のようにデータを取得します。

# 学習用データ
dataset = WineDataset()
train_x, train_y = dataset[:]

# テスト用データ
dataset = WineDataset(test_mode=True)
test_x , test_y  = dataset[:]

不恰好な気もします。。