nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-05-21 11:59:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

[feat] update logging and bucket

Browse source
This commit is contained in:
zhengzangw 2024-05-09 08:07:56 +00:00
parent 89a90e6b3d
commit 2a55d672a7
19 changed files with 376 additions and 306 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
configs/opensora-v1-1/train/stage1.py
View file

@ -16,15 +16,15 @@ bucket_config = { # 1s/it
"1024": {1: (0.3, 3)},
}
mask_ratios = {
"mask_no": 0.75,
"mask_quarter_random": 0.025,
"mask_quarter_head": 0.025,
"mask_quarter_tail": 0.025,
"mask_quarter_head_tail": 0.05,
"mask_image_random": 0.025,
"mask_image_head": 0.025,
"mask_image_tail": 0.025,
"mask_image_head_tail": 0.05,
"identity": 0.75,
"quarter_random": 0.025,
"quarter_head": 0.025,
"quarter_tail": 0.025,
"quarter_head_tail": 0.05,
"image_random": 0.025,
"image_head": 0.025,
"image_tail": 0.025,
"image_head_tail": 0.05,
}
# Define acceleration

18
configs/opensora-v1-1/train/stage2.py
View file

@ -18,15 +18,15 @@ bucket_config = { # 7s/it
"1080p": {1: (0.4, 8)},
}
mask_ratios = {
"mask_no": 0.75,
"mask_quarter_random": 0.025,
"mask_quarter_head": 0.025,
"mask_quarter_tail": 0.025,
"mask_quarter_head_tail": 0.05,
"mask_image_random": 0.025,
"mask_image_head": 0.025,
"mask_image_tail": 0.025,
"mask_image_head_tail": 0.05,
"identity": 0.75,
"quarter_random": 0.025,
"quarter_head": 0.025,
"quarter_tail": 0.025,
"quarter_head_tail": 0.05,
"image_random": 0.025,
"image_head": 0.025,
"image_tail": 0.025,
"image_head_tail": 0.05,
}
# Define acceleration

18
configs/opensora-v1-1/train/stage3.py
View file

@ -18,15 +18,15 @@ bucket_config = { # 13s/it
"1024": {1: (0.3, 40)},
}
mask_ratios = {
"mask_no": 0.75,
"mask_quarter_random": 0.025,
"mask_quarter_head": 0.025,
"mask_quarter_tail": 0.025,
"mask_quarter_head_tail": 0.05,
"mask_image_random": 0.025,
"mask_image_head": 0.025,
"mask_image_tail": 0.025,
"mask_image_head_tail": 0.05,
"identity": 0.75,
"quarter_random": 0.025,
"quarter_head": 0.025,
"quarter_tail": 0.025,
"quarter_head_tail": 0.05,
"image_random": 0.025,
"image_head": 0.025,
"image_tail": 0.025,
"image_head_tail": 0.05,
}
# Define acceleration

50
configs/opensora-v1-2/train/stage1-gc.py
View file

@ -4,13 +4,34 @@ dataset = dict(
transform_name="resize_crop",
frame_interval=1,
)
# h800
# bucket_config = { # 20s/it
# "144p": {1: (1.0, 100), 51: (1.0, 30), 102: (1.0, 20), 204: (1.0, 8), 408: (1.0, 4)},
# # ---
# "256": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
# "240p": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
# # ---
# "360p": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
# "512": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
# # ---
# "480p": {1: (0.5, 40), 51: (0.3, 6), 102: (0.3, 3), 204: (0.3, 1), 408: (0.0, None)},
# # ---
# "720p": {1: (0.2, 20), 51: (0.3, 2), 102: (0.3, 1), 204: (0.0, None)},
# "1024": {1: (0.1, 20), 51: (0.3, 2), 102: (0.3, 1), 204: (0.0, None)},
# # ---
# "1080p": {1: (0.1, 10)},
# # ---
# "2048": {1: (0.1, 5)},
# }
# webvid
bucket_config = { # 20s/it
"144p": {1: (1.0, 100), 51: (1.0, 30), 102: (1.0, 20), 204: (1.0, 8), 408: (1.0, 4)},
"144p": {1: (1.0, 100), 51: (1.0, 30), 102: (1.0, 20), 204: ((1.0, 0.2), 8), 408: ((1.0, 0.1), 4)},
# ---
"256": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
"240p": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
"256": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: ((0.3, 0.2), 4), 408: ((0.3, 0.1), 2)},
"240p": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: ((0.3, 0.2), 4), 408: ((0.3, 0.1), 2)},
# ---
"360p": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
"360p": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: ((0.3, 0.2), 2), 408: ((0.3, 0.1), 1)},
"512": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
# ---
"480p": {1: (0.5, 40), 51: (0.3, 6), 102: (0.3, 3), 204: (0.3, 1), 408: (0.0, None)},
@ -22,6 +43,7 @@ bucket_config = { # 20s/it
# ---
"2048": {1: (0.1, 5)},
}
grad_checkpoint = True
batch_size = None
@ -71,16 +93,16 @@ scheduler = dict(
# Mask settings
mask_ratios = {
"mask_random": 0.2,
"mask_intepolate": 0.01,
"mask_quarter_random": 0.01,
"mask_quarter_head": 0.01,
"mask_quarter_tail": 0.01,
"mask_quarter_head_tail": 0.01,
"mask_image_random": 0.05,
"mask_image_head": 0.1,
"mask_image_tail": 0.05,
"mask_image_head_tail": 0.05,
"random": 0.2,
"intepolate": 0.01,
"quarter_random": 0.01,
"quarter_head": 0.01,
"quarter_tail": 0.01,
"quarter_head_tail": 0.01,
"image_random": 0.05,
"image_head": 0.1,
"image_tail": 0.05,
"image_head_tail": 0.05,
}
# Log settings

20
configs/opensora-v1-2/train/stage1.py
View file

@ -69,16 +69,16 @@ scheduler = dict(
# Mask settings
mask_ratios = {
"mask_random": 0.2,
"mask_intepolate": 0.01,
"mask_quarter_random": 0.01,
"mask_quarter_head": 0.01,
"mask_quarter_tail": 0.01,
"mask_quarter_head_tail": 0.01,
"mask_image_random": 0.05,
"mask_image_head": 0.1,
"mask_image_tail": 0.05,
"mask_image_head_tail": 0.05,
"random": 0.2,
"intepolate": 0.01,
"quarter_random": 0.01,
"quarter_head": 0.01,
"quarter_tail": 0.01,
"quarter_head_tail": 0.01,
"image_random": 0.05,
"image_head": 0.1,
"image_tail": 0.05,
"image_head_tail": 0.05,
}
# Log settings

4
configs/opensora/train/16x256x256-mask.py
View file

@ -24,8 +24,8 @@ model = dict(
enable_layernorm_kernel=True,
)
mask_ratios = {
"mask_no": 0.7,
"mask_random": 0.15,
"identity": 0.7,
"random": 0.15,
"mask_head": 0.05,
"mask_tail": 0.05,
"mask_head_tail": 0.05,

4
configs/opensora/train/16x256x256-spee-rflow.py
View file

@ -28,8 +28,8 @@ model = dict(
)
# mask_ratios = [0.5, 0.29, 0.07, 0.07, 0.07]
# mask_ratios = {
# "mask_no": 0.9,
# "mask_random": 0.06,
# "identity": 0.9,
# "random": 0.06,
# "mask_head": 0.01,
# "mask_tail": 0.01,
# "mask_head_tail": 0.02,

4
configs/opensora/train/16x256x256-spee.py
View file

@ -24,8 +24,8 @@ model = dict(
enable_layernorm_kernel=True,
)
mask_ratios = {
"mask_no": 0.5,
"mask_random": 0.29,
"identity": 0.5,
"random": 0.29,
"mask_head": 0.07,
"mask_tail": 0.07,
"mask_head_tail": 0.07,

18
docs/config.md
View file

@ -152,15 +152,15 @@ bucket_config = {
}
# mask ratio in training
mask_ratios = {
"mask_no": 0.75, # 75% no mask
"mask_quarter_random": 0.025, # 2.5% random mask with 1 frame to 1/4 #frames
"mask_quarter_head": 0.025, # 2.5% mask at the beginning with 1 frame to 1/4 #frames
"mask_quarter_tail": 0.025, # 2.5% mask at the end with 1 frame to 1/4 #frames
"mask_quarter_head_tail": 0.05, # 5% mask at the beginning and end with 1 frame to 1/4 #frames
"mask_image_random": 0.025, # 2.5% random mask with 1 image to 1/4 #images
"mask_image_head": 0.025, # 2.5% mask at the beginning with 1 image to 1/4 #images
"mask_image_tail": 0.025, # 2.5% mask at the end with 1 image to 1/4 #images
"mask_image_head_tail": 0.05, # 5% mask at the beginning and end with 1 image to 1/4 #images
"identity": 0.75, # 75% no mask
"quarter_random": 0.025, # 2.5% random mask with 1 frame to 1/4 #frames
"quarter_head": 0.025, # 2.5% mask at the beginning with 1 frame to 1/4 #frames
"quarter_tail": 0.025, # 2.5% mask at the end with 1 frame to 1/4 #frames
"quarter_head_tail": 0.05, # 5% mask at the beginning and end with 1 frame to 1/4 #frames
"image_random": 0.025, # 2.5% random mask with 1 image to 1/4 #images
"image_head": 0.025, # 2.5% mask at the beginning with 1 image to 1/4 #images
"image_tail": 0.025, # 2.5% mask at the end with 1 image to 1/4 #images
"image_head_tail": 0.05, # 5% mask at the beginning and end with 1 image to 1/4 #images
}
# Define acceleration

31
opensora/datasets/aspect.py
View file

@ -1,18 +1,6 @@
import math
# Ours
def get_h_w(a, ts, eps=1e-4):
h = (ts * a) ** 0.5
h = h + eps
h = math.ceil(h) if math.ceil(h) % 2 == 0 else math.floor(h)
w = h / a
w = w + eps
w = math.ceil(w) if math.ceil(w) % 2 == 0 else math.floor(w)
return h, w
# computation
AR = (
3 / 8,
9 / 21,
@ -32,7 +20,17 @@ AR = (
2 / 1,
1 / 0.48,
)
ARV = (0.375, 0.43, 0.48, 0.50, 0.53, 0.54, 0.56, 0.62, 0.67, 0.75, 1, 1.33, 1.50, 1.78, 1.89, 2, 2.08)
AR_fraction = (0.375, 0.43, 0.48, 0.50, 0.53, 0.54, 0.56, 0.62, 0.67, 0.75, 1, 1.33, 1.50, 1.78, 1.89, 2, 2.08)
def get_h_w(a, ts, eps=1e-4):
h = (ts * a) ** 0.5
h = h + eps
h = math.ceil(h) if math.ceil(h) % 2 == 0 else math.floor(h)
w = h / a
w = w + eps
w = math.ceil(w) if math.ceil(w) % 2 == 0 else math.floor(w)
return h, w
def get_aspect_ratios_dict(ts=360 * 640, ars=AR):
@ -40,6 +38,7 @@ def get_aspect_ratios_dict(ts=360 * 640, ars=AR):
return est
# computed from above code
# S = 8294400
ASPECT_RATIO_4K = {
"0.38": (1764, 4704),
@ -451,3 +450,7 @@ ASPECT_RATIOS = {
"2880": (8294400, ASPECT_RATIO_2880),
"4k": (8294400, ASPECT_RATIO_4K),
}
def get_num_pixels(name):
return ASPECT_RATIOS[name][0]

14
opensora/datasets/bucket.py
View file

@ -2,6 +2,8 @@ from collections import OrderedDict
import numpy as np
from opensora.utils.misc import get_logger
from .aspect import ASPECT_RATIOS, get_closest_ratio
@ -67,7 +69,7 @@ class Bucket:
self.t_criteria = t_criteria
self.ar_criteria = ar_criteria
self.num_bucket = num_bucket
print(f"Number of buckets: {num_bucket}")
get_logger().info("Number of buckets: %s", num_bucket)
def get_bucket_id(self, T, H, W, frame_interval=1, seed=None):
resolution = H * W
@ -92,6 +94,11 @@ class Bucket:
# otherwise, find suitable t_id for video
t_fail = True
for t_id, prob in t_criteria.items():
rng = np.random.default_rng(seed + self.bucket_id[hw_id][t_id])
if isinstance(prob, tuple):
prob_t = prob[1]
if rng.random() > prob_t:
continue
if T > t_id * frame_interval and t_id != 1:
t_fail = False
break
@ -99,8 +106,9 @@ class Bucket:
continue
# leave the loop if prob is high enough
rng = np.random.default_rng(seed + self.bucket_id[hw_id][t_id])
if prob == 1 or rng.random() < prob:
if isinstance(prob, tuple):
prob = prob[0]
if prob >= 1 or rng.random() < prob:
fail = False
break
if fail:

33
opensora/datasets/dataloader.py
View file

@ -1,40 +1,13 @@
import random
from typing import Iterator, Optional
from typing import Optional
import numpy as np
import torch
from torch.distributed import ProcessGroup
from torch.distributed.distributed_c10d import _get_default_group
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
from .sampler import VariableVideoBatchSampler
class StatefulDistributedSampler(DistributedSampler):
def __init__(
self,
dataset: Dataset,
num_replicas: Optional[int] = None,
rank: Optional[int] = None,
shuffle: bool = True,
seed: int = 0,
drop_last: bool = False,
) -> None:
super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
self.start_index: int = 0
def __iter__(self) -> Iterator:
iterator = super().__iter__()
indices = list(iterator)
indices = indices[self.start_index :]
return iter(indices)
def __len__(self) -> int:
return self.num_samples - self.start_index
def set_start_index(self, start_index: int) -> None:
self.start_index = start_index
from .sampler import StatefulDistributedSampler, VariableVideoBatchSampler
def prepare_dataloader(

98
opensora/datasets/sampler.py
View file

@ -1,19 +1,22 @@
import warnings
from collections import OrderedDict, defaultdict
from pprint import pprint
from pprint import pformat
from typing import Iterator, List, Optional
import torch
import torch.distributed as dist
from pandarallel import pandarallel
from torch.utils.data import DistributedSampler
from torch.utils.data import Dataset, DistributedSampler
from opensora.utils.misc import format_numel_str, get_logger
from .aspect import get_num_pixels
from .bucket import Bucket
from .datasets import VariableVideoTextDataset
# HACK: use pandarallel
# pandarallel should only access local variables
# use pandarallel to accelerate bucket processing
# NOTE: pandarallel should only access local variables
def apply(data, method=None, frame_interval=None, seed=None, num_bucket=None):
return method(
data["num_frames"],
@ -24,6 +27,32 @@ def apply(data, method=None, frame_interval=None, seed=None, num_bucket=None):
)
class StatefulDistributedSampler(DistributedSampler):
def __init__(
self,
dataset: Dataset,
num_replicas: Optional[int] = None,
rank: Optional[int] = None,
shuffle: bool = True,
seed: int = 0,
drop_last: bool = False,
) -> None:
super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
self.start_index: int = 0
def __iter__(self) -> Iterator:
iterator = super().__iter__()
indices = list(iterator)
indices = indices[self.start_index :]
return iter(indices)
def __len__(self) -> int:
return self.num_samples - self.start_index
def set_start_index(self, start_index: int) -> None:
self.start_index = start_index
class VariableVideoBatchSampler(DistributedSampler):
def __init__(
self,
@ -53,6 +82,7 @@ class VariableVideoBatchSampler(DistributedSampler):
bucket_sample_dict = OrderedDict()
pandarallel.initialize(nb_workers=self.num_bucket_build_workers, progress_bar=False)
get_logger().info("Building buckets...")
bucket_ids = self.dataset.data.parallel_apply(
apply,
axis=1,
@ -199,29 +229,51 @@ class VariableVideoBatchSampler(DistributedSampler):
def load_state_dict(self, state_dict: dict) -> None:
self.__dict__.update(state_dict)
def _print_bucket_info(self, bucket_sample_dict: dict, verbose=True) -> None:
def _print_bucket_info(self, bucket_sample_dict: dict) -> None:
# collect statistics
total_samples = 0
num_batch = 0
num_dict = {}
num_aspect_dict = defaultdict(int)
num_hwt_dict = defaultdict(int)
total_batch = 0
num_aspect_dict = defaultdict(lambda: [0, 0])
num_hwt_dict = defaultdict(lambda: [0, 0])
for k, v in bucket_sample_dict.items():
size = len(v)
num_batch = size // self.bucket.get_batch_size(k[:-1])
total_samples += size
num_dict[k] = size
num_aspect_dict[k[-1]] += size
num_hwt_dict[k[:-1]] += size
num_batch += size // self.bucket.get_batch_size(k[:-1])
if dist.get_rank() == 0 and verbose:
print(f"Total training samples: {total_samples}, num buckets: {len(num_dict)}")
print("Bucket samples:")
pprint(num_dict)
print("Bucket samples by aspect ratio:")
pprint(num_aspect_dict)
print("Bucket samples by HxWxT:")
pprint(num_hwt_dict)
print(f"Number of batches: {num_batch}")
self.approximate_num_batch = num_batch
total_batch += num_batch
num_aspect_dict[k[-1]][0] += size
num_aspect_dict[k[-1]][1] += num_batch
num_hwt_dict[k[:-1]][0] += size
num_hwt_dict[k[:-1]][1] += num_batch
# sort
num_aspect_dict = dict(sorted(num_aspect_dict.items(), key=lambda x: x[0]))
num_hwt_dict = dict(
sorted(num_hwt_dict.items(), key=lambda x: (get_num_pixels(x[0][0]), x[0][1]), reverse=True)
)
num_hwt_img_dict = {k: v for k, v in num_hwt_dict.items() if k[1] == 1}
num_hwt_vid_dict = {k: v for k, v in num_hwt_dict.items() if k[1] > 1}
# log
if dist.get_rank() == 0 and self.verbose:
get_logger().info("Bucket Info:")
get_logger().info(
"Bucket [#sample, #batch] by aspect ratio:\n%s", pformat(num_aspect_dict, sort_dicts=False)
)
get_logger().info(
"Image Bucket [#sample, #batch] by HxWxT:\n%s", pformat(num_hwt_img_dict, sort_dicts=False)
)
get_logger().info(
"Video Bucket [#sample, #batch] by HxWxT:\n%s", pformat(num_hwt_vid_dict, sort_dicts=False)
)
get_logger().info(
"#training batch: %s, #training sample: %s, #non empty bucket: %s",
format_numel_str(total_batch),
format_numel_str(total_samples),
len(bucket_sample_dict),
)
self.approximate_num_batch = total_batch
def set_epoch(self, epoch: int) -> None:
super().set_epoch(epoch)

44
opensora/utils/ckpt_utils.py
View file

@ -13,6 +13,8 @@ from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler
from torchvision.datasets.utils import download_url
from .misc import get_logger
hf_endpoint = os.environ.get("HF_ENDPOINT")
if hf_endpoint is None:
hf_endpoint = "https://huggingface.co"
@ -41,7 +43,7 @@ pretrained_models = {
def reparameter(ckpt, name=None, model=None):
name = os.path.basename(name)
if not dist.is_initialized() or dist.get_rank() == 0:
print("loading pretrained model:", name)
get_logger().info("loading pretrained model: %s", name)
if name in ["DiT-XL-2-512x512.pt", "DiT-XL-2-256x256.pt"]:
ckpt["x_embedder.proj.weight"] = ckpt["x_embedder.proj.weight"].unsqueeze(2)
del ckpt["pos_embed"]
@ -80,16 +82,20 @@ def reparameter(ckpt, name=None, model=None):
# different text length
if "y_embedder.y_embedding" in ckpt:
if ckpt["y_embedder.y_embedding"].shape[0] < model.y_embedder.y_embedding.shape[0]:
print(
f"Extend y_embedding from {ckpt['y_embedder.y_embedding'].shape[0]} to {model.y_embedder.y_embedding.shape[0]}"
get_logger().info(
"Extend y_embedding from %s to %s",
ckpt["y_embedder.y_embedding"].shape[0],
model.y_embedder.y_embedding.shape[0],
)
additional_length = model.y_embedder.y_embedding.shape[0] - ckpt["y_embedder.y_embedding"].shape[0]
new_y_embedding = torch.zeros(additional_length, model.y_embedder.y_embedding.shape[1])
new_y_embedding[:] = ckpt["y_embedder.y_embedding"][-1]
ckpt["y_embedder.y_embedding"] = torch.cat([ckpt["y_embedder.y_embedding"], new_y_embedding], dim=0)
elif ckpt["y_embedder.y_embedding"].shape[0] > model.y_embedder.y_embedding.shape[0]:
print(
f"Shrink y_embedding from {ckpt['y_embedder.y_embedding'].shape[0]} to {model.y_embedder.y_embedding.shape[0]}"
get_logger().info(
"Shrink y_embedding from %s to %s",
ckpt["y_embedder.y_embedding"].shape[0],
model.y_embedder.y_embedding.shape[0],
)
ckpt["y_embedder.y_embedding"] = ckpt["y_embedder.y_embedding"][: model.y_embedder.y_embedding.shape[0]]
# stdit3 special case
@ -169,16 +175,6 @@ def model_gathering(model: torch.nn.Module, model_shape_dict: dict):
dist.barrier()
def load_json(file_path: str):
with open(file_path, "r") as f:
return json.load(f)
def save_json(data, file_path: str):
with open(file_path, "w") as f:
json.dump(data, f, indent=4)
def remove_padding(tensor: torch.Tensor, original_shape: Tuple) -> torch.Tensor:
return tensor[: functools.reduce(operator.mul, original_shape)]
@ -194,19 +190,29 @@ def load_checkpoint(model, ckpt_path, save_as_pt=False, model_name="model"):
if ckpt_path.endswith(".pt") or ckpt_path.endswith(".pth"):
state_dict = find_model(ckpt_path, model=model)
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print(f"Missing keys: {missing_keys}")
print(f"Unexpected keys: {unexpected_keys}")
get_logger().info("Missing keys: %s", missing_keys)
get_logger().info("Unexpected keys: %s", unexpected_keys)
elif os.path.isdir(ckpt_path):
load_from_sharded_state_dict(model, ckpt_path, model_name)
print(f"Model checkpoint loaded from {ckpt_path}")
get_logger().info("Model checkpoint loaded from %s", ckpt_path)
if save_as_pt:
save_path = os.path.join(ckpt_path, model_name + "_ckpt.pt")
torch.save(model.state_dict(), save_path)
print(f"Model checkpoint saved to {save_path}")
get_logger().info("Model checkpoint saved to %s", save_path)
else:
raise ValueError(f"Invalid checkpoint path: {ckpt_path}")
def load_json(file_path: str):
with open(file_path, "r") as f:
return json.load(f)
def save_json(data, file_path: str):
with open(file_path, "w") as f:
json.dump(data, f, indent=4)
# save and load for training

26
opensora/utils/config_utils.py
View file

@ -86,16 +86,17 @@ def merge_args(cfg, args, training=False):
if args.data_path is not None:
cfg.dataset["data_path"] = args.data_path
args.data_path = None
if not training and args.image_size is not None and "dataset" in cfg:
cfg.dataset["image_size"] = args.image_size
if not training and args.num_frames is not None and "dataset" in cfg:
cfg.dataset["num_frames"] = args.num_frames
if not training and args.cfg_scale is not None:
cfg.scheduler["cfg_scale"] = args.cfg_scale
args.cfg_scale = None
if not training and args.num_sampling_steps is not None:
cfg.scheduler["num_sampling_steps"] = args.num_sampling_steps
args.num_sampling_steps = None
if not training:
if args.image_size is not None and "dataset" in cfg:
cfg.dataset["image_size"] = args.image_size
if args.num_frames is not None and "dataset" in cfg:
cfg.dataset["num_frames"] = args.num_frames
if args.cfg_scale is not None:
cfg.scheduler["cfg_scale"] = args.cfg_scale
args.cfg_scale = None
if args.num_sampling_steps is not None:
cfg.scheduler["num_sampling_steps"] = args.num_sampling_steps
args.num_sampling_steps = None
for k, v in vars(args).items():
if v is not None:
@ -128,11 +129,6 @@ def merge_args(cfg, args, training=False):
cfg["prompt"] = cfg["prompt"][: args.end_index]
if "multi_resolution" not in cfg:
cfg["multi_resolution"] = False
else:
# Training only
# - Allow not set
if "transform_name" not in cfg.dataset:
cfg.dataset["transform_name"] = "center"
return cfg

161
opensora/utils/misc.py
View file

@ -1,7 +1,6 @@
import collections
import importlib
import logging
import os
import time
from collections import OrderedDict
from collections.abc import Sequence
@ -12,6 +11,35 @@ import numpy as np
import torch
import torch.distributed as dist
# ======================================================
# Logging
# ======================================================
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
def get_logger():
return logging.getLogger(__name__)
def print_rank(var_name, var_value, rank=0):
if dist.get_rank() == rank:
@ -23,12 +51,9 @@ def print_0(*args, **kwargs):
print(*args, **kwargs)
def requires_grad(model: torch.nn.Module, flag: bool = True) -> None:
"""
Set requires_grad flag for all parameters in a model.
"""
for p in model.parameters():
p.requires_grad = flag
# ======================================================
# String
# ======================================================
def format_numel_str(numel: int) -> str:
@ -45,48 +70,6 @@ def format_numel_str(numel: int) -> str:
return f"{numel}"
def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
tensor.div_(dist.get_world_size())
return tensor
def get_model_numel(model: torch.nn.Module) -> Tuple[int, int]:
num_params = 0
num_params_trainable = 0
for p in model.parameters():
num_params += p.numel()
if p.requires_grad:
num_params_trainable += p.numel()
return num_params, num_params_trainable
def try_import(name):
"""Try to import a module.
Args:
name (str): Specifies what module to import in absolute or relative
terms (e.g. either pkg.mod or ..mod).
Returns:
ModuleType or None: If importing successfully, returns the imported
module, otherwise returns None.
"""
try:
return importlib.import_module(name)
except ImportError:
return None
def transpose(x):
"""
transpose a list of list
Args:
x (list[list]):
"""
ret = list(map(list, zip(*x)))
return ret
def get_timestamp():
timestamp = time.strftime("%Y%m%d-%H%M%S", time.localtime(time.time()))
return timestamp
@ -125,6 +108,39 @@ def format_time(seconds):
return f
# ======================================================
# PyTorch
# ======================================================
def requires_grad(model: torch.nn.Module, flag: bool = True) -> None:
"""
Set requires_grad flag for all parameters in a model.
"""
for p in model.parameters():
p.requires_grad = flag
def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
tensor.div_(dist.get_world_size())
return tensor
def get_model_numel(model: torch.nn.Module) -> Tuple[int, int]:
num_params = 0
num_params_trainable = 0
for p in model.parameters():
num_params += p.numel()
if p.requires_grad:
num_params_trainable += p.numel()
return num_params, num_params_trainable
def count_params(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def to_tensor(data):
"""Convert objects of various python types to :obj:`torch.Tensor`.
@ -186,10 +202,6 @@ def to_torch_dtype(dtype):
raise ValueError
def count_params(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def _ntuple(n):
def parse(x):
if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
@ -252,6 +264,11 @@ def inverse_sigmoid(x, eps=1e-5):
return torch.log(x1 / x2)
# ======================================================
# Python
# ======================================================
def count_columns(df, columns):
cnt_dict = OrderedDict()
num_samples = len(df)
@ -265,23 +282,27 @@ def count_columns(df, columns):
return cnt_dict
def build_logger(work_dir, cfgname):
log_file = cfgname + ".log"
log_path = os.path.join(work_dir, log_file)
def try_import(name):
"""Try to import a module.
logger = logging.getLogger(cfgname)
logger.setLevel(logging.INFO)
# formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
formatter = logging.Formatter("%(asctime)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
Args:
name (str): Specifies what module to import in absolute or relative
terms (e.g. either pkg.mod or ..mod).
Returns:
ModuleType or None: If importing successfully, returns the imported
module, otherwise returns None.
"""
try:
return importlib.import_module(name)
except ImportError:
return None
handler1 = logging.FileHandler(log_path)
handler1.setFormatter(formatter)
handler2 = logging.StreamHandler()
handler2.setFormatter(formatter)
logger.addHandler(handler1)
logger.addHandler(handler2)
logger.propagate = False
return logger
def transpose(x):
"""
transpose a list of list
Args:
x (list[list]):
"""
ret = list(map(list, zip(*x)))
return ret

74
opensora/utils/train_utils.py
View file

@ -1,4 +1,3 @@
import logging
import math
import random
from collections import OrderedDict
@ -10,6 +9,8 @@ from colossalai.booster.plugin import LowLevelZeroPlugin
from opensora.acceleration.parallel_states import set_data_parallel_group, set_sequence_parallel_group
from opensora.acceleration.plugin import ZeroSeqParallelPlugin
from .misc import get_logger
def create_colossalai_plugin(plugin, dtype, grad_clip, sp_size):
if plugin == "zero2":
@ -66,17 +67,17 @@ def update_ema(
class MaskGenerator:
def __init__(self, mask_ratios):
valid_mask_names = [
"mask_no",
"mask_quarter_random",
"mask_quarter_head",
"mask_quarter_tail",
"mask_quarter_head_tail",
"mask_image_random",
"mask_image_head",
"mask_image_tail",
"mask_image_head_tail",
"mask_random",
"mask_intepolate",
"identity",
"quarter_random",
"quarter_head",
"quarter_tail",
"quarter_head_tail",
"image_random",
"image_head",
"image_tail",
"image_head_tail",
"random",
"intepolate",
]
assert all(
mask_name in valid_mask_names for mask_name in mask_ratios.keys()
@ -88,12 +89,12 @@ class MaskGenerator:
mask_ratio <= 1 for mask_ratio in mask_ratios.values()
), f"mask_ratio should be less than or equal to 1, got {mask_ratios.values()}"
# sum of mask_ratios should be 1
if "mask_no" not in mask_ratios:
mask_ratios["mask_no"] = 1.0 - sum(mask_ratios.values())
if "identity" not in mask_ratios:
mask_ratios["identity"] = 1.0 - sum(mask_ratios.values())
assert math.isclose(
sum(mask_ratios.values()), 1.0, abs_tol=1e-6
), f"sum of mask_ratios should be 1, got {sum(mask_ratios.values())}"
print(f"mask ratios: {mask_ratios}")
get_logger().info("mask ratios: %s", mask_ratios)
self.mask_ratios = mask_ratios
def get_mask(self, x):
@ -114,39 +115,39 @@ class MaskGenerator:
if num_frames <= 1:
return mask
if mask_name == "mask_quarter_random":
if mask_name == "quarter_random":
random_size = random.randint(1, condition_frames_max)
random_pos = random.randint(0, x.shape[2] - random_size)
mask[random_pos : random_pos + random_size] = 0
elif mask_name == "mask_image_random":
elif mask_name == "image_random":
random_size = 1
random_pos = random.randint(0, x.shape[2] - random_size)
mask[random_pos : random_pos + random_size] = 0
elif mask_name == "mask_quarter_head":
elif mask_name == "quarter_head":
random_size = random.randint(1, condition_frames_max)
mask[:random_size] = 0
elif mask_name == "mask_image_head":
elif mask_name == "image_head":
random_size = 1
mask[:random_size] = 0
elif mask_name == "mask_quarter_tail":
elif mask_name == "quarter_tail":
random_size = random.randint(1, condition_frames_max)
mask[-random_size:] = 0
elif mask_name == "mask_image_tail":
elif mask_name == "image_tail":
random_size = 1
mask[-random_size:] = 0
elif mask_name == "mask_quarter_head_tail":
elif mask_name == "quarter_head_tail":
random_size = random.randint(1, condition_frames_max)
mask[:random_size] = 0
mask[-random_size:] = 0
elif mask_name == "mask_image_head_tail":
elif mask_name == "image_head_tail":
random_size = 1
mask[:random_size] = 0
mask[-random_size:] = 0
elif mask_name == "mask_intepolate":
elif mask_name == "intepolate":
random_start = random.randint(0, 1)
mask[random_start::2] = 0
elif mask_name == "mask_random":
mask_ratio = random.uniform(0.3, 0.7)
elif mask_name == "random":
mask_ratio = random.uniform(0.1, 0.9)
mask = torch.rand(num_frames, device=x.device) > mask_ratio
# if mask is all False, set the last frame to True
if not mask.any():
@ -161,24 +162,3 @@ class MaskGenerator:
masks.append(mask)
masks = torch.stack(masks, dim=0)
return masks
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

33
scripts/train.py
View file

@ -23,8 +23,15 @@ from opensora.utils.config_utils import (
parse_configs,
save_training_config,
)
from opensora.utils.misc import all_reduce_mean, format_numel_str, get_model_numel, requires_grad, to_torch_dtype
from opensora.utils.train_utils import MaskGenerator, create_colossalai_plugin, create_logger, update_ema
from opensora.utils.misc import (
all_reduce_mean,
create_logger,
format_numel_str,
get_model_numel,
requires_grad,
to_torch_dtype,
)
from opensora.utils.train_utils import MaskGenerator, create_colossalai_plugin, update_ema
DEFAULT_DATASET_NAME = "VideoTextDataset"
@ -79,6 +86,7 @@ def main():
# ======================================================
# 2. build dataset and dataloader
# ======================================================
logger.info("Building dataset...")
# == build dataset ==
dataset = build_module(cfg.dataset, DATASETS)
logger.info("Dataset contains %s samples.", len(dataset))
@ -98,16 +106,21 @@ def main():
dataloader = prepare_dataloader(**dataloader_args)
total_batch_size = cfg.batch_size * dist.get_world_size() // cfg.get("sp_size", 1)
logger.info("Total batch size: %s", total_batch_size)
num_steps_per_epoch = len(dataloader)
sampler_to_io = None
else:
dataloader = prepare_variable_dataloader(
bucket_config=cfg.get("bucket_config", None),
num_bucket_build_workers=cfg.get("num_bucket_build_workers", 1),
**dataloader_args,
)
num_steps_per_epoch = dataloader.batch_sampler.get_num_batch() // dist.get_world_size()
sampler_to_io = None if cfg.get("start_from_scratch ", False) else dataloader.batch_sampler
# ======================================================
# 3. build model
# ======================================================
logger.info("Building models...")
# == build text-encoder and vae ==
text_encoder = build_module(cfg.text_encoder, MODELS, device=device)
vae = build_module(cfg.vae, MODELS).to(device, dtype)
@ -159,8 +172,9 @@ def main():
mask_generator = MaskGenerator(cfg.mask_ratios)
# =======================================================
# 5. distributed training preparation with colossalai
# 4. distributed training preparation with colossalai
# =======================================================
logger.info("Preparing for distributed training...")
# == boosting ==
# NOTE: we set dtype first to make initialization of model consistent with the dtype; then reset it to the fp32 as we make diffusion scheduler in fp32
torch.set_default_dtype(dtype)
@ -172,19 +186,12 @@ def main():
)
torch.set_default_dtype(torch.float)
logger.info("Boosting model for distributed training")
if cfg.dataset.type == DEFAULT_DATASET_NAME:
num_steps_per_epoch = len(dataloader)
sampler_to_io = None
else:
num_steps_per_epoch = dataloader.batch_sampler.get_num_batch() // dist.get_world_size()
sampler_to_io = None if cfg.get("start_from_scratch ", False) else dataloader.batch_sampler
cfg_epochs = cfg.get("epochs", 1000)
logger.info("Training for %s epochs with %s steps per epoch", cfg_epochs, num_steps_per_epoch)
# == global variables ==
cfg_epochs = cfg.get("epochs", 1000)
start_epoch = start_step = log_step = sampler_start_idx = acc_step = 0
running_loss = 0.0
logger.info("Training for %s epochs with %s steps per epoch", cfg_epochs, num_steps_per_epoch)
# == resume ==
if cfg.get("load", None) is not None:
@ -207,7 +214,7 @@ def main():
model_sharding(ema)
# =======================================================
# 6. training loop
# 5. training loop
# =======================================================
dist.barrier()
for epoch in range(start_epoch, cfg_epochs):

14
scripts/train_vae.py
View file

@ -25,8 +25,8 @@ from opensora.utils.config_utils import (
parse_configs,
save_training_config,
)
from opensora.utils.misc import all_reduce_mean, format_numel_str, get_model_numel, to_torch_dtype
from opensora.utils.train_utils import create_colossalai_plugin, create_logger
from opensora.utils.misc import all_reduce_mean, create_logger, format_numel_str, get_model_numel, to_torch_dtype
from opensora.utils.train_utils import create_colossalai_plugin
DEFAULT_DATASET_NAME = "VideoTextDataset"
@ -81,6 +81,7 @@ def main():
# ======================================================
# 2. build dataset and dataloader
# ======================================================
logger.info("Building dataset...")
# == build dataset ==
assert cfg.dataset.type == DEFAULT_DATASET_NAME, "Only support VideoTextDataset for vae training"
dataset = build_module(cfg.dataset, DATASETS)
@ -100,10 +101,12 @@ def main():
dataloader = prepare_dataloader(**dataloader_args)
total_batch_size = cfg.batch_size * dist.get_world_size() // cfg.get("sp_size", 1)
logger.info("Total batch size: %s", total_batch_size)
num_steps_per_epoch = len(dataloader)
# ======================================================
# 3. build model
# ======================================================
logger.info("Building models...")
# == build vae model ==
model = build_module(cfg.model, MODELS).to(device, dtype)
model.train()
@ -177,6 +180,7 @@ def main():
# =======================================================
# 4. distributed training preparation with colossalai
# =======================================================
logger.info("Preparing for distributed training...")
# == boosting ==
# NOTE: we set dtype first to make initialization of model consistent with the dtype; then reset it to the fp32 as we make diffusion scheduler in fp32
torch.set_default_dtype(dtype)
@ -194,14 +198,12 @@ def main():
)
torch.set_default_dtype(torch.float)
logger.info("Boosting model for distributed training")
num_steps_per_epoch = len(dataloader)
cfg_epochs = cfg.get("epochs", 1000)
logger.info("Training for %s epochs with %s steps per epoch", cfg_epochs, num_steps_per_epoch)
# == global variables ==
cfg_epochs = cfg.get("epochs", 1000)
start_epoch = start_step = log_step = sampler_start_idx = acc_step = 0
running_loss = running_disc_loss = 0.0
logger.info("Training for %s epochs with %s steps per epoch", cfg_epochs, num_steps_per_epoch)
# == resume ==
if cfg.get("load", None) is not None: