(
dataset: torch.utils.data.Dataset,
model: DiffusionTrainingModule,
model_logger: ModelLogger,
optimizer: torch.optim.Optimizer,
scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
batch_size: int = 1,
clip_grad_norm: float = 1.0,
num_workers: int = 8,
save_steps: int = None,
num_epochs: int = 1,
gradient_accumulation_steps: int = 1,
find_unused_parameters: bool = False,
log_dir: Optional[str] = None,
prefetch_factor: int = 2,
resume_from=None, # Resume directory (points to a previous output_path).
)
| 841 | return json.load(f) |
| 842 | |
| 843 | def launch_training_task( |
| 844 | dataset: torch.utils.data.Dataset, |
| 845 | model: DiffusionTrainingModule, |
| 846 | model_logger: ModelLogger, |
| 847 | optimizer: torch.optim.Optimizer, |
| 848 | scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None, |
| 849 | batch_size: int = 1, |
| 850 | clip_grad_norm: float = 1.0, |
| 851 | num_workers: int = 8, |
| 852 | save_steps: int = None, |
| 853 | num_epochs: int = 1, |
| 854 | gradient_accumulation_steps: int = 1, |
| 855 | find_unused_parameters: bool = False, |
| 856 | log_dir: Optional[str] = None, |
| 857 | prefetch_factor: int = 2, |
| 858 | resume_from=None, # Resume directory (points to a previous output_path). |
| 859 | ): |
| 860 | def collate_skip_none(batch): |
| 861 | batch = [b for b in batch if b is not None] |
| 862 | if len(batch) == 0: |
| 863 | return None |
| 864 | return torch.utils.data.dataloader.default_collate(batch) |
| 865 | |
| 866 | dataloader = torch.utils.data.DataLoader( |
| 867 | dataset, batch_size=batch_size, shuffle=True, pin_memory=True, |
| 868 | num_workers=num_workers, prefetch_factor=prefetch_factor, |
| 869 | in_order=True, collate_fn=collate_skip_none |
| 870 | ) |
| 871 | # Enable logging with Accelerator if log_dir is provided |
| 872 | log_with = None |
| 873 | if log_dir is not None: |
| 874 | log_with = "tensorboard" |
| 875 | accelerator = Accelerator( |
| 876 | gradient_accumulation_steps=gradient_accumulation_steps, |
| 877 | kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=find_unused_parameters)], |
| 878 | log_with=log_with, |
| 879 | project_dir=log_dir, |
| 880 | ) |
| 881 | from accelerate.utils import set_seed |
| 882 | set_seed(42, device_specific=True) |
| 883 | # Build the scheduler before accelerator.prepare so world_size is known. |
| 884 | if scheduler is None: |
| 885 | steps_per_epoch = math.ceil(len(dataset) / max(1, batch_size)) |
| 886 | optim_steps_per_epoch = math.ceil(steps_per_epoch / max(1, gradient_accumulation_steps)) |
| 887 | total_steps = max(1, optim_steps_per_epoch * max(1, num_epochs)) |
| 888 | warmup_steps = min(100, max(0, total_steps - 1)) |
| 889 | decay_steps = max(1, int(round(total_steps * 0.10))) |
| 890 | stable_steps = max(0, total_steps - warmup_steps - decay_steps) |
| 891 | scheduler = get_wsd_schedule( |
| 892 | optimizer=optimizer, |
| 893 | num_warmup_steps=warmup_steps, |
| 894 | num_stable_steps=stable_steps, |
| 895 | num_decay_steps=decay_steps, |
| 896 | ) |
| 897 | # accelerator.prepare wraps the model in DDP (for multi-GPU), shards the |
| 898 | # dataloader across ranks, and handles mixed-precision casting. |
| 899 | model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler) |
| 900 |
no test coverage detected