トップ > AI > 「Video Diffusion」による動画の学習と生成

「Video Diffusion」による動画の学習と生成

AI

概要

今回は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()