[自分用メモ]Google Colab ProでどのGPUを使えばいい？

背景

現在Kaggleのコンペに参加していて、KaggleのNotebookに参加しています。
Kaggle NotebookだとGPUの処理速度に限界を感じた為、Google Colab pro(月額約1,100円）に課金。処理速度を上げるべくGPUを使用するよう設定変更を行いました。

各GPUによって処理時間がどう変わるのか

最強GPUはA100とのことでまずは設定。
Kaggle NotebookをGoogle Colabで開くことで自動生成される、以下のコードを各GPUで実行してみた・

# IMPORTANT: RUN THIS CELL IN ORDER TO IMPORT YOUR KAGGLE DATA SOURCES
# TO THE CORRECT LOCATION (/kaggle/input) IN YOUR NOTEBOOK,
# THEN FEEL FREE TO DELETE THIS CELL.
# NOTE: THIS NOTEBOOK ENVIRONMENT DIFFERS FROM KAGGLE'S PYTHON
# ENVIRONMENT SO THERE MAY BE MISSING LIBRARIES USED BY YOUR
# NOTEBOOK.

import os
import sys
from tempfile import NamedTemporaryFile
from urllib.request import urlopen
from urllib.parse import unquote, urlparse
from urllib.error import HTTPError
from zipfile import ZipFile
import tarfile
import shutil

CHUNK_SIZE = 40960
DATA_SOURCE_MAPPING = 'leash-BELKA:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-competitions-data%2Fkaggle-v2%2F67356%2F8006601%2Fbundle%2Farchive.zip%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240520%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240520T104931Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3D76e1ae2ff311e8444c7a61b856709702c0a3c7513bee8b19d612acda189dd2a960d7e2c8fd7d3a4c301aecdf545439e7a2792ab2c29538854b16845694762afa7072f00623ade0c89b093228421262937a2d7f174c8f9bceb53ef63290eb4c3e5f591cb07662e5ba0c9bededd22c13dae40d23bf8a03136a4fbc4da354c9d77f2cde49eae7dda8f480ef17257a34992eff3131018c161aedf55dae4b719db767f4402c7e90bcc625cf8db375463da0b559d4e936b6e22b64c66a589d34dd98b036c4e0b4c67f6dad4f3e614e8d6f7232e4e8ef2e8163704f12af4239a35a5de905c0040fc84914231d150f66e49e864e425ea03d9f8324a3536958280a6b4c32'

KAGGLE_INPUT_PATH='/kaggle/input'
KAGGLE_WORKING_PATH='/kaggle/working'
KAGGLE_SYMLINK='kaggle'

!umount /kaggle/input/ 2> /dev/null
shutil.rmtree('/kaggle/input', ignore_errors=True)
os.makedirs(KAGGLE_INPUT_PATH, 0o777, exist_ok=True)
os.makedirs(KAGGLE_WORKING_PATH, 0o777, exist_ok=True)

try:
  os.symlink(KAGGLE_INPUT_PATH, os.path.join("..", 'input'), target_is_directory=True)
except FileExistsError:
  pass
try:
  os.symlink(KAGGLE_WORKING_PATH, os.path.join("..", 'working'), target_is_directory=True)
except FileExistsError:
  pass

for data_source_mapping in DATA_SOURCE_MAPPING.split(','):
    directory, download_url_encoded = data_source_mapping.split(':')
    download_url = unquote(download_url_encoded)
    filename = urlparse(download_url).path
    destination_path = os.path.join(KAGGLE_INPUT_PATH, directory)
    try:
        with urlopen(download_url) as fileres, NamedTemporaryFile() as tfile:
            total_length = fileres.headers['content-length']
            print(f'Downloading {directory}, {total_length} bytes compressed')
            dl = 0
            data = fileres.read(CHUNK_SIZE)
            while len(data) > 0:
                dl += len(data)
                tfile.write(data)
                done = int(50 * dl / int(total_length))
                sys.stdout.write(f"\r[{'=' * done}{' ' * (50-done)}] {dl} bytes downloaded")
                sys.stdout.flush()
                data = fileres.read(CHUNK_SIZE)
            if filename.endswith('.zip'):
              with ZipFile(tfile) as zfile:
                zfile.extractall(destination_path)
            else:
              with tarfile.open(tfile.name) as tarfile:
                tarfile.extractall(destination_path)
            print(f'\nDownloaded and uncompressed: {directory}')
    except HTTPError as e:
        print(f'Failed to load (likely expired) {download_url} to path {destination_path}')
        continue
    except OSError as e:
        print(f'Failed to load {download_url} to path {destination_path}')
        continue

print('Data source import complete.')

自分用メモ

・A100:5分~6分程度
・L4:5分
・TPU:15分程
A100はセッション(後述)によっては、割り当ててもらえず、その際はL4というGPUに接続されていました。

セッションとは？

Google ColabでNotebookを実行することをセッションという。
例えば、"アクティブなセッション"というとアクティブになっている＝つまり、実行中になっているセッション（Notebook）という意味になります。

おまけ

testデータにおいて前処理とDataFrameの作成を行うコードは
TPUで21分かかった。これは普通なのでしょうか。。。