WSL+ローカルGPUによる深層学習の環境構築

はじめに

私の研究室は学生数が多く、全員が機械学習・深層学習を使った研究を行っているため、空前のGPU不足に悩まされています。

応急的な対策としてGoogle ColabなどのGPUを使うことが挙げられますが、セキュリティ問題や利用時間の制限などが伴います。

また、研究用途ではなく趣味の一環で、ネット上にデータを上げたくないけど機械学習・深層学習に挑戦してみたいという方もいると思います。

そこで本記事では、WSL(Windows Subsystem for Linux)と、PCに備えられているGPU（ローカルGPU）を用いた機械学習・深層学習の環境構築について紹介します。

特に大学生・大学院生で、ローカルGPUの性能が良い人は、研究室でのGPUの取り合いから解放される可能性がありますので、導入を検討してみてください。

同じような試み自体は、既にネット上にたくさんの記事がありますが、どれも少しずつ方法が違うので、本記事もそのバリエーションの一つとして捉えてください。

環境構築

参考までに、私のPC環境は以下のとおりです。

• OS：Windows11
• GPU：NVIDIA GeForce RTX 3060Ti

本記事は、私が環境構築をしたときの経験を元に執筆しています。
私が参考にしたWebページを示しておくので、困ったら参照してください。
私の記事よりそちらのほうが分かりやすいと思います。

• WSL上のUbuntuでのNVIDIA CUDAツールキットなどのインストールと動作確認
• WSL2でGPUを使う

WSLの導入

以下のページに従ってWindows環境にWSLを導入してください。

• WSL2のインストール、WSL2上へのUbuntuのインストールと利用

WSLの導入が完了すれば、WindowsやVSCodeなどのターミナルでwslを入力することで起動できます。

NVIDIAドライバのインストール

NVIDIAドライバのダウンロードページにアクセスし、ご自身の環境に合うインストーラをダウンロードしてください。

ここでは、OSはWindowsを選択します。

ダウンロードしたファイルを実行すると、ドライバがインストールされます。

ローカルGPUの詳細は、以下の手順で確認できます。

1. タスクバーを右クリック
2. タスクマネージャーを選択
3. タブから「パフォーマンス」を選択
4. 「GPU」を選択

Python環境の整備

VSCodeなどのターミナルからWSLに入り、一通りのPython実行環境を整えてください。

これには色々な方法がありますし、私はそこまで詳しくないので、ご自身でお好みの方法を調べていただければと思います。

ちなみに私は、1年以上前にMediaPipeを使ったときに、知らない間に仮想環境を作っていたようで、惰性でそれをそのまま使い続けています。
厳密には、環境は分けたほうが良いような気もしますが……。

一応、私が見たMediaPipeの導入記事はこれです。

PyTorchのインストール

機械学習・深層学習のフレームワークとしてはTensorFlowとPyTorchが有名です。
これも私は詳しくないのですが、どうやらPyTorchのほうが環境構築は簡単そうだったので、本記事ではPyTorchで説明していきます。

PyTorchのWebページにアクセスし、「INSTALL PYTORCH」欄で次のように選択してください。

• PyTorch Build：Stable
• Your OS：Linux（WSL上で動かすのでWindowsではありません）
• Package：Pip
• Language：Python
• Compute Platform：お好みのバージョン（動かしたい深層学習モデルに依存する可能性あり）

コマンドが表示されるので、WSL（あるいは仮想環境）で実行してください。
権限の問題でコマンドが通らないときはsudoをつけてください。

環境構築の確認

以上の手順が正常にできていれば、環境構築は終了です。（簡単！）

正常にできているかどうかは、下記コマンドをWSL（あるいは仮想環境）で実行すれば分かります。

二つ目のコマンドに関しては、PyTorchのバージョンとTrueが表示されたらOKです。

• WSLがGPUを認識しているか：nvidia-smi
• PyTorchとGPUを利用できるか：python3 -c "import torch; print(torch.__version__, torch.cuda.is_available())"

Hello Deep Learning!

MNISTの手書き文字認識を想定した簡易的なニューラルネットワークを部分的にChatGPTを使って作成したので、これをCPUとGPUで実行したときの計算時間を比較します。

ネットワークが単純なため、CPUとGPUの性能の兼ね合いによってはデフォルト値でGPUの演算性能を実感できません。
必要に応じて、実行時のオプション引数からパラメータを変えてみてください。

なお、学習中のGPUの稼働状況はタスクマネージャーから確認できます。

gpu_test.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import argparse
import time

##################################################
def parseArgs():
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", type=str, default="cuda", help="set device as 'cuda' or 'cpu' [default: cuda]")
    parser.add_argument("--data_dimension", type=int, default=65536, help="dimension of input data [default: 65536, which means a 256*256 image]")
    parser.add_argument("--data_samples", type=int, default=10000, help="amount of training data [default: 10000]")
    parser.add_argument("--batch_size", type=int, default=64, help="batch size during training [default: 64]")
    parser.add_argument("--epoch", type=int, default=5, help="epoch to run [default: 5]")
    return parser.parse_args()

##################################################
# define simple neural network
class SimpleNN(nn.Module):
    def __init__(self, data_dimension):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(data_dimension, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)  # 出力層: 10クラス（MNISTを想定）

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

##################################################
def computeNN(args):
    # set device
    device = torch.device(args.device)
    print(f"Using device: {device}")

    # prepare dataset
    x_train = torch.randn(args.data_samples, args.data_dimension)
    y_train = torch.randint(0, 10, (args.data_samples,))  # サンプル数に応じたラベル（MNIST想定なので0～9）

    # define data loader
    dataset = TensorDataset(x_train, y_train)
    train_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

    # set up model
    model = SimpleNN(args.data_dimension).to(device)

    # define loss function and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # start of training
    start = time.time()

    # train
    for epoch in range(args.epoch):
        model.train()
        running_loss = 0.0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)  # データをデバイスに転送
        
            # 順伝播
            optimizer.zero_grad()  # 勾配を初期化
            outputs = model(inputs)
            loss = criterion(outputs, labels)
        
            # 逆伝播と最適化
            loss.backward()
            optimizer.step()
        
        running_loss += loss.item()
        print(f"Epoch {epoch+1}/{args.epoch}, Loss: {running_loss/len(train_loader)}")

    # end of training
    end = time.time()

    model.eval()
    print("Training complete!")
    print(f"time: {end - start}[s]")

##################################################
if __name__=="__main__":
    args = parseArgs()
    computeNN(args)

画像中のtimeの単位は「秒」です。
スクリーンショットを撮ったときは忘れていたのですが、本記事のプログラム中では単位も表示するようにしています。

注意事項

WSLのメモリ

GPUメモリにモデルとデータが乗るかどうかに加え、WSLを使う場合はWSLに割り当てられたメモリに空きがあるかどうかも確認する必要があります。

以下に私の場合を示します。

freeが空き容量なので、この場合は合計で18GB使えることになります。

容量不足になると学習が一切進まなくなるので、利用可能なメモリ容量とご自身のデータセットのサイズを事前に確認しておきましょう。

WSL+ローカルGPUのメリット

• 手軽に環境構築ができる。
• ローカルで学習を回すため、Google Colabと比べてセキュリティ的に安全、かつ利用時間が無制限。
• 研究室に所属する大学生・大学院生にとっては、ローカルGPUの性能が良ければGPUの取り合いから解放される可能性あり。（そもそもローカルGPUを使っている場合を除く）

WSL+ローカルGPUのデメリット

• Ubuntuを直接使っていないので、学習が遅い。（参考ページ）
• RAMが少なく、データロードでメモリが一杯になる場合、学習を回しているときにPCを使えない。

所感

PointNet++(PyTorch)のセマンティックセグメンテーションを試してみました。
私の環境では、S3DISデータセットがメモリ自体には乗ったものの、WSLの容量が不足しており、学習できませんでした。
恥ずかしながら、それに気づくまでに数日要しました。

データを間引いてみると正常に学習が進んでいたので、この環境は大きなデータセットを使いたい人には向いていないかもしれません。

学習速度については、遅いようには私は感じませんでした。
ただし、何かと比較しているわけではなく、あくまで個人的な体感に過ぎないので、その点はご注意ください。

おわりに

本記事では、WSLとローカルGPUを使った深層学習の環境構築、およびそのメリット・デメリットなどについて紹介しました。

質問や疑問などは、コメント欄かX(Twitter)でご連絡ください。

本記事が少しでもお役に立てば幸いです。