Function forward_step

sat/train_video.py:191–215

(data_iterator, model, args, timers, data_class=None)

def forward_step(data_iterator, model, args, timers, data_class=None):
    if mpu.get_model_parallel_rank() == 0:
        timers("data loader").start()
        batch = next(data_iterator)
        timers("data loader").stop()
        for key in batch:
            if isinstance(batch[key], torch.Tensor):
                batch[key] = batch[key].cuda()

        if torch.distributed.get_rank() == 0:
            if not os.path.exists(os.path.join(args.save, "training_config.yaml")):
                configs = [OmegaConf.load(cfg) for cfg in args.base]
                config = OmegaConf.merge(*configs)
                os.makedirs(args.save, exist_ok=True)
                OmegaConf.save(config=config, f=os.path.join(args.save, "training_config.yaml"))
    else:
        batch = {"mp4": None, "fps": None, "num_frames": None, "txt": None}

    batch["global_step"] = args.iteration

    broad_cast_batch(batch)

    loss, loss_dict = model.shared_step(batch)

    return loss, loss_dict


if __name__ == "__main__":