[PyTorch]modelとdatasetがcudaモードになっているか確認する

はじめに

PyTorchはテンソルに対してhoge.to(device)などで簡単にcpuとgpuモードを切り替えられますが，よくこのデータセットやモデルがcpuかgpuなのかわからなくなったので，確認する方法を書き残しときます．

確認方法

前提としてデータセットとモデルの準備は

IMAGE_SIZE=224
BATCH_SIZE=20
TRAIN = 'train'
DATA_DIR = 'dataset/predata/'
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

data_transforms = {
    TRAIN: transforms.Compose([
        transforms.Resize(IMAGE_SIZE),
        transforms.ToTensor(),
    ])
}

# データの前処理
image_datasets = {x: datasets.ImageFolder(os.path.join(DATA_DIR, x), data_transforms[x]) for x in [TRAIN]}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], 
                                              batch_size=BATCH_SIZE, shuffle=True, num_workers=4) for x in [TRAIN]}

is_cuda を使えば bool 型で返ってくるのでここ(1)を参考にして

data, label = iter(dataloaders[TRAIN]).next()
data = data.to(DEVICE)
label = label.to(DEVICE)
print(data.is_cuda)
print(label.is_cuda)

model = model.to(DEVICE)
print(next(model.parameters()).is_cuda)

これで確認できます．

さいごに

ローカル変数抜ける前と抜けた後だとモードってどうなるんだろうね
確認してみるのも面白いかも

参考資料

(1) How to check if Model is on cuda