ついにWATLABブログから書籍「いきなりプログラミングPython」が発売しました！
機械学習モデルの構築はPythonを使うのが一般的です。しかしモバイルアプリやその他組み込み環境に機械学習ができるレベルのPython環境を整えるのはかなりシンドイと思います。ここではPyTorchのモデルをTorchScript形式に変換し、容易にデプロイができる状態を目指します。
こんにちは。wat（@watlablog）です。これまで色々とPyTorchで機械学習モデルを作ってきましたが、TorchScript形式に変換することでよりデプロイを意識します！
TorchScriptとは？
TorchScriptは、PyTorchのモデルを中間表現（Intermediate Representation: IR）に変換し、それをシリアライズ可能（保存可能）にすることで、PyTorchの実行環境（ランタイム）がある場所であればPythonなしでモデルをロード・推論できるようにします。
例えばモバイルアプリのフレームワークであるFlutterはGoogleが開発したDart言語を使いますが、flutter_pytorchプラグインを使うことでTorchScript形式で書かれたpytorchモデルを読み込めるとのこと（ChatGPTより）。これは面白そうだと思ったので本当にできるものなのか後で検証しようと思います。
別にモバイルに限らず、Webアプリのバックエンドで動作させるのも良いかも知れません。さらに、TorchScript形式のモデルを使うことで速度UPも期待できる[1]とのことです。TorchScript化する時に直面する課題もあると思うので、まずはやってみるというコンセプトでこの記事を書きました。
題材：2つのサイン波を分類する音声認識モデル
この記事では誰でも簡単にできる題材として、100Hzと1000Hzの音声を分類する音声認識モデルを作りたいと思います。たった2つの周波数の音声であれば超シンプルなCNNモデルで表現できると思うので、機能調査には打って付けではないでしょうか。WATLABブログは音声系が得意なので、音声認識モデルを題材にしておけば今後Flutterでアプリ化する時にマイク入力等も一緒に学べて一石二鳥です。
サンプルwavファイル
読者の皆様がすぐにテストできるように、100Hz.wavと1000Hz.wavをダウンロードできるようにしておきました。これもPythonで作っています。
機械学習モデルの作成
TorchScript無し
まずはTorchScriptを使わず、これまで当ブログで書いてきた書き方で音声認識モデルを書きます。
学習コード
次のコードが学習（トレーニング）を行うためのPythonコードです。 wavというフォルダの中に 100Hz.wavと 1000Hz.wavを入れてコードを実行すると、プログラム実行フォルダの直下に audio_classifier.ptというPyTorchのモデルが保存されます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import torch import torch.nn as nn import torch.optim as optim import torchaudio import os class AudioClassifier(nn.Module): """CNNモデルの定義""" def __init__(self): super(AudioClassifier, self).__init__() self.conv = nn.Sequential( nn.Conv2d(1, 8, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(8, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.AdaptiveAvgPool2d((8, 8)) ) self.fc = nn.Sequential( nn.Linear(16 * 8 * 8, 32), nn.ReLU(), nn.Linear(32, 2) ) def forward(self, x): # 入力 x の形状: [B, 1, F, T] x = self.conv(x) x = x.flatten(1) # [B, 16*8*8] return self.fc(x) def load_wav_to_tensor(path): """wavファイルを読み込み、スペクトログラムに変換する関数""" waveform, sr = torchaudio.load(path) spec = to_spectrogram(waveform) # [1, freq, time] spec = (spec - spec.mean()) / spec.std() return spec # [1, F, T] if __name__ == "__main__": """メイン""" # CPU/GPUを自動判定 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # スペクトログラム変換 to_spectrogram = torchaudio.transforms.Spectrogram( n_fft=256, win_length=256, hop_length=128, power=2.0 ) # トレーニングデータの準備 data = [] labels_map = {'100Hz.wav': 0, '1000Hz.wav': 1} for fname, lbl in labels_map.items(): fp = os.path.join('wav', fname) spec = load_wav_to_tensor(fp) # [1, F, T] data.append((spec, lbl)) # DataLoader loader = torch.utils.data.DataLoader( data, batch_size=2, shuffle=True ) # モデル・損失・オプティマイザ初期化 model = AudioClassifier().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) # トレーニングループ epochs = 30 for epoch in range(1, epochs + 1): model.train() total_loss = 0.0 for inputs, targets in loader: # inputs: [B,1,F,T], targets: [B] inputs = inputs.to(device) targets = targets.to(device) optimizer.zero_grad() preds = model(inputs) loss = criterion(preds, targets) loss.backward() optimizer.step() total_loss += loss.item() print(f'Epoch {epoch:02d}/{epochs}, Loss: {total_loss:.4f}') # モデル保存 torch.save(model.state_dict(), 'audio_classifier.pt') |
テストコード
そしてこちらが学習モデル audio_classifier.pt をロードして推論を行うテストコードです。コード内でwavファイルをパス指定して実行すると、そのファイルの音声が100Hzなのか1000Hzなのかを分類します。学習に使ったデータをそのまま読み込めば100%分類ができていることが確認できるでしょう。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import torch import torchaudio import os class AudioClassifier(torch.nn.Module): """モデル定義（トレーニング時と同じ構造を保持）""" def __init__(self): super(AudioClassifier, self).__init__() self.conv = torch.nn.Sequential( torch.nn.Conv2d(1, 8, kernel_size=3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2), torch.nn.Conv2d(8, 16, kernel_size=3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2), torch.nn.AdaptiveAvgPool2d((8, 8)) ) self.fc = torch.nn.Sequential( torch.nn.Linear(16 * 8 * 8, 32), torch.nn.ReLU(), torch.nn.Linear(32, 2) ) def forward(self, x): x = self.conv(x) x = x.flatten(1) return self.fc(x) # ラベルマップ label_map = {0: '100Hz', 1: '1000Hz'} def predict(wav_path, model, transform, device): """推論関数""" # wavファイルの読み込み waveform, sample_rate = torchaudio.load(wav_path) # スペクトログラム変換 spec = transform(waveform) # [1, freq_bins, time_steps] # 標準化 spec = (spec - spec.mean()) / spec.std() # バッチ次元とチャンネル次元を追加 input_tensor = spec.unsqueeze(0).to(device) # [1,1,freq_bins,time_steps] # 推論 model.eval() with torch.no_grad(): output = model(input_tensor) pred_idx = output.argmax(dim=1).item() return label_map[pred_idx] if __name__ == '__main__': """メイン""" # デバイス設定 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # スペクトログラム変換の定義 spectrogram_transform = torchaudio.transforms.Spectrogram( n_fft=256, win_length=256, hop_length=128, power=2.0 ) # モデルの初期化と重みのロード model = AudioClassifier().to(device) model_path = 'audio_classifier.pt' if not os.path.exists(model_path): raise FileNotFoundError(f"Model file not found: {model_path}") model.load_state_dict(torch.load(model_path, map_location=device)) # 推論するwavファイルパスを指定 wav_file = 'wav/100Hz.wav' # または 'wav/1000Hz.wav' # 推論実行 predicted_label = predict(wav_file, model, spectrogram_transform, device) print(f'Predicted: {predicted_label}') |
TorchScript形式で機械学習モデルを出力するPythonコード
それではTorchScript形式でモデルを保存するコードを書いて検証します。
学習コード
TorchScriptには torchaudioの関数を含めることが可能です。今回のコードではデータの前処理として、 torchaudio.transformsでスペクトログラム変換をしています。せっかくモデルをTorchScript形式としてIRに落とし込んでも、モバイルアプリやWebアプリのバックエンドで複雑な前処理を書き直すのは面倒です。すべての関数が対応可能というわけではないようですが、できるだけ torchの機能でコード化しておけば、モデルにデータを投入するだけで結果を出すという非常にシンプルな構造にできます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import torch import torch.nn as nn import torch.optim as optim import torchaudio import os class AudioClassifier(nn.Module): """CNNモデル（前処理を追加）""" def __init__(self): super(AudioClassifier, self).__init__() # スペクトログラム変換 self.spec = torchaudio.transforms.Spectrogram( n_fft=256, win_length=256, hop_length=128, power=2.0 ) # 畳み込み層 self.conv = nn.Sequential( nn.Conv2d(1, 8, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(8, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.AdaptiveAvgPool2d((8, 8)) ) # 全結合層 self.fc = nn.Sequential( nn.Linear(16 * 8 * 8, 32), nn.ReLU(), nn.Linear(32, 2) ) def forward(self, waveform: torch.Tensor) -> torch.Tensor: # 1) スペクトログラム変換 -> [batch, channel, freq_bins, time_steps] spec = self.spec(waveform) # 2) バッチごと・チャネルごとの mean/std で正規化 mean = spec.mean(dim=[2,3], keepdim=True) std = spec.std(dim=[2,3], keepdim=True) + 1e-6 spec = (spec - mean) / std # 3) CNN に渡す (入力チャネルは 1) x = self.conv(spec) # 4) Flatten x = x.flatten(1) # -> [batch, 16*8*8] return self.fc(x) def load_wav_to_tensor(path: str) -> torch.Tensor: """wavファイルを読み込む関数""" waveform, sr = torchaudio.load(path) # モノラル化 if waveform.size(0) > 1: waveform = waveform.mean(dim=0, keepdim=True) return waveform if __name__ == '__main__': """メイン""" # デバイス設定 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # トレーニングデータ準備 data = [] labels_map = {'100Hz.wav': 0, '1000Hz.wav': 1} for fname, lbl in labels_map.items(): fp = os.path.join('wav', fname) waveform = load_wav_to_tensor(fp) # [1, time] data.append((waveform, lbl)) # DataLoader loader = torch.utils.data.DataLoader(data, batch_size=2, shuffle=True) # モデル・損失関数・オプティマイザ初期化 model = AudioClassifier().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) # トレーニングループ epochs = 30 for epoch in range(1, epochs + 1): model.train() total_loss = 0.0 for inputs, targets in loader: # inputs shape: [batch, channel, time] inputs = inputs.to(device) targets = targets.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.step() total_loss += loss.item() print(f'Epoch {epoch:02d}/{epochs}, Loss: {total_loss:.4f}') # TorchScript形式で保存 model.eval() scripted_model = torch.jit.script(model) scripted_model.save('audio_classifier_scripted.pt') print('TorchScriptモデルを保存しました: audio_classifier_scripted.pt') |
テストコード
こちらがテストコードです。推論する関数にもはやネットワーク構造を記載する必要はなく、形を整えたwavファイルを入力するだけで分類ができます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import torch import torchaudio import os def predict(wav_path, model): """推論する関数（モデルレス）""" # wavの読み込み waveform, sr = torchaudio.load(wav_path) # モノラル化 if waveform.size(0) > 1: waveform = waveform.mean(dim=0, keepdim=True) # バッチ次元追加: [1, channel, time] input_tensor = waveform.unsqueeze(0) # 推論 with torch.no_grad(): output = model(input_tensor) pred_idx = output.argmax(dim=1).item() return LABEL_MAP[pred_idx] # 推論ラベルマップ LABEL_MAP = {0: '100Hz', 1: '1000Hz'} if __name__ == '__main__': """メイン""" # TorchScriptモデルのロード model_path = 'audio_classifier_scripted.pt' if not os.path.exists(model_path): raise FileNotFoundError(f"モデルファイルが見つかりません: {model_path}") model = torch.jit.load(model_path) model.eval() # 推論したいwavファイルを指定 wav_file = 'wav/100Hz.wav' # または 'wav/1000Hz.wav' # 推論実行 predicted_label = predict(wav_file, model) print(f'Predicted label: {predicted_label}') |
特に画像等はないですが、コードを実行すると正確に分類できていることがわかりました。そして気持ち速度も上がっているようです。
…書いていて思ったのですが、もしかして別に機械学習モデルに限らなくてもPyTorchの機能でPythonコードを書いてTorchScriptでモデルにすれば、色々な環境下でPythonコードを走らせることができる？
Pythonの強みである外部ライブラリは変換できないかもしれないけど、PyTorchには多くの便利な関数があるので、TorchScriptを使ったハック的なものができるかも！
後でやってみよう。
まとめ
この記事ではPyTorchの機械学習モデルをさまざまなプラットフォーム上で利用できるTorchScript形式に変換する方法を紹介しました。IR（中間表現）にも
torchaudioのスペクトログラム変換を含めることができることも確認しています。
前処理をモデルに含めることができるとわかったので、今後試す予定のモバイルアプリやWebアプリ側での処理がかなり簡便になるはずです。次回はアプリ側で機械学習モデルを利用する方法について調査してみようと思います。
参考文献
[1] Qiita：TorchScriptを使用してPyTorchのモデルを保存するTorchScriptすごい！！前処理もモデルに含めることができるのは知らなかった！
Xでも関連情報をつぶやいているので、wat(@watlablog)のフォローお待ちしています！