エンジニアリング

StyleGAN3環境構築 Win10 RTX3070 Anaconda

2021年10月StyleGAN3がリリースされたので、早速遊んでみました。GPUとしては普及しているであろうRTX3070ですが、環境構築に少し手間取ったので、手順をメモとしてまとめておきます。
※2023年1月時点でCUDA Toolkit 11.7で構築します。

環境は以下の通り。
Windows 10
RTX 3070
Anaconda

執筆者

Win10+RTX3000シリーズの環境にPytorchをインストールしてStyleGAN3を動かしてみます!

StyleGAN3とは?

2018年12月にNvidiaの研究者によって導入された、リアルな画像を合成するための生成的逆ニューラルネットワーク(GAN)に基づく機械学習システムです。複数の画像の特徴の補間に基づいて、新しい画像を合成することができます。バージョン2はTensorFlowでしたが、バージョン3はPytorchで実装されました。

StyleGAN3環境構築

StyleGAN3に必要なものはPytorchとCUDAです。RTX3070の場合はcudatoolのバージョンが11.7.0で、正しくGPUが認識されていることが必要です。手順を追えば無事に動いてくれます。

GeForce Game Ready ドライバー

NVIDIA公式からGeForce Game Ready ドライバー「516.94-desktop-win10-win11-64bit-international-dch-whql.exe」をダウンロード&インストールします。

CUDA Toolkit

NVIDIA公式からCUDA Toolkit「cuda_11.7.0_516.01_windows.exe」をダウンロード&インストールします。

cuDNN

NVIDIA公式からCUDA Toolkit「cudnn-windows-x86_64-8.5.0.96_cuda11-archive.zip」をダウンロードします。

zipを解凍したら「bin」「include」「lib」フォルダを「C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7」に該当するフォルダにそれぞれコピーして配置します。

以前はPATHを通さないといけませんでしたが、11.7はCUDAのPATHがすでに通っているので環境変数の編集は不要です。

GPU確認

コマンドプロンプトで、nvidia-smi(nVIDIAシステム管理インタフェース)コマンドとCUDA開発に利用するnvcc等のコマンドを実行し、ドライバが正常にインストールされていることとGPUデバイスが問題なく認識できていることを確認します。

nvidia-smi
| NVIDIA-SMI 516.94       Driver Version: 516.94       CUDA Version: 11.7     |
where nvcc
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin\nvcc.exe
nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Jun__8_16:59:34_Pacific_Daylight_Time_2022
Cuda compilation tools, release 11.7, V11.7.99
Build cuda_11.7.r11.7/compiler.31442593_0

Done!うまくいきました。

Pytorch

仮想環境はDockerでもなんでもいいですが、私はAnacondaを使ってインストール。
https://www.anaconda.com/products/individual

Python3.8でstyleganという環境を作ってアクティブ、ローカルにstyleganというフォルダを作成します。作業フォルダにcdしたらPytorchをpipインストールします。

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
pip list
torch                  1.13.0

インストールDoneです。

StyleGAN3

pytorchがインストールができたら、StyleGAN3のファイルを作業フォルダにgit pullして持ってきます。
https://github.com/NVlabs/stylegan3

コンパイルにはMicrosoft VisualStudioが必要なので、Visual Studio Community Editionをインストールし、PATHを通します。

C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Auxiliary\Build\vcvars64.bat

stylegan/torch_utilis/custom_ops.py

StyleGAN3の方も通しておきます。

def _find_compiler_bindir():
    patterns = [
        'C:/Program Files (x86)/Microsoft Visual Studio/2019/Community/VC/Tools/MSVC/14.29.30133/bin/Hostx64/x64',
    ]

構築完了

オフィシャルのチュートリアルを試して完了!

python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

学習データからlerp.mp4という動画が生成されます。

ヌルヌル動いて、StyleGAN3を使ってAI生成の動画が作れました。ローカル環境なので、Google Colabよりも早く機械学習が行えます。Win10+RTX3000シリーズでPytrorch環境を構築することができたので、いろいろなライブラリを試してみたい方など、ぜひ参考になれば幸いです!