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

added missing function call (#18)

Browse source
This commit is contained in:
Frank Lee 2024-02-29 14:27:25 +08:00 committed by GitHub
parent 56cabd8e5e
commit c1767416bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 10 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
train.py
View file

@ -28,12 +28,7 @@ from tqdm import tqdm
from open_sora.diffusion import create_diffusion
from open_sora.modeling import DiT_models
from open_sora.utils.data import (
create_video_compressor,
load_datasets,
make_batch,
preprocess_batch,
)
from open_sora.utils.data import create_video_compressor, load_datasets, make_batch, preprocess_batch
#################################################################################
# Training Helper Functions #
@ -95,6 +90,7 @@ def main(args):
launch_from_torch({})
coordinator = DistCoordinator()
logger = get_dist_logger()
configure_backends()
# Step 2: set up acceleration plugins
plugin = LowLevelZeroPlugin(stage=2, precision="fp16")
@ -123,9 +119,7 @@ def main(args):
model.enable_gradient_checkpointing()
# Step 5: create diffusion pipeline
diffusion = create_diffusion(
timestep_respacing=""
) # default: 1000 steps, linear noise schedule
diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule
# Step 6: setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
opt = HybridAdam(model.parameters(), lr=args.lr, weight_decay=0)
@ -139,9 +133,7 @@ def main(args):
shuffle=True,
drop_last=True,
)
lr_scheduler = CosineAnnealingLR(
opt, args.epochs * len(dataloader) // args.accumulation_steps
)
lr_scheduler = CosineAnnealingLR(opt, args.epochs * len(dataloader) // args.accumulation_steps)
logger.info(f"Dataset contains {len(dataset)} samples", ranks=[0])
# Step 8: setup booster
@ -176,9 +168,7 @@ def main(args):
(video_inputs.shape[0],),
device=video_inputs.device,
)
loss_dict = diffusion.training_losses(
model, video_inputs, t, batch, mask=mask
)
loss_dict = diffusion.training_losses(model, video_inputs, t, batch, mask=mask)
loss = loss_dict["loss"].mean() / args.accumulation_steps
total_loss.add_(loss.data)
booster.backward(loss, opt)
@ -192,9 +182,7 @@ def main(args):
all_reduce_mean(total_loss)
pbar.set_postfix({"Loss": f"{total_loss.item():.4f}"})
if coordinator.is_master():
global_step = (epoch * num_steps_per_epoch) + (
step + 1
) // args.accumulation_steps
global_step = (epoch * num_steps_per_epoch) + (step + 1) // args.accumulation_steps
writer.add_scalar(
tag="Loss",
scalar_value=total_loss.item(),
@ -205,12 +193,9 @@ def main(args):
# Save DiT checkpoint:
if args.save_interval > 0 and (
(step + 1) % (args.save_interval * args.accumulation_steps) == 0
or (step + 1) == len(dataloader)
(step + 1) % (args.save_interval * args.accumulation_steps) == 0 or (step + 1) == len(dataloader)
):
save_path = os.path.join(
args.checkpoint_dir, f"epoch-{epoch}-step-{step}"
)
save_path = os.path.join(args.checkpoint_dir, f"epoch-{epoch}-step-{step}")
save_checkpoints(booster, model, opt, ema, save_path, coordinator)
logger.info(f"Saved checkpoint to {save_path}", ranks=[0])
@ -227,9 +212,7 @@ def main(args):
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(
"-m", "--model", type=str, choices=list(DiT_models.keys()), default="DiT-S/8"
)
parser.add_argument("-m", "--model", type=str, choices=list(DiT_models.keys()), default="DiT-S/8")
parser.add_argument("-d", "--dataset", nargs="+", default=[])
parser.add_argument("-v", "--video_dir", type=str, required=True)
parser.add_argument("-e", "--epochs", type=int, default=10)
@ -240,9 +223,7 @@ if __name__ == "__main__":
parser.add_argument("--save_interval", type=int, default=20)
parser.add_argument("--checkpoint_dir", type=str, default="checkpoints")
parser.add_argument("--tensorboard_dir", type=str, default="runs")
parser.add_argument(
"-c", "--compressor", choices=["raw", "vqvae", "vae"], default="raw"
)
parser.add_argument("-c", "--compressor", choices=["raw", "vqvae", "vae"], default="raw")
parser.add_argument("--load_model", default=None)
parser.add_argument("--load_optimizer", default=None)
args = parser.parse_args()