概要
今回はVideo Diffusion Modelsを使用した動画の学習と生成を試した。
Video Diffusion Modelsは拡散モデルを用いた動画を生成する手法。 github.com
環境
開発環境は以下の通り。
- OS : Windows11
- GPU : NVIDIA GeForce RTX 3080Ti
セットアップ
Anacondaで仮想環境の作成
conda create -n video-diffusion-pytorch python=3.9 conda activate video-diffusion-pytorch
Pytorchのインストール
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
video-diffusion-pytorchのインストール
pip install video-diffusion-pytorch
学習
データセットの用意
今回、学習に使用したデータセットは
動き回る手書き数字で作成されたデータセットであるMoving MNIST。
http://www.cs.toronto.edu/~nitish/unsupervised_video/
上記リンクから、「mnist_test_seq.npy」をファイルをダウンロード。
以下のコードでnpyファイルをgifに変換する。
import numpy as np import imageio data = np.load('data/mnist_test_seq.npy') for i in range(data.shape[1]): images = data[:, i, :, :] path = f'data/mnist/moving_mnist_{i}.gif' print(path) imageio.mimsave(path, images)
学習の開始
以下のコードで学習を開始。
import torch from video_diffusion_pytorch import Unet3D, GaussianDiffusion, Trainer model = Unet3D( dim = 64, dim_mults = (1, 2, 4, 8), ) diffusion = GaussianDiffusion( model, image_size = 64, num_frames = 10, timesteps = 1000, loss_type = 'l1' # L1 or L2 ).cuda() trainer = Trainer( diffusion, './data/mnist', train_batch_size = 4, train_lr = 1e-4, save_and_sample_every = 1000, train_num_steps = 700000, gradient_accumulate_every = 2, ema_decay = 0.995, amp = True ) trainer.train()
学習過程
今回は70000step学習を行った。
1000step
5000step
2000step
70000step
動画の生成
以下のコードで動画を生成することができる。
import torch from video_diffusion_pytorch import Unet3D, GaussianDiffusion import matplotlib.pyplot as plt from matplotlib.animation import ArtistAnimation NUM_FRAMES = 10 IMG_SIZE = 64 SEED = 0 torch.manual_seed(SEED) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(device) model = Unet3D( dim = 64, dim_mults = (1, 2, 4, 8) ) diffusion = GaussianDiffusion( model, image_size = IMG_SIZE, num_frames = NUM_FRAMES, timesteps = 1000, loss_type = 'l1' # L1 or L2 ).to(device=device) diffusion.load_state_dict(torch.load("results/model-70.pt")['ema']) # 動画生成 sampled_videos = diffusion.sample(batch_size = 1) # 動画データをnumpy配列に変換 sampled_videos = sampled_videos.cpu().numpy() sampled_videos = (sampled_videos + 1) / 2 # -1から+1の範囲を0から1に変換 # フレームごとのグラフを作成し、Artistオブジェクトのリストに追加 ims = [] fig = plt.figure() for i in range(NUM_FRAMES): im = plt.imshow(sampled_videos[0, :, i, :, :].transpose(1,2,0), animated=True) ims.append([im]) # ArtistAnimationオブジェクトを作成 ani = ArtistAnimation(fig, ims, interval=50) # 動画再生 plt.show()