PyTorchの実行速度 CUDA on Windows10 vs WSL2

Windows10 Insider Program ではWSL2からCUDAが使えるとのことで、下記を参考にインストールしました。qiita.com

その上で、WindowsからCUDAを叩くのと、WSL2から叩くのとでどれぐらい実行時間に差があるのか試しました。

環境

• Windows10 Pro Insider Preview
  • Dev Channel: Build 21343.rs_prerelease.210320-1757
• 32GB RAM
• Intel Core i9-10900K
• NVIDIA GeForce RTX 3070
• PyTorch 1.7.1+cu110

コード

torch_test.py

#!python3
import torch
import torchvision.models as models
import argparse
import time

def run(use_gpu: bool, dtype: torch.dtype):
  # Prepare model and data
  model = models.resnet18()
  tensor = torch.rand((40, 3, 480, 640), dtype = dtype)

  # To GPU
  start_time_total = time.time()
  if use_gpu:
    start_time_gpu = time.time()
    model = model.cuda()
    tensor = tensor.cuda()
    elapsed_time_gpu = time.time() - start_time_gpu
    print('toGPU time = {:3f}[s] ({})'.format(
      elapsed_time_gpu,
      dtype
      )
    )

  # Process
  start_time_proc = time.time()
  _ = model(tensor)
  finish_time = time.time()
  elapsed_time_proc = finish_time - start_time_proc
  elapsed_time_total = finish_time - start_time_total
  
  print('{} proc time = {:3f}[s] ({}), total time = {:3f}[s]'.format(
    'GPU' if use_gpu else 'CPU',
    elapsed_time_proc,
    dtype,
    elapsed_time_total
    )
  )

if __name__ == '__main__':
  print('PyTorch %s' % torch.__version__)

  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  print('device: %s' % device)

  run(False, torch.float32)
  run(True, torch.float32)

実行結果

それぞれ3回実行しました。

Windows

>python torch_test.py
PyTorch 1.7.1+cu110
device: cuda:0
CPU proc time = 2.412968[s] (torch.float32), total time = 2.412968[s]
toGPU time = 1.028965[s] (torch.float32)
GPU proc time = 1.586547[s] (torch.float32), total time = 2.615512[s]

>python torch_test.py
PyTorch 1.7.1+cu110
device: cuda:0
CPU proc time = 2.465967[s] (torch.float32), total time = 2.465967[s]
toGPU time = 1.031965[s] (torch.float32)
GPU proc time = 1.586000[s] (torch.float32), total time = 2.617966[s]

>python torch_test.py
PyTorch 1.7.1+cu110
device: cuda:0
CPU proc time = 2.976967[s] (torch.float32), total time = 2.976967[s]
toGPU time = 1.005968[s] (torch.float32)
GPU proc time = 1.590999[s] (torch.float32), total time = 2.596968[s]

WSL2

$ python torch_test.py
PyTorch 1.7.1+cu110
device: cuda:0
CPU proc time = 2.276961[s] (torch.float32), total time = 2.276962[s]
toGPU time = 2.526399[s] (torch.float32)
GPU proc time = 0.220261[s] (torch.float32), total time = 2.746719[s]

$ python torch_test.py
PyTorch 1.7.1+cu110
device: cuda:0
CPU proc time = 2.329590[s] (torch.float32), total time = 2.329591[s]
toGPU time = 2.428853[s] (torch.float32)
GPU proc time = 0.219387[s] (torch.float32), total time = 2.648300[s]

$ python torch_test.py
PyTorch 1.7.1+cu110
device: cuda:0
CPU proc time = 2.325431[s] (torch.float32), total time = 2.325431[s]
toGPU time = 2.907129[s] (torch.float32)
GPU proc time = 0.219199[s] (torch.float32), total time = 3.126382[s]

結論

おおむね所要時間は同じでした。
ただ細かいことを言うと、CUDA使用時は

• cuda()を呼ぶとき
• 推論するとき

の所要時間が異なっていました（2つを合計すればほぼ一緒のようです）
WinとWSL2（またはLinux全般？）で書き方の違いで時間がかかるタイミングが異なるかもしれないことに注意しないといけないのですね…