ライブラリ経由で事前学習済みモデルをダウンロードする際のフォルダ指定

お疲れ様です。

Pythonでディープラーニングのプログラムを書く際、その際に様々なライブラリにお世話になるかと思います。
その中で、モデルアーキテクチャの定義と同時に特定の事前学習済みモデルの重みファイルを自動でダウンロードしてくれる機能があるライブラリもあります。
それらダウンロードされる事前学習済みモデルのダウンロード先フォルダをこちらで指定する方法を調べたのでこちらにも備忘録としておいておきます。

私はなんとなくプロジェクトごとに使ったモデルがわかるようにしたくて、プロジェクトフォルダの直下にダウンロードした重みファイル用のフォルダを作って管理しています。
GitHubのリモートリポジトリにはpushされないように.gitignoreで指定はしていますが…。

ダウンロードフォルダを指定しない場合

Windowsの場合、ユーザフォルダ直下に「.cache」というフォルダが作成され、その中に保存されます。
これはpytorch系やhuggingfaceいずれの場合も同じです。

ダウンロードフォルダの指定

torchvision, torch.hub

環境変数"TORCH_HOME"か、torch.hub.set_dirに保存先フォルダのパスを指定することで可能です。環境変数の場合はコンソールでsetコマンドを使って指定することも可能かと思います。

os.environ["TORCH_HOME"] = "weights"

torch.hub.set_dir("weights")

余談ですが、torchvisionで事前学習モデルをダウンロードする際の指定の仕方が最近のバージョンでは変更されているようです。（最近になって知りました…。）

from torchvision import models

models.resnet50(pretrained=True) # 非推奨（2024/8現在は実行すると警告が出る）

models.resnet50(weights=models.ResNet50_Weights.DEFAULT) # 推奨の指定方法

Huggingface

環境変数"HF_HOME"を指定します。環境変数なのでsetコマンドでも可能です。

os.environ["HF_HOME"] = "weights"

こちらの場合少し注意が必要で、huggingfaceのライブラリ（transformers等）をimportする前にあらかじめ指定しておかないと反映されません。

import os
os.environ["HF_HOME"] = "weights" # 先に事前学習モデルの保存先を指定

from transformers import DetrConfig, DetrForObjectDetection

model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")

timm

ダウンロード先がPytorch系の場合とHuggingfaceの場合の2パターンがあるので、利用したいモデルに合わせて上記の方法をお試しください。
（最近だとHuggingfaceからのダウンロードが多いのかな…。）

誰かの技術置き場