109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
import sys
|
|
sys.path.append(".")
|
|
import torch
|
|
import random
|
|
import numpy as np
|
|
from opensora.models.ae.videobase import (
|
|
CausalVAEModel,
|
|
)
|
|
from torch.utils.data import DataLoader
|
|
from opensora.models.ae.videobase.dataset_videobase import VideoDataset
|
|
import argparse
|
|
from transformers import HfArgumentParser
|
|
from dataclasses import dataclass, field, asdict
|
|
import torch.distributed as dist
|
|
import os
|
|
import pytorch_lightning as pl
|
|
from pytorch_lightning.loggers import WandbLogger
|
|
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
|
|
|
|
@dataclass
|
|
class TrainingArguments:
|
|
exp_name: str = field(default="causalvae")
|
|
batch_size: int = field(default=1)
|
|
precision: str = field(default="bf16")
|
|
max_steps: int = field(default=100000)
|
|
save_steps: int = field(default=2000)
|
|
output_dir: str = field(default="results/causalvae")
|
|
video_path: str = field(default="/remote-home1/dataset/data_split_tt")
|
|
video_num_frames: int = field(default=17)
|
|
sample_rate: int = field(default=1)
|
|
dynamic_sample: bool = field(default=False)
|
|
model_config: str = field(default="scripts/causalvae/288.yaml")
|
|
n_nodes: int = field(default=1)
|
|
devices: int = field(default=8)
|
|
resolution: int = field(default=64)
|
|
num_workers: int = field(default=8)
|
|
resume_from_checkpoint: str = field(default=None)
|
|
load_from_checkpoint: str = field(default=None)
|
|
|
|
|
|
def set_seed(seed=1006):
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
|
|
def load_callbacks_and_logger(args):
|
|
checkpoint_callback = ModelCheckpoint(
|
|
dirpath=args.output_dir,
|
|
filename="model-{epoch:02d}-{step}",
|
|
every_n_train_steps=args.save_steps,
|
|
save_top_k=-1,
|
|
save_on_train_epoch_end=False,
|
|
)
|
|
lr_monitor = LearningRateMonitor(logging_interval="step")
|
|
logger = WandbLogger(name=args.exp_name, log_model=False)
|
|
return [checkpoint_callback, lr_monitor], logger
|
|
|
|
def train(args):
|
|
set_seed()
|
|
# Load Config
|
|
model = CausalVAEModel()
|
|
if args.load_from_checkpoint is not None:
|
|
model = CausalVAEModel.from_pretrained(args.load_from_checkpoint)
|
|
else:
|
|
model = CausalVAEModel.from_config(args.model_config)
|
|
|
|
if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
|
|
print(model)
|
|
|
|
# Load Dataset
|
|
dataset = VideoDataset(args.video_path, sequence_length=args.video_num_frames, resolution=args.resolution, sample_rate=args.sample_rate, dynamic_sample=args.dynamic_sample)
|
|
train_loader = DataLoader(
|
|
dataset,
|
|
shuffle=True,
|
|
num_workers=args.num_workers,
|
|
batch_size=args.batch_size,
|
|
pin_memory=True,
|
|
)
|
|
# Load Callbacks and Logger
|
|
callbacks, logger = load_callbacks_and_logger(args)
|
|
# Load Trainer
|
|
trainer = pl.Trainer(
|
|
accelerator="cuda",
|
|
devices=args.devices,
|
|
num_nodes=args.n_nodes,
|
|
callbacks=callbacks,
|
|
logger=logger,
|
|
log_every_n_steps=5,
|
|
precision=args.precision,
|
|
max_steps=args.max_steps,
|
|
strategy="ddp_find_unused_parameters_true"
|
|
)
|
|
trainer_kwargs = {}
|
|
if args.resume_from_checkpoint:
|
|
trainer_kwargs['ckpt_path'] = args.resume_from_checkpoint
|
|
|
|
trainer.fit(
|
|
model,
|
|
train_loader,
|
|
**trainer_kwargs
|
|
)
|
|
# Save Huggingface Model
|
|
model.save_pretrained(os.path.join(args.output_dir, "hf"))
|
|
|
|
if __name__ == "__main__":
|
|
parser = HfArgumentParser(TrainingArguments)
|
|
args = parser.parse_args_into_dataclasses()
|
|
train(args[0])
|