nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[refactor] added more methods in abstract classes

Browse Source 
Former-commit-id: a535e3a06eb766bebe7baa3d9df64b04d04d40e6
This commit is contained in:
qqingzheng
2024-03-06 17:23:40 +00:00
parent 7a78ef2a97
commit b409ef399a
9 changed files with 300 additions and 461 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/rec_video.py
+8 -4
View File
@@ -1,3 +1,6 @@
import sys
sys.path.append(".")
import random
import cv2
import numpy as np
@@ -81,8 +84,10 @@ def main(args):
sample_fps = args.sample_fps
sample_rate = args.sample_rate
device = torch.device('cuda')
vqvae = VQVAEModel.download_and_load_model(args.ckpt)
if args.ckpt in ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']:
vqvae = VQVAEModel.download_and_load_model(args.ckpt)
else:
vqvae = VQVAEModel.load_from_checkpoint(args.ckpt)
vqvae.eval()
vqvae = vqvae.to(device)
@@ -99,8 +104,7 @@ if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--video-path', type=str, default='')
parser.add_argument('--rec-path', type=str, default='')
parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4',
choices=['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4'])
parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4')
parser.add_argument('--sample-fps', type=int, default=30)
parser.add_argument('--resolution', type=int, default=336)
parser.add_argument('--crop-size', type=int, default=None)
opensora/models/ae/videobase/configuration_videobase.py
+16 -4
View File
@@ -1,5 +1,4 @@
from dataclasses import dataclass
import json
"""
We are currently working on the abstraction of the code structure,
@@ -7,6 +6,19 @@ so there is no content here for the moment.
It will be gradually optimized over time.
Contributions from the open-source community are welcome.
"""
@dataclass
class VideoBaseConfiguration:
...
def __init__(self, **kwargs):
pass
def to_json_string(self):
json_string = json.dumps(vars(self))
return json_string
def to_dict(self):
return vars(self)
@classmethod
def load_from_file(cls, config_path):
with open(config_path, 'r') as json_file:
config_dict = json.load(json_file)
return cls(**config_dict)
opensora/models/ae/videobase/modeling_videobase.py
+6 -2
View File
@@ -10,8 +10,12 @@ class VideoBaseAE(nn.Module):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def download_and_load_model(self):
return None
@classmethod
def load_from_checkpoint(cls, model_path):
pass
@classmethod
def download_and_load_model(cls, model_name, cache_dir=None):
pass
def encode(self, x, *args, **kwargs):
pass
opensora/models/ae/videobase/vqvae/configuration_vqvae.py
+29 -15
View File
@@ -1,17 +1,31 @@
from ..configuration_videobase import VideoBaseConfiguration
from dataclasses import dataclass, field
from typing import Tuple, Union
from typing import Union, Tuple
@dataclass
class VQVAEConfiguration:
embedding_dim: int = field(default=256)
n_codes: int = field(default=2048)
n_hiddens: int = field(default=240)
n_res_layers: int = field(default=4)
resolution: int = field(default=128)
sequence_length: int = field(default=16)
downsample: Union[str, Tuple[int, int, int]] = field(default=(4,4,4))
no_pos_embd: bool = field(default=True)
def __post_init__(self):
if isinstance(self.downsample, str):
self.downsample = tuple(map(int, self.downsample.split(',')))
class VQVAEConfiguration(VideoBaseConfiguration):
def __init__(
self,
embedding_dim: int = 256,
n_codes: int = 2048,
n_hiddens: int = 240,
n_res_layers: int = 4,
resolution: int = 128,
sequence_length: int = 16,
downsample: Union[Tuple[int, int, int], str] = (4, 4, 4),
no_pos_embd: bool = True,
**kwargs,
):
super().__init__(**kwargs)
self.embedding_dim = embedding_dim
self.n_codes = n_codes
self.n_hiddens = n_hiddens
self.n_res_layers = n_res_layers
self.resolution = resolution
self.sequence_length = sequence_length
if isinstance(downsample, str):
self.downsample = tuple(map(int, downsample.split(",")))
else:
self.downsample = downsample
self.no_pos_embd = no_pos_embd
opensora/models/ae/videobase/vqvae/modeling_vqvae.py
+207 -133
View File
@@ -5,8 +5,11 @@ import numpy as np
import torch.distributed as dist
import torch.nn.functional as F
import math
import os
import json
from .configuration_vqvae import VQVAEConfiguration
# Copied from https://github.com/wilson1yan/VideoGPT
def view_range(x, i, j, shape):
shape = tuple(shape)
@@ -26,6 +29,7 @@ def view_range(x, i, j, shape):
target_shape = x_shape[:i] + shape + x_shape[j:]
return x.view(target_shape)
# Copied from https://github.com/wilson1yan/VideoGPT
def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True):
n_dims = len(x.shape)
@@ -49,36 +53,41 @@ def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True):
x = x.contiguous()
return x
# Copied from https://github.com/wilson1yan/VideoGPT
def scaled_dot_product_attention(q, k, v, mask=None, attn_dropout=0., training=True):
def scaled_dot_product_attention(q, k, v, mask=None, attn_dropout=0.0, training=True):
# Performs scaled dot-product attention over the second to last dimension dn
# (b, n_head, d1, ..., dn, d)
attn = torch.matmul(q, k.transpose(-1, -2))
attn = attn / np.sqrt(q.shape[-1])
if mask is not None:
attn = attn.masked_fill(mask == 0, float('-inf'))
attn = attn.masked_fill(mask == 0, float("-inf"))
attn_float = F.softmax(attn, dim=-1)
attn = attn_float.type_as(attn) # b x n_head x d1 x ... x dn x d
attn = attn_float.type_as(attn) # b x n_head x d1 x ... x dn x d
attn = F.dropout(attn, p=attn_dropout, training=training)
a = torch.matmul(attn, v) # b x n_head x d1 x ... x dn x d
a = torch.matmul(attn, v) # b x n_head x d1 x ... x dn x d
return a
# Copied from https://github.com/wilson1yan/VideoGPT
class AxialBlock(nn.Module):
def __init__(self, n_hiddens, n_head):
super().__init__()
kwargs = dict(shape=(0,) * 3, dim_q=n_hiddens,
dim_kv=n_hiddens, n_head=n_head,
n_layer=1, causal=False, attn_type='axial')
self.attn_w = MultiHeadAttention(attn_kwargs=dict(axial_dim=-2),
**kwargs)
self.attn_h = MultiHeadAttention(attn_kwargs=dict(axial_dim=-3),
**kwargs)
self.attn_t = MultiHeadAttention(attn_kwargs=dict(axial_dim=-4),
**kwargs)
kwargs = dict(
shape=(0,) * 3,
dim_q=n_hiddens,
dim_kv=n_hiddens,
n_head=n_head,
n_layer=1,
causal=False,
attn_type="axial",
)
self.attn_w = MultiHeadAttention(attn_kwargs=dict(axial_dim=-2), **kwargs)
self.attn_h = MultiHeadAttention(attn_kwargs=dict(axial_dim=-3), **kwargs)
self.attn_t = MultiHeadAttention(attn_kwargs=dict(axial_dim=-4), **kwargs)
def forward(self, x):
x = shift_dim(x, 1, -1)
@@ -86,6 +95,7 @@ class AxialBlock(nn.Module):
x = shift_dim(x, -1, 1)
return x
# Copied from https://github.com/wilson1yan/VideoGPT
class AttentionResidualBlock(nn.Module):
def __init__(self, n_hiddens):
@@ -99,19 +109,20 @@ class AttentionResidualBlock(nn.Module):
SamePadConv3d(n_hiddens // 2, n_hiddens, 1, bias=False),
nn.BatchNorm3d(n_hiddens),
nn.ReLU(),
AxialBlock(n_hiddens, 2)
AxialBlock(n_hiddens, 2),
)
def forward(self, x):
return x + self.block(x)
# Copied from https://github.com/wilson1yan/VideoGPT
class Codebook(nn.Module):
def __init__(self, n_codes, embedding_dim):
super().__init__()
self.register_buffer('embeddings', torch.randn(n_codes, embedding_dim))
self.register_buffer('N', torch.zeros(n_codes))
self.register_buffer('z_avg', self.embeddings.data.clone())
self.register_buffer("embeddings", torch.randn(n_codes, embedding_dim))
self.register_buffer("N", torch.zeros(n_codes))
self.register_buffer("z_avg", self.embeddings.data.clone())
self.n_codes = n_codes
self.embedding_dim = embedding_dim
@@ -133,7 +144,7 @@ class Codebook(nn.Module):
y = self._tile(flat_inputs)
d = y.shape[0]
_k_rand = y[torch.randperm(y.shape[0])][:self.n_codes]
_k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
if dist.is_initialized():
dist.broadcast(_k_rand, 0)
self.embeddings.data.copy_(_k_rand)
@@ -145,9 +156,11 @@ class Codebook(nn.Module):
if self._need_init and self.training:
self._init_embeddings(z)
flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2)
distances = (flat_inputs ** 2).sum(dim=1, keepdim=True) \
- 2 * flat_inputs @ self.embeddings.t() \
+ (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)
distances = (
(flat_inputs**2).sum(dim=1, keepdim=True)
- 2 * flat_inputs @ self.embeddings.t()
+ (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)
)
encoding_indices = torch.argmin(distances, dim=1)
encode_onehot = F.one_hot(encoding_indices, self.n_codes).type_as(flat_inputs)
@@ -175,7 +188,7 @@ class Codebook(nn.Module):
self.embeddings.data.copy_(encode_normalized)
y = self._tile(flat_inputs)
_k_rand = y[torch.randperm(y.shape[0])][:self.n_codes]
_k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
if dist.is_initialized():
dist.broadcast(_k_rand, 0)
@@ -187,13 +200,18 @@ class Codebook(nn.Module):
avg_probs = torch.mean(encode_onehot, dim=0)
perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
return dict(embeddings=embeddings_st, encodings=encoding_indices,
commitment_loss=commitment_loss, perplexity=perplexity)
return dict(
embeddings=embeddings_st,
encodings=encoding_indices,
commitment_loss=commitment_loss,
perplexity=perplexity,
)
def dictionary_lookup(self, encodings):
embeddings = F.embedding(encodings, self.embeddings)
return embeddings
# Copied from https://github.com/wilson1yan/VideoGPT
class Encoder(nn.Module):
def __init__(self, n_hiddens, n_res_layers, downsample):
@@ -210,10 +228,9 @@ class Encoder(nn.Module):
self.conv_last = SamePadConv3d(in_channels, n_hiddens, kernel_size=3)
self.res_stack = nn.Sequential(
*[AttentionResidualBlock(n_hiddens)
for _ in range(n_res_layers)],
*[AttentionResidualBlock(n_hiddens) for _ in range(n_res_layers)],
nn.BatchNorm3d(n_hiddens),
nn.ReLU()
nn.ReLU(),
)
def forward(self, x):
@@ -224,10 +241,12 @@ class Encoder(nn.Module):
h = self.res_stack(h)
return h
# Copied from https://github.com/wilson1yan/VideoGPT
class MultiHeadAttention(nn.Module):
def __init__(self, shape, dim_q, dim_kv, n_head, n_layer,
causal, attn_type, attn_kwargs):
def __init__(
self, shape, dim_q, dim_kv, n_head, n_layer, causal, attn_type, attn_kwargs
):
super().__init__()
self.causal = causal
self.shape = shape
@@ -236,30 +255,30 @@ class MultiHeadAttention(nn.Module):
self.d_v = dim_kv // n_head
self.n_head = n_head
self.w_qs = nn.Linear(dim_q, n_head * self.d_k, bias=False) # q
self.w_qs = nn.Linear(dim_q, n_head * self.d_k, bias=False) # q
self.w_qs.weight.data.normal_(std=1.0 / np.sqrt(dim_q))
self.w_ks = nn.Linear(dim_kv, n_head * self.d_k, bias=False) # k
self.w_ks = nn.Linear(dim_kv, n_head * self.d_k, bias=False) # k
self.w_ks.weight.data.normal_(std=1.0 / np.sqrt(dim_kv))
self.w_vs = nn.Linear(dim_kv, n_head * self.d_v, bias=False) # v
self.w_vs = nn.Linear(dim_kv, n_head * self.d_v, bias=False) # v
self.w_vs.weight.data.normal_(std=1.0 / np.sqrt(dim_kv))
self.fc = nn.Linear(n_head * self.d_v, dim_q, bias=True) # c
self.fc = nn.Linear(n_head * self.d_v, dim_q, bias=True) # c
self.fc.weight.data.normal_(std=1.0 / np.sqrt(dim_q * n_layer))
if attn_type == 'full':
if attn_type == "full":
self.attn = FullAttention(shape, causal, **attn_kwargs)
elif attn_type == 'axial':
assert not causal, 'causal axial attention is not supported'
elif attn_type == "axial":
assert not causal, "causal axial attention is not supported"
self.attn = AxialAttention(len(shape), **attn_kwargs)
elif attn_type == 'sparse':
elif attn_type == "sparse":
self.attn = SparseAttention(shape, n_head, causal, **attn_kwargs)
self.cache = None
def forward(self, q, k, v, decode_step=None, decode_idx=None):
""" Compute multi-head attention
"""Compute multi-head attention
Args
q, k, v: a [b, d1, ..., dn, c] tensor or
a [b, 1, ..., 1, c] tensor if decode_step is not None
@@ -286,34 +305,40 @@ class MultiHeadAttention(nn.Module):
if self.causal:
k_shape = (q.shape[0], n_head, *self.shape, self.d_k)
v_shape = (q.shape[0], n_head, *self.shape, self.d_v)
self.cache = dict(k=torch.zeros(k_shape, dtype=k.dtype, device=q.device),
v=torch.zeros(v_shape, dtype=v.dtype, device=q.device))
self.cache = dict(
k=torch.zeros(k_shape, dtype=k.dtype, device=q.device),
v=torch.zeros(v_shape, dtype=v.dtype, device=q.device),
)
else:
# cache only once in the non-causal case
self.cache = dict(k=k.clone(), v=v.clone())
if self.causal:
idx = (slice(None, None), slice(None, None), *[slice(i, i+ 1) for i in decode_idx])
self.cache['k'][idx] = k
self.cache['v'][idx] = v
k, v = self.cache['k'], self.cache['v']
idx = (
slice(None, None),
slice(None, None),
*[slice(i, i + 1) for i in decode_idx],
)
self.cache["k"][idx] = k
self.cache["v"][idx] = v
k, v = self.cache["k"], self.cache["v"]
a = self.attn(q, k, v, decode_step, decode_idx)
# (b, *d_shape, n_head, d) -> (b, *d_shape, n_head * d)
a = shift_dim(a, 1, -2).flatten(start_dim=-2)
a = self.fc(a) # (b x seq_len x embd_dim)
a = self.fc(a) # (b x seq_len x embd_dim)
return a
# Copied from https://github.com/wilson1yan/VideoGPT
class Decoder(nn.Module):
def __init__(self, n_hiddens, n_res_layers, upsample):
super().__init__()
self.res_stack = nn.Sequential(
*[AttentionResidualBlock(n_hiddens)
for _ in range(n_res_layers)],
*[AttentionResidualBlock(n_hiddens) for _ in range(n_res_layers)],
nn.BatchNorm3d(n_hiddens),
nn.ReLU()
nn.ReLU(),
)
n_times_upsample = np.array([int(math.log2(d)) for d in upsample])
@@ -322,8 +347,7 @@ class Decoder(nn.Module):
for i in range(max_us):
out_channels = 3 if i == max_us - 1 else n_hiddens
us = tuple([2 if d > 0 else 1 for d in n_times_upsample])
convt = SamePadConvTranspose3d(n_hiddens, out_channels, 4,
stride=us)
convt = SamePadConvTranspose3d(n_hiddens, out_channels, 4, stride=us)
self.convts.append(convt)
n_times_upsample -= 1
@@ -335,6 +359,7 @@ class Decoder(nn.Module):
h = F.relu(h)
return h
# Copied from https://github.com/wilson1yan/VideoGPT
class SamePadConv3d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True):
@@ -347,17 +372,19 @@ class SamePadConv3d(nn.Module):
# assumes that the input shape is divisible by stride
total_pad = tuple([k - s for k, s in zip(kernel_size, stride)])
pad_input = []
for p in total_pad[::-1]: # reverse since F.pad starts from last dim
for p in total_pad[::-1]: # reverse since F.pad starts from last dim
pad_input.append((p // 2 + p % 2, p // 2))
pad_input = sum(pad_input, tuple())
self.pad_input = pad_input
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size,
stride=stride, padding=0, bias=bias)
self.conv = nn.Conv3d(
in_channels, out_channels, kernel_size, stride=stride, padding=0, bias=bias
)
def forward(self, x):
return self.conv(F.pad(x, self.pad_input))
# Copied from https://github.com/wilson1yan/VideoGPT
class SamePadConvTranspose3d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True):
@@ -369,18 +396,24 @@ class SamePadConvTranspose3d(nn.Module):
total_pad = tuple([k - s for k, s in zip(kernel_size, stride)])
pad_input = []
for p in total_pad[::-1]: # reverse since F.pad starts from last dim
for p in total_pad[::-1]: # reverse since F.pad starts from last dim
pad_input.append((p // 2 + p % 2, p // 2))
pad_input = sum(pad_input, tuple())
self.pad_input = pad_input
self.convt = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
stride=stride, bias=bias,
padding=tuple([k - 1 for k in kernel_size]))
self.convt = nn.ConvTranspose3d(
in_channels,
out_channels,
kernel_size,
stride=stride,
bias=bias,
padding=tuple([k - 1 for k in kernel_size]),
)
def forward(self, x):
return self.convt(F.pad(x, self.pad_input))
# Copied from https://github.com/wilson1yan/VideoGPT
class FullAttention(nn.Module):
def __init__(self, shape, causal, attn_dropout):
@@ -390,7 +423,7 @@ class FullAttention(nn.Module):
seq_len = np.prod(shape)
if self.causal:
self.register_buffer('mask', torch.tril(torch.ones(seq_len, seq_len)))
self.register_buffer("mask", torch.tril(torch.ones(seq_len, seq_len)))
def forward(self, q, k, v, decode_step, decode_idx):
mask = self.mask if self.causal else None
@@ -402,12 +435,13 @@ class FullAttention(nn.Module):
k = k.flatten(start_dim=2, end_dim=-2)
v = v.flatten(start_dim=2, end_dim=-2)
out = scaled_dot_product_attention(q, k, v, mask=mask,
attn_dropout=self.attn_dropout,
training=self.training)
out = scaled_dot_product_attention(
q, k, v, mask=mask, attn_dropout=self.attn_dropout, training=self.training
)
return view_range(out, 2, 3, old_shape)
# Copied from https://github.com/wilson1yan/VideoGPT
class AxialAttention(nn.Module):
def __init__(self, n_dim, axial_dim):
@@ -415,7 +449,7 @@ class AxialAttention(nn.Module):
if axial_dim < 0:
axial_dim = 2 + n_dim + 1 + axial_dim
else:
axial_dim += 2 # account for batch, head, dim
axial_dim += 2 # account for batch, head, dim
self.axial_dim = axial_dim
def forward(self, q, k, v, decode_step, decode_idx):
@@ -430,12 +464,14 @@ class AxialAttention(nn.Module):
out = shift_dim(out, -2, self.axial_dim)
return out
# Copied from https://github.com/wilson1yan/VideoGPT
class StridedSparsityConfig(object):
"""
Strided Sparse configuration specified in https://arxiv.org/abs/1904.10509 that
generalizes to arbitrary dimensions
"""
def __init__(self, shape, n_head, causal, block, num_local_blocks):
self.n_head = n_head
self.shape = shape
@@ -443,8 +479,8 @@ class StridedSparsityConfig(object):
self.block = block
self.num_local_blocks = num_local_blocks
assert self.num_local_blocks >= 1, 'Must have at least 1 local block'
assert self.seq_len % self.block == 0, 'seq len must be divisible by block size'
assert self.num_local_blocks >= 1, "Must have at least 1 local block"
assert self.seq_len % self.block == 0, "seq len must be divisible by block size"
self._block_shape = self._compute_block_shape()
self._block_shape_cum = self._block_shape_cum_sizes()
@@ -462,8 +498,8 @@ class StridedSparsityConfig(object):
for row in range(0, num_blocks):
end = min(row + self.num_local_blocks, num_blocks)
for col in range(
max(0, row - self.num_local_blocks),
(row + 1 if self.causal else end)):
max(0, row - self.num_local_blocks), (row + 1 if self.causal else end)
):
layout[:, row, col] = 1
return layout
@@ -487,7 +523,9 @@ class StridedSparsityConfig(object):
return layout
def make_layout(self):
layout = torch.zeros((self.n_head, self.num_blocks, self.num_blocks), dtype=torch.int64)
layout = torch.zeros(
(self.n_head, self.num_blocks, self.num_blocks), dtype=torch.int64
)
layout = self.set_local_layout(layout)
layout = self.set_global_layout(layout)
return layout
@@ -505,7 +543,7 @@ class StridedSparsityConfig(object):
elem = block_layout[h, i, j].item()
if elem == 1:
assert i >= j
if i == j: # need to mask within block on diagonals
if i == j: # need to mask within block on diagonals
attn_mask[counter] = torch.tril(attn_mask[counter])
counter += 1
assert counter == num_dense_blocks
@@ -514,9 +552,9 @@ class StridedSparsityConfig(object):
def get_non_block_layout_row(self, block_layout, row):
block_row = row // self.block
block_row = block_layout[:, [block_row]] # n_head x 1 x n_blocks
block_row = block_layout[:, [block_row]] # n_head x 1 x n_blocks
block_row = block_row.repeat_interleave(self.block, dim=-1)
block_row[:, :, row + 1:] = 0.
block_row[:, :, row + 1 :] = 0.0
return block_row
############# Helper functions ##########################
@@ -540,7 +578,9 @@ class StridedSparsityConfig(object):
return tuple(np.flip(np.cumprod(bs)[:-1])) + (1,)
def _to_flattened_idx(self, idx):
assert len(idx) == len(self._block_shape), f"{len(idx)} != {len(self._block_shape)}"
assert len(idx) == len(
self._block_shape
), f"{len(idx)} != {len(self._block_shape)}"
flat_idx = 0
for i in range(len(self._block_shape)):
flat_idx += idx[i] * self._block_shape_cum[i]
@@ -554,60 +594,82 @@ class StridedSparsityConfig(object):
flat_idx %= self._block_shape_cum[i]
return tuple(idx)
# Copied from https://github.com/wilson1yan/VideoGPT
class SparseAttention(nn.Module):
ops = dict()
attn_mask = dict()
block_layout = dict()
def __init__(self, shape, n_head, causal, num_local_blocks=4, block=32,
attn_dropout=0.): # does not use attn_dropout
def __init__(
self, shape, n_head, causal, num_local_blocks=4, block=32, attn_dropout=0.0
): # does not use attn_dropout
super().__init__()
self.causal = causal
self.shape = shape
self.sparsity_config = StridedSparsityConfig(shape=shape, n_head=n_head,
causal=causal, block=block,
num_local_blocks=num_local_blocks)
self.sparsity_config = StridedSparsityConfig(
shape=shape,
n_head=n_head,
causal=causal,
block=block,
num_local_blocks=num_local_blocks,
)
if self.shape not in SparseAttention.block_layout:
SparseAttention.block_layout[self.shape] = self.sparsity_config.make_layout()
SparseAttention.block_layout[self.shape] = (
self.sparsity_config.make_layout()
)
if causal and self.shape not in SparseAttention.attn_mask:
SparseAttention.attn_mask[self.shape] = self.sparsity_config.make_sparse_attn_mask()
SparseAttention.attn_mask[self.shape] = (
self.sparsity_config.make_sparse_attn_mask()
)
def get_ops(self):
try:
from deepspeed.ops.sparse_attention import MatMul, Softmax
except:
raise Exception('Error importing deepspeed. Please install using `DS_BUILD_SPARSE_ATTN=1 pip install deepspeed`')
raise Exception(
"Error importing deepspeed. Please install using `DS_BUILD_SPARSE_ATTN=1 pip install deepspeed`"
)
if self.shape not in SparseAttention.ops:
sparsity_layout = self.sparsity_config.make_layout()
sparse_dot_sdd_nt = MatMul(sparsity_layout,
self.sparsity_config.block,
'sdd',
trans_a=False,
trans_b=True)
sparse_dot_sdd_nt = MatMul(
sparsity_layout,
self.sparsity_config.block,
"sdd",
trans_a=False,
trans_b=True,
)
sparse_dot_dsd_nn = MatMul(sparsity_layout,
self.sparsity_config.block,
'dsd',
trans_a=False,
trans_b=False)
sparse_dot_dsd_nn = MatMul(
sparsity_layout,
self.sparsity_config.block,
"dsd",
trans_a=False,
trans_b=False,
)
sparse_softmax = Softmax(sparsity_layout, self.sparsity_config.block)
SparseAttention.ops[self.shape] = (sparse_dot_sdd_nt,
sparse_dot_dsd_nn,
sparse_softmax)
SparseAttention.ops[self.shape] = (
sparse_dot_sdd_nt,
sparse_dot_dsd_nn,
sparse_softmax,
)
return SparseAttention.ops[self.shape]
def forward(self, q, k, v, decode_step, decode_idx):
if self.training and self.shape not in SparseAttention.ops:
self.get_ops()
SparseAttention.block_layout[self.shape] = SparseAttention.block_layout[self.shape].to(q)
SparseAttention.block_layout[self.shape] = SparseAttention.block_layout[
self.shape
].to(q)
if self.causal:
SparseAttention.attn_mask[self.shape] = SparseAttention.attn_mask[self.shape].to(q).type_as(q)
SparseAttention.attn_mask[self.shape] = (
SparseAttention.attn_mask[self.shape].to(q).type_as(q)
)
attn_mask = SparseAttention.attn_mask[self.shape] if self.causal else None
old_shape = q.shape[2:-1]
@@ -616,39 +678,43 @@ class SparseAttention(nn.Module):
v = v.flatten(start_dim=2, end_dim=-2)
if decode_step is not None:
mask = self.sparsity_config.get_non_block_layout_row(SparseAttention.block_layout[self.shape], decode_step)
out = scaled_dot_product_attention(q, k, v, mask=mask, training=self.training)
mask = self.sparsity_config.get_non_block_layout_row(
SparseAttention.block_layout[self.shape], decode_step
)
out = scaled_dot_product_attention(
q, k, v, mask=mask, training=self.training
)
else:
if q.shape != k.shape or k.shape != v.shape:
raise Exception('SparseAttention only support self-attention')
raise Exception("SparseAttention only support self-attention")
sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax = self.get_ops()
scaling = float(q.shape[-1]) ** -0.5
attn_output_weights = sparse_dot_sdd_nt(q, k)
if attn_mask is not None:
attn_output_weights = attn_output_weights.masked_fill(attn_mask == 0,
float('-inf'))
attn_output_weights = sparse_softmax(
attn_output_weights,
scale=scaling
)
attn_output_weights = attn_output_weights.masked_fill(
attn_mask == 0, float("-inf")
)
attn_output_weights = sparse_softmax(attn_output_weights, scale=scaling)
out = sparse_dot_dsd_nn(attn_output_weights, v)
return view_range(out, 2, 3, old_shape)
# Modified from https://github.com/wilson1yan/VideoGPT
class VQVAEModel(VideoBaseAE):
DOWNLOADED_VQVAE = {
'bair_stride4x2x2': '1iIAYJ2Qqrx5Q94s5eIXQYJgAydzvT_8L',
'ucf101_stride4x4x4': '1uuB_8WzHP_bbBmfuaIV7PK_Itl3DyHY5',
'kinetics_stride4x4x4': '1DOvOZnFAIQmux6hG7pN_HkyJZy3lXbCB',
'kinetics_stride2x4x4': '1jvtjjtrtE4cy6pl7DK_zWFEPY3RZt2pB'
"bair_stride4x2x2": "1iIAYJ2Qqrx5Q94s5eIXQYJgAydzvT_8L",
"ucf101_stride4x4x4": "1uuB_8WzHP_bbBmfuaIV7PK_Itl3DyHY5",
"kinetics_stride4x4x4": "1DOvOZnFAIQmux6hG7pN_HkyJZy3lXbCB",
"kinetics_stride2x4x4": "1jvtjjtrtE4cy6pl7DK_zWFEPY3RZt2pB",
}
def __init__(self, config: VQVAEConfiguration):
super().__init__()
self.config = config
self.embedding_dim = config.embedding_dim
self.n_codes = config.n_codes
self.encoder = Encoder(config.n_hiddens, config.n_res_layers, config.downsample)
@@ -656,45 +722,53 @@ class VQVAEModel(VideoBaseAE):
self.pre_vq_conv = SamePadConv3d(config.n_hiddens, config.embedding_dim, 1)
self.post_vq_conv = SamePadConv3d(config.embedding_dim, config.n_hiddens, 1)
self.codebook = Codebook(config.n_codes, config.embedding_dim)
def forward(self, x):
z = self.pre_vq_conv(self.encoder(x))
vq_output = self.codebook(z)
x_recon = self.decoder(self.post_vq_conv(vq_output['embeddings']))
x_recon = self.decoder(self.post_vq_conv(vq_output["embeddings"]))
recon_loss = F.mse_loss(x_recon, x) / 0.06
return recon_loss, x_recon, vq_output
def encode(self, x, include_embeddings=False):
h = self.pre_vq_conv(self.encoder(x))
vq_output = self.codebook(h)
if include_embeddings:
return vq_output['encodings'], vq_output['embeddings']
return vq_output["encodings"], vq_output["embeddings"]
else:
return vq_output['encodings']
return vq_output["encodings"]
def decode(self, encodings):
h = F.embedding(encodings, self.codebook.embeddings)
h = self.post_vq_conv(shift_dim(h, -1, 1))
return self.decoder(h)
@staticmethod
def load_from_checkpoint(model_path):
model_cpkt = torch.load(model_path)
# Compatible with old videogpt model formats.
if 'hyper_parameters' in model_cpkt:
hyper_parameters = vars(model_cpkt.get("hyper_parameters").get("args"))
print(hyper_parameters)
state_dict = model_cpkt.get("state_dict")
model = VQVAEModel(config=VQVAEConfiguration(**hyper_parameters))
@classmethod
def load_from_checkpoint(cls, model_path):
if not os.path.isdir(model_path):
"""model downloaded from internet"""
model_cpkt = torch.load(model_path)
# Compatible with old videogpt model formats.
if "hyper_parameters" in model_cpkt:
hyper_parameters = vars(model_cpkt.get("hyper_parameters").get("args"))
state_dict = model_cpkt.get("state_dict")
model = cls(config=VQVAEConfiguration(**hyper_parameters))
model.load_state_dict(state_dict)
return model
else:
raise RuntimeError("Model checkpoint has a wrong format.")
else:
with open(os.path.join(model_path, "config.json"), "r") as file:
config = json.load(file)
state_dict = torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location="cpu")
model = cls(config=VQVAEConfiguration(**config))
model.load_state_dict(state_dict)
return model
else:
raise RuntimeError("Model checkpoint has a wrong format.")
@staticmethod
def download_and_load_model(model_name, cache_dir=None):
@classmethod
def download_and_load_model(cls, model_name, cache_dir=None):
from .....utils.downloader import gdown_download
path = gdown_download(VQVAEModel.DOWNLOADED_VQVAE[model_name], model_name, cache_dir=cache_dir)
return VQVAEModel.load_from_checkpoint(path)
path = gdown_download(
cls.DOWNLOADED_VQVAE[model_name], model_name, cache_dir=cache_dir
)
return cls.load_from_checkpoint(path)
opensora/models/ae/videobase/vqvae/trainer_vqvae.py
+18 -4
View File
@@ -1,15 +1,29 @@
from torch import nn
from transformers import Trainer
import torch.nn.functional as F
from typing import Optional
import os
import torch
from transformers.utils import WEIGHTS_NAME
import json
class VQVAETrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
model = model.module
x = inputs.get("video")
z = model.pre_vq_conv(model.encoder(x))
vq_output = model.codebook(z)
x_recon = model.decoder(model.post_vq_conv(vq_output['embeddings']))
x_recon = model.decoder(model.post_vq_conv(vq_output["embeddings"]))
recon_loss = F.mse_loss(x_recon, x) / 0.06
return recon_loss
return recon_loss
def _save(self, output_dir: Optional[str] = None, state_dict=None):
output_dir = output_dir if output_dir is not None else self.args.output_dir
os.makedirs(output_dir, exist_ok=True)
if state_dict is None:
state_dict = self.model.state_dict()
with open(os.path.join(output_dir, "config.json"), "w") as file:
json.dump(self.model.config.to_dict(), file)
torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
train_scripts/train_videogpt.py → opensora/train/train_videogpt.py
+14 -3
View File
@@ -1,5 +1,6 @@
import sys
sys.path.append(".")
from opensora.models.ae.videobase import (
VQVAEModel,
VQVAEConfiguration,
@@ -10,11 +11,19 @@ import argparse
from typing import Optional
from accelerate.utils import set_seed
from transformers import HfArgumentParser, TrainingArguments
from dataclasses import dataclass, field
from dataclasses import dataclass, field, asdict
@dataclass
class VQVAEArgument(VQVAEConfiguration):
class VQVAEArgument:
embedding_dim: int = field(default=256),
n_codes: int = field(default=2048),
n_hiddens: int = field(default=240),
n_res_layers: int = field(default=4),
resolution: int = field(default=128),
sequence_length: int = field(default=16),
downsample: str = field(default="4,4,4"),
no_pos_embd: bool = True,
data_path: str = field(default=None, metadata={"help": "data path"})
@dataclass
@@ -24,8 +33,10 @@ class VQVAETrainingArgument(TrainingArguments):
)
def train(args, vqvae_args, training_args):
# Load Config
config = VQVAEConfiguration(**asdict(vqvae_args))
# Load Model
model = VQVAEModel(vqvae_args)
model = VQVAEModel(config)
# Load Dataset
dataset = VQVAEDataset(args.data_path, sequence_length=args.sequence_length)
scripts/train_videogpt.sh
+2 -2
View File
@@ -1,6 +1,6 @@
accelerate launch \
--config_file scripts/accelerate_configs/default_config.yaml \
train_scripts/train_videogpt.py \
opensora/train/train_videogpt.py \
--do_train \
--seed 1234 \
--data_path UCF-101 \
@@ -15,7 +15,7 @@ accelerate launch \
--save_strategy steps \
--save_total_limit 5 \
--logging_steps 5 \
--save_steps 500 \
--save_steps 5 \
--n_codes 2048 \
--n_hiddens 240 \
--n_res_layers 4 \
train_scripts/train.py
-294
View File
@@ -1,294 +0,0 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""
A minimal training script for DiT using PyTorch DDP.
"""
import torch
# the first flag below was False when we tested this script but True makes A100 training a lot faster:
from torch import nn
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from collections import OrderedDict
from copy import deepcopy
from glob import glob
from time import time
import argparse
import logging
import os
from opensora.dataset import getdataset
from opensora.models.ae import getae
from opensora.utils.dataset_utils import Collate
from opensora.models.diffusion.dit.models import DiT_models as Diffusion_models
from opensora.models.diffusion.dit.diffusion import create_diffusion
#################################################################################
# Training Helper Functions #
#################################################################################
@torch.no_grad()
def update_ema(ema_model, model, decay=0.9999):
"""
Step the EMA model towards the current model.
"""
ema_params = OrderedDict(ema_model.named_parameters())
model_params = OrderedDict(model.named_parameters())
for name, param in model_params.items():
# TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
def requires_grad(model, flag=True):
"""
Set requires_grad flag for all parameters in a model.
"""
for p in model.parameters():
p.requires_grad = flag
def cleanup():
"""
End DDP training.
"""
dist.destroy_process_group()
def create_logger(logging_dir):
"""
Create a logger that writes to a log file and stdout.
"""
if dist.get_rank() == 0: # real logger
logging.basicConfig(
level=logging.INFO,
format='[\033[34m%(asctime)s\033[0m] %(message)s',
datefmt='%Y-%m-%d %H:%M:%S',
handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
)
logger = logging.getLogger(__name__)
else: # dummy logger (does nothing)
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())
return logger
#################################################################################
# Training Loop #
#################################################################################
def main(args):
"""
Trains a new DiT model.
"""
assert torch.cuda.is_available(), "Training currently requires at least one GPU."
# Setup DDP:
dist.init_process_group("nccl")
assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size."
rank = dist.get_rank()
device = rank % torch.cuda.device_count()
seed = args.global_seed * dist.get_world_size() + rank
torch.manual_seed(seed)
torch.cuda.set_device(device)
print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
# Setup an experiment folder:
if rank == 0:
print(args)
os.makedirs(args.results_dir, exist_ok=True) # Make results folder (holds all experiment subfolders)
experiment_index = len(glob(f"{args.results_dir}/*"))
model_string_name = args.model.replace("/", "-") # e.g., DiT-XL/2 --> DiT-XL-2 (for naming folders)
experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" # Create an experiment folder
checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints
os.makedirs(checkpoint_dir, exist_ok=True)
logger = create_logger(experiment_dir)
logger.info(f"Experiment directory created at {experiment_dir}")
else:
logger = create_logger(None)
# Create model:
vae_stride_t, vae_stride_h, vae_stride_w = [int(i) for i in args.ae[-5:].split('x')]
args.vae_stride_t, args.vae_stride_h, args.vae_stride_w = vae_stride_t, vae_stride_h, vae_stride_w
args.vae_stride = args.vae_stride_h
patch_size = args.model[-3:]
patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2])
args.patch_size = patch_size_h
args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w
assert args.num_frames % vae_stride_t == 0, "Image size must be divisible by 8 (for the VAE encoder)."
assert args.max_image_size % vae_stride_h == 0, "Image size must be divisible by 8 (for the VAE encoder)."
assert vae_stride_h == vae_stride_w, "Support now."
assert patch_size_h == patch_size_w, "Support now."
latent_size = (args.num_frames // vae_stride_t, args.max_image_size // vae_stride_h, args.max_image_size // vae_stride_w)
model = Diffusion_models[args.model](
input_size=latent_size,
num_classes=args.num_classes
)
model.gradient_checkpointing = args.gradient_checkpointing
if args.pt_ckpt is not None:
state_dict = torch.load(args.pt_ckpt, map_location='cpu')
if state_dict.get('model', None) is not None:
state_dict = state_dict['model']
del state_dict['x_embedder.proj.weight']
del state_dict['final_layer.linear.weight']
del state_dict['final_layer.linear.bias']
del state_dict['pos_embed']
del state_dict['y_embedder.embedding_table.weight']
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print('missing_keys:', missing_keys, 'unexpected_keys:', unexpected_keys)
assert len(missing_keys) == 4
logger.info(f"{missing_keys}, {unexpected_keys}")
# Note that parameter initialization is done within the DiT constructor
ema = deepcopy(model).to(device) # Create an EMA of the model for use after training
requires_grad(ema, False)
model = DDP(model.to(device), device_ids=[rank])
diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule
ae = getae(args).to(device)
logger.info(f"{model}")
logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
for n, p in model.named_parameters():
if p.requires_grad:
logger.info(f"Training Parameters: {n}")
# Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0)
# Setup data:
dataset = getdataset(args)
sampler = DistributedSampler(
dataset,
num_replicas=dist.get_world_size(),
rank=rank,
shuffle=True,
seed=args.global_seed
)
loader = DataLoader(
dataset,
batch_size=int(args.global_batch_size // dist.get_world_size()),
shuffle=False,
sampler=sampler,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True,
collate_fn=Collate(args)
)
logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})")
# Prepare models for training:
update_ema(ema, model.module, decay=0) # Ensure EMA is initialized with synced weights
model.train() # important! This enables embedding dropout for classifier-free guidance
ema.eval() # EMA model should always be in eval mode
# Variables for monitoring/logging purposes:
train_steps = 0
log_steps = 0
running_loss = 0
start_time = time()
logger.info(f"Training for {args.epochs} epochs...")
for epoch in range(args.epochs):
sampler.set_epoch(epoch)
logger.info(f"Beginning epoch {epoch}...")
for x, y, attn_mask in loader:
x = x.to(device)
y = y.to(device)
attn_mask = attn_mask.to(device)
with torch.no_grad():
x = ae.encode(x) # Map input images to latent space + normalize latents
t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
model_kwargs = dict(y=y, attention_mask=attn_mask)
loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
loss = loss_dict["loss"].mean()
opt.zero_grad()
loss.backward()
if args.clip_grad_norm is not None:
nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm)
opt.step()
update_ema(ema, model.module)
# Log loss values:
running_loss += loss.item()
log_steps += 1
train_steps += 1
if train_steps % args.log_every == 0:
# Measure training speed:
torch.cuda.synchronize()
end_time = time()
steps_per_sec = log_steps / (end_time - start_time)
# Reduce loss history over all processes:
avg_loss = torch.tensor(running_loss / log_steps, device=device)
dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
avg_loss = avg_loss.item() / dist.get_world_size()
logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
# Reset monitoring variables:
running_loss = 0
log_steps = 0
start_time = time()
# Save DiT checkpoint:
if train_steps % args.ckpt_every == 0 and train_steps > 0:
if rank == 0:
checkpoint = {
"model": model.module.state_dict(),
"ema": ema.state_dict(),
"opt": opt.state_dict(),
"args": args
}
checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
torch.save(checkpoint, checkpoint_path)
logger.info(f"Saved checkpoint to {checkpoint_path}")
dist.barrier()
model.eval() # important! This disables randomized embedding dropout
# do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
logger.info("Done!")
cleanup()
if __name__ == "__main__":
# Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters).
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", type=str, required=True)
parser.add_argument("--data-path", type=str, required=True)
parser.add_argument("--results-dir", type=str, default="results")
parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="DiT-XL/122")
parser.add_argument("--num-classes", type=int, default=1000)
parser.add_argument("--epochs", type=int, default=1400)
parser.add_argument("--global-batch-size", type=int, default=256)
parser.add_argument("--global-seed", type=int, default=0)
parser.add_argument("--num-workers", type=int, default=8)
parser.add_argument("--log-every", type=int, default=100)
parser.add_argument("--ckpt-every", type=int, default=50_000)
# --------------------------------------
parser.add_argument("--ae", type=str, choices=['bair_stride4x2x2', 'ucf101_stride4x4x4',
'kinetics_stride4x4x4', 'kinetics_stride2x4x4'],
default="ucf101_stride4x4x4")
parser.add_argument("--pt-ckpt", type=str, default=None)
parser.add_argument("--sample-rate", type=int, default=4)
parser.add_argument("--num-frames", type=int, default=16)
parser.add_argument("--max-image-size", type=int, default=128)
parser.add_argument("--dynamic-frames", action="store_true")
parser.add_argument("--gradient-checkpointing", action="store_true")
parser.add_argument("--lr", type=float, default=1e-4)
parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)")
# --------------------------------------
args = parser.parse_args()
main(args)