Google Colab を使って快適にKaggleをする

Kaggleをするとき、ネックになるのがGPUの実行時間です。デフォルトでは週30時間、月1000円のColab Pro に入っても週45時間です。ここで、Colab Proを契約すれば1000円で月50時間ほどのクオータを追加することができます。16GBのGPUを自宅に用意することを考えればかなりお手頃です。

Colabの場合は実験管理などがKaggleと同じようにはいきません。ということで、なるべく快適にColabでも使えるように行っているスニペットなどを紹介します。

1. Kaggle から Colab を開く

Kaggle Notebookのエディタから「File」「Open in Colab」を選ぶことで、同じノートブックをColabで開いてくれます。このとき、先頭にいくつかセルが追加されています。

１つ目はKaggleのログインを行うもので、２つ目はKaggleからコンペやimportしたデータをダウンロードしてくれるものです。

まずKaggleのログイン画面ですが、 環境変数に KAGGLE_KEY と KAGGLE_USERNAME を入れておくことでスキップできます！ 何度も実験を回すときに毎回ログインなんてやってられないので、はじめに設定するのがおすすめです。

さて、ダウンロードしたデータはそれぞれに自動的にパスが割り振られます。このパスはデータセットのURLが変わらなければ同じなので、データ読み出しなどはこのパスから行います。

たとえば、コンペのディレクトリから train.csv を取ってきたいとします。このときは、

COMPETITION_DIR = rsna_intracranial_aneurysm_detection_path #Colab
#COMPETITION_DIR = "/kaggle/input/rsna-intracranial-aneurysm-detection/" #Kaggle
train_df = pd.read_csv(COMPETITION_DIR + "train.csv")

#COMPETITION_DIR = rsna_intracranial_aneurysm_detection_path #Colab
COMPETITION_DIR = "/kaggle/input/rsna-intracranial-aneurysm-detection/" #Kaggle
train_df = pd.read_csv(COMPETITION_DIR + "train.csv")

のように書けば、ColabでもKaggleでも対応できます。
しかし、毎回実行環境を気にしてどちらを使うか変えるのは面倒です。実行環境を自動で判定できるようにしましょう。

2. 実行環境を自動判定する

import sys
if "/content" in sys.path:
    IS_RUNNING_IN_COLAB = True
else:
    IS_RUNNING_IN_COLAB = False

とても単純です。/content フォルダは Colab にはあっても Kaggle 環境にはないので、これで実行環境がどちらかの判定ができます。走らせるたびに手動で切り替えてもいいのですが、面倒なのでこれをセルの頭に書くことが多いです。

さっきのコードはこのように書き直せます。

if IS_RUNNING_IN_COLAB:
    COMPETITION_DIR = rsna_intracranial_aneurysm_detection_path #Colab
else:
    COMPETITION_DIR = "/kaggle/input/rsna-intracranial-aneurysm-detection/" #Kaggle
train_df = pd.read_csv(COMPETITION_DIR + "train.csv")

3. 自動でGoogle Driveに出力を保存させる

Kaggle Notebookでは出力 (モデルの重みなど) はノートブックにひも付きますが、Google Drive ではランタイムを切ると消えてしまいます。ということで、Google Drive 上にフォルダを作って保存してもらいます。

Google Drive はデフォルトでは /content/drive/MyDrive/ 以下にアタッチされますので、その下にKaggleディレクトリ、コンペのディレクトリ、実験名のディレクトリを作り、そこに保存します。

EXPERIMENT_NAME = "001_baseline" #各実験の名前
if IS_RUNNING_IN_COLAB:
    from google.colab import drive
    drive.mount('/content/drive')
    EXPORT_DIR = f"/content/drive/MyDrive/kaggle/COMPETITION_NAME/{EXPERIMENT_NAME}/"
    os.makedirs(EXPORT_DIR, exist_ok = True)
    assert os.path.exists(EXPORT_DIR), "The EXPORT_DIR does not exist. "
else:
    EXPORT_DIR = "./"

#重みの保存時は次のように書く

torch.save(model.state_dict(), EXPORT_DIR + "model.pt")

モデルによっては一回の実験で数百MBになることもあると思います。ストレージに課金しているか、Geminiの有料プランがあるなら気にすることはないですが、Google Drive は無課金だと上限は15GB程度なので注意が必要です。

Kaggleからデータを読み込むときは、一旦ローカルに落としてDatasetやModelなどにするのがよいと思います。

4. 自動で環境変数を取得する

WandbやHF, OpenAI API を使いたいときは環境変数にトークンを格納すると思います。環境変数の読み方もKaggle と Colab で違うので、次のように一元化すると楽です。この例はWandb ですが、他のものでも同じように使えます。

if IS_RUNNING_IN_COLAB:
  from google.colab import userdata
  api_key = userdata.get('wandb_api')
else:
  from kaggle_secrets import UserSecretsClient
  user_secrets = UserSecretsClient()
  api_key = user_secrets.get_secret("wandb_api")
wandb.login(key=api_key)

5. 足りないパッケージを自動で取ってこさせる

あんまりないですが、たまにKaggle環境にあるけどColab環境にないパッケージがあります。逆もあるかもしれません。
よく使いそうなものでは、医療画像に使う DICOM を読むための pydicom がColabにはありません。ということで、このようにします。

if IS_RUNNING_IN_COLAB:
  !pip install pydicom
import pydicom

! のマジックコマンドはコード中に突然書いても動いてくれます。便利です。

さて、このようなスニペットをコードに仕込んでおくと、クオータが足りないときにも Open in Colab を押して実行をすれば動いてくれるのでとても便利です。GPUをぶん回すときにはぜひ使ってみてください。