日本語解説!PyTorch Hub を活用した学習済みモデルの効率的な利用法

2025-05-16

機能

  1. URLからのダウンロード
    指定されたURLから.pth.ptといった拡張子の付いたファイル(通常はPyTorchのstate_dictを保存したファイル)を自動的にダウンロードします。
  2. state_dictのロード
    ダウンロードしたファイルから、モデルのパラメータが格納された辞書形式のデータ構造であるstate_dictをロードします。
  3. キャッシュ機能
    一度ダウンロードしたファイルはローカルにキャッシュされ、次回以降同じURLを指定した際には再度ダウンロードすることなく、キャッシュされたファイルが利用されます。これにより、何度も同じモデルをロードする際の時間を節約できます。

基本的な使い方

import torch

# モデルの定義 (ここでは簡単な例として)
class SimpleModel(torch.nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.linear = torch.nn.Linear(100, num_classes)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

# 学習済みモデルのstate_dictのURL
url = 'https://download.pytorch.org/models/resnet18-f37072fd.pth' # ResNet18の例

# URLからstate_dictをロード
state_dict = torch.hub.load_state_dict_from_url(url)

# ロードしたstate_dictをモデルに適用
model.load_state_dict(state_dict, strict=False) # strict=Falseはキーが完全に一致しなくてもロードを試みるオプション

print("学習済みモデルの重みをロードしました。")

引数

  • file_name (str, オプション): ダウンロードしたファイルをローカルに保存する際のファイル名を指定します。デフォルトはURLから推測されます。
  • check_hash (bool, オプション): ダウンロードしたファイルのハッシュ値を検証するかどうかを指定します。デフォルトは False です。
  • progress (bool, オプション): ダウンロードの進行状況を表示するかどうかを指定します。デフォルトは True です。
  • map_location (callable or torch.device, オプション): テンソルをどこにロードするかを指定します。例えば、torch.device('cpu')torch.device('cuda') を指定できます。デフォルトは現在のデバイスです。
  • model_dir (str, オプション): ダウンロードしたファイルを保存するローカルディレクトリを指定します。デフォルトはPyTorch Hubのキャッシュディレクトリです。
  • url (str): ダウンロードするstate_dictファイルのURLを指定します。

利点

  • 再利用性
    公開されている多くの学習済みモデルを容易に利用できます。
  • 簡便性
    事前学習済みモデルの重みを数行のコードで簡単にロードできます。
  • インターネット接続が必要です。
  • ロードするstate_dictは、モデルの構造と互換性がある必要があります。完全に同じ構造でなくても、strict=False オプションを使うことで、一部の層の重みだけをロードすることも可能です。


ネットワーク関連のエラー (Network Errors)

  • トラブルシューティング
    • インターネット接続が正常であることを確認してください。
    • 指定したURLが正しいかどうか再度確認してください。ブラウザなどでURLを開いてアクセスできるか試してみるのも有効です。
    • 一時的なサーバーの問題である可能性もあるため、しばらく時間をおいて再度試してみてください。
    • プロキシ環境下で作業している場合は、プロキシの設定がPyTorchに反映されているか確認してください。環境変数 (HTTP_PROXY, HTTPS_PROXY) の設定が必要な場合があります。
  • 原因
    指定されたURLにアクセスできない場合に発生します。ネットワーク接続の問題、URLの誤り、サーバー側の問題などが考えられます。
  • エラー
    urllib.error.URLErrorrequests.exceptions.ConnectionError など。

ファイル関連のエラー (File-Related Errors)

  • トラブルシューティング
    • model_dir 引数を明示的に指定している場合は、そのディレクトリが存在し、書き込み可能であることを確認してください。
    • デフォルトのキャッシュディレクトリに問題がある場合は、PyTorch Hubのキャッシュディレクトリをクリアしてみるのも有効かもしれません。キャッシュディレクトリの場所は torch.hub.get_dir() で確認できます。
  • 原因
    ダウンロードしたファイルを保存するディレクトリが存在しない、または書き込み権限がない場合に発生することがあります。
  • エラー
    FileNotFoundError (キャッシュディレクトリ関連など)。

ハッシュ値の不一致エラー (Hash Mismatch Error)

  • トラブルシューティング
    • check_hash=False を設定してハッシュ値のチェックを無効にする(セキュリティ上のリスクがあるため、信頼できるソースからのファイルでのみ推奨されます)。
    • 再度ファイルをダウンロードしてみる。
    • もし期待されるハッシュ値が分かっている場合は、それを torch.hub.load_state_dict_from_url() の引数として指定することもできます(ただし、この機能は現在のPyTorchのバージョンでは直接提供されていないかもしれません)。
    • モデルの提供元にハッシュ値の情報を確認してください。
  • 原因
    check_hash=True (デフォルトは False) が設定されている場合に、ダウンロードしたファイルのハッシュ値が期待されるハッシュ値と一致しない場合に発生します。これは、ファイルが破損しているか、URL先のファイルが更新された可能性があります。
  • エラー
    RuntimeError: File downloaded from <URL> has different hash. Expected <expected_hash>, got <actual_hash>

state_dict のキーの不一致エラー (Key Mismatch Errors)

  • トラブルシューティング
    • model.load_state_dict(state_dict, strict=False) のように、strict=False オプションを使用して、キーが完全に一致しない場合でもロードを試みるようにします。この場合、一致するキーの重みだけがロードされ、不一致のキーは無視されます。
    • ロードしようとしている state_dict が、自分のモデルの構造と互換性があるかどうかを確認してください。
    • 必要に応じて、自分のモデルの層の名前を事前学習済みモデルの層の名前に合わせる、または事前学習済みモデルの一部だけをロードして利用するなどの調整を行います。
  • 原因
    ダウンロードした state_dict が、現在使用しているモデルの構造と完全に一致しない場合に発生します。これは、事前学習済みモデルと自分のモデルの層の名前や構造が異なる場合に起こります。
  • エラー
    RuntimeError: Error(s) in loading state_dict for <YourModel>: から始まる、ロードする state_dict のキーがモデルの構造と完全に一致しないことに関するエラーメッセージ。

デバイス関連のエラー (Device-Related Errors)

  • トラブルシューティング
    • 利用可能なデバイス ('cpu''cuda') を map_location に正しく指定してください。
    • CUDAを利用する場合は、CUDAが正しくインストールされ、GPUが認識されていることを確認してください。
  • 原因
    map_location 引数で指定したデバイスが利用できない場合に発生することがあります。例えば、CUDAが利用できない環境で map_location='cuda' を指定した場合などです。
  • エラー
    CUDA関連のエラーなど。
  • シンプルなケースで試す
    まずは簡単なモデルや、よく知られた学習済みモデルのURLで試してみて、基本的な動作を確認すると良いでしょう。
  • PyTorchのバージョンを確認
    PyTorchのバージョンによって、関数の挙動や利用可能な機能が異なる場合があります。使用しているPyTorchのバージョンを確認し、必要に応じてドキュメントを参照してください。
  • エラーメッセージをよく読む
    エラーメッセージには、問題の原因や解決のヒントが含まれていることが多いです。


基本的な例: 事前学習済みResNet18の重みをロードする

この例では、PyTorch Hubで公開されている事前学習済みの ResNet18 モデルの重みをダウンロードし、ロードします。

import torch
import torchvision.models as models

# ResNet18モデルのインスタンスを作成 (構造だけ)
model = models.resnet18()

# 事前学習済みResNet18のstate_dictのURL
url = 'https://download.pytorch.org/models/resnet18-f37072fd.pth'

# URLからstate_dictをロード
state_dict = torch.hub.load_state_dict_from_url(url)

# ロードしたstate_dictをモデルに適用
model.load_state_dict(state_dict)

# モデルを評価モードにする (推論に備える)
model.eval()

print("事前学習済みResNet18の重みをロードし、モデルを評価モードにしました。")

# 適当な入力データで推論を試す (オプション)
dummy_input = torch.randn(1, 3, 224, 224)
output = model(dummy_input)
print(output.shape)

解説

  1. import torchimport torchvision.models as models で必要なライブラリをインポートします。
  2. models.resnet18() で ResNet18 モデルのインスタンスを作成します。この時点では、モデルの構造だけが定義されており、重みはランダムに初期化されています。
  3. url 変数に、事前学習済み ResNet18 の重みが保存されているURLを指定します。
  4. torch.hub.load_state_dict_from_url(url) を呼び出すことで、指定されたURLから重みがダウンロードされ、state_dict としてロードされます。
  5. model.load_state_dict(state_dict) を使用して、ロードした重みをモデルの対応する層に適用します。
  6. model.eval() は、モデルを評価モードに設定します。これにより、バッチ正規化層やドロップアウト層が推論に適した挙動になります。
  7. オプションとして、適当な入力データ (dummy_input) を作成し、ロードしたモデルで推論を実行してみることで、モデルが正常に動作しているかを確認できます。

応用例1: 特定のデバイスに重みをロードする

map_location 引数を使用すると、ダウンロードした重みを特定のデバイス(CPUまたはGPU)にロードできます。

import torch
import torchvision.models as models

model = models.resnet18()
url = 'https://download.pytorch.org/models/resnet18-f37072fd.pth'

# CUDAが利用可能な場合はGPUに、そうでない場合はCPUにロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
state_dict = torch.hub.load_state_dict_from_url(url, map_location=device)
model.load_state_dict(state_dict)
model.to(device) # モデル自体も同じデバイスに移動

print(f"重みを {device} にロードしました。")

解説

  • model.to(device) を使用して、モデル自体も同じデバイスに移動させる必要があります。
  • load_state_dict_from_url()map_location 引数に決定したデバイスを指定することで、ダウンロードされた重みがそのデバイスに直接ロードされます。
  • torch.device() を使用して、利用可能なデバイス ('cuda' または 'cpu') を決定します。

応用例2: strict=False を使用して一部の重みをロードする

モデルの構造が完全に一致しない場合でも、strict=False オプションを使用すると、名前が一致する層の重みだけをロードできます。

import torch
import torch.nn as nn

# 事前学習済みモデルと異なる構造のモデル
class CustomModel(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        # ResNet18のlayer1に相当する層 (名前は異なる)
        self.layer_one = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes) # 出力クラス数を変更

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer_one(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

model = CustomModel(num_classes=10) # 出力クラス数を10に変更
url = 'https://download.pytorch.org/models/resnet18-f37072fd.pth'
state_dict = torch.hub.load_state_dict_from_url(url)

# strict=False を指定して、名前が一致する層の重みだけをロード
model.load_state_dict(state_dict, strict=False)

print("一部の互換性のある重みをロードしました。")
  • model.load_state_dict(state_dict, strict=False) とすることで、state_dict のキーと CustomModel の層の名前が一致する部分の重みだけがロードされます。不一致の部分は無視されます。これは、事前学習済みモデルの特徴抽出部分を利用し、最後の分類層だけを自分のタスクに合わせて再学習する場合などに役立ちます。
  • CustomModel は、ResNet18 と完全に同じ構造ではありません。特に、最終的な全結合層の出力クラス数が異なっていたり、中間層の名前が変更されていたりします。


torch.load() を直接使用してローカルファイルをロードする

もし学習済みモデルの重みファイル (.pth.pt) が既にローカルに保存されている場合、torch.load() 関数を直接使用して state_dict をロードできます。

import torch
import torchvision.models as models

# ResNet18モデルのインスタンスを作成
model = models.resnet18()

# ローカルに保存されたstate_dictファイルのパス
local_path = 'path/to/your/resnet18-f37072fd.pth' # 実際のパスに置き換えてください

# ローカルファイルからstate_dictをロード
state_dict = torch.load(local_path)

# ロードしたstate_dictをモデルに適用
model.load_state_dict(state_dict)

model.eval()
print("ローカルファイルから学習済みモデルの重みをロードしました。")

利点

  • ダウンロードの失敗や中断のリスクがありません。
  • ファイルがローカルにあるため、ロードが高速になる可能性があります。
  • インターネット接続が不要です。

欠点

  • ファイルの管理が必要になります。
  • 事前に手動でファイルをダウンロードしておく必要があります。

torch.hub.load() を使用してモデル全体をロードする

import torch

# PyTorch HubからResNet18モデル(学習済み重み込み)をロード
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
model.eval()
print("PyTorch Hubから事前学習済みResNet18モデルをロードしました。")

# 特定のバージョンを指定する場合
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True, source='github')

# ローカルのキャッシュディレクトリを指定する場合
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True, cache_dir='./my_cache_dir')

解説

  • cache_dir 引数で、ダウンロードしたファイルを保存するローカルディレクトリを指定できます。
  • source 引数で、モデルのソース(デフォルトは 'github')を指定できます。
  • torch.hub.load('リポジトリ名', 'モデル名', pretrained=True, ...) のように、リポジトリ名(例: 'pytorch/vision:v0.10.0')、モデル名(例: 'resnet18')、pretrained=True を指定することで、学習済みの重みが自動的にダウンロードされ、モデルにロードされます。

利点

  • 特定のバージョンやキャッシュディレクトリなどを指定できます。
  • PyTorch Hub がモデルの管理や提供を行っているため、信頼性が高いことが多いです。
  • モデルの定義と重みのロードを一度に行えるため、コードが簡潔になります。

欠点

  • モデルの構造が自分の期待するものと完全に一致しない場合があります。
  • load_state_dict_from_url() ほど柔軟に重みのロード方法をカスタマイズできません(例えば、一部の層だけをロードするなど)。

自分で state_dict を作成してロードする

より高度なケースでは、自分で state_dict を作成し、それをモデルにロードすることも可能です。これは、例えば、複数の事前学習済みモデルの重みを組み合わせて新しいモデルを初期化する場合などに使用されます。

import torch
import torch.nn as nn
import torchvision.models as models

# 2つの事前学習済みモデルのstate_dictをロード
url1 = 'https://download.pytorch.org/models/resnet18-f37072fd.pth'
state_dict1 = torch.hub.load_state_dict_from_url(url1)

url2 = 'https://download.pytorch.org/models/alexnet-owt-d8f56d90.pth'
state_dict2 = torch.hub.load_state_dict_from_url(url2)

# 新しいモデルの定義 (ここでは簡単な例として)
class CombinedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        # ... 他の層 ...
        self.fc = nn.Linear(256 * 6 * 6, 10) # AlexNetのfc6の入力サイズを仮定

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        # ... 他の処理 ...
        x = x.view(-1, 256 * 6 * 6)
        x = self.fc(x)
        return x

model = CombinedModel()
new_state_dict = {}

# ResNet18の最初の畳み込み層の重みを新しいモデルにコピー
if 'conv1.weight' in state_dict1 and 'conv1.weight' in model.state_dict():
    new_state_dict['conv1.weight'] = state_dict1['conv1.weight']
    if 'bn1.weight' in state_dict1 and 'bn1.weight' in model.state_dict():
        new_state_dict['bn1.weight'] = state_dict1['bn1.weight']
        new_state_dict['bn1.bias'] = state_dict1['bn1.bias']
        new_state_dict['bn1.running_mean'] = state_dict1['bn1.running_mean']
        new_state_dict['bn1.running_var'] = state_dict1['bn1.running_var']
        new_state_dict['bn1.num_batches_tracked'] = state_dict1['bn1.num_batches_tracked']

# AlexNetの特定の全結合層の重みを新しいモデルにコピー (名前を調整)
if 'classifier.6.weight' in state_dict2 and 'fc.weight' in model.state_dict():
    new_state_dict['fc.weight'] = state_dict2['classifier.6.weight']
    new_state_dict['fc.bias'] = state_dict2['classifier.6.bias']

# 作成したstate_dictをモデルにロード (strict=False を使用することが多い)
model.load_state_dict(new_state_dict, strict=False)

print("複数の事前学習済みモデルの重みを組み合わせてロードしました。")

解説

  • strict=False を使用して、新しい state_dict に含まれていないキーがあってもエラーにならないようにします。
  • 層の名前が異なる場合は、手動でマッピングを行う必要があります。
  • 複数の state_dict をロードし、それらのキーと値を調べて、新しいモデルの対応する層に重みをコピーします。

利点

  • 複数の事前学習済みモデルの知識を組み合わせることができます。
  • 非常に柔軟に重みを操作できます。
  • モデルの構造や層の名前を深く理解している必要があります。
  • 実装が複雑になる可能性があります。