p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/tokenizer_knoms_time_measurement.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_500_shards'

# FIXME: the `time_measurement_interval` line is only for debugging and should be removed
srun python train_tokenizer.py \
    --time_measurement_interval 10 \
    --batch_size 48 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/tokenizer_knoms_time_measurement_grain.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/dummy_arrayrecords'

# FIXME: the `time_measurement_interval` line is only for debugging and should be removed
srun python train_tokenizer.py \
    --time_measurement_interval 10 \
    --batch_size 48 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] train_tokenizer.py (python)
            videos_sharding = NamedSharding(
                mesh, PartitionSpec("data", None, None, None, None)
            )
            videos = jax.make_array_from_process_local_data(videos_sharding, videos)
[content] train_tokenizer.py (python)
            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    wandb.log(
                        {
                            "loss": loss,
                            "step": step,
                            **metrics,
                        }
                    )
                if step % args.log_image_interval == 0:
                    gt_seq = inputs["videos"][0]
                    recon_seq = recon[0].clip(0, 1)
                    comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                    comparison_seq = einops.rearrange(
                        comparison_seq * 255, "t h w c -> h (t w) c"
                    )
                    # NOTE: Process-dependent control flow deliberately happens
                    # after indexing operation since it must not contain code
                    # sections that lead to cross-accelerator communication.
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[0])),
                            recon=wandb.Image(np.asarray(recon_seq[0])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
            if step % args.log_checkpoint_interval == 0:
                ckpt = {"model": train_state}
                orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
                save_args = orbax_utils.save_args_from_target(ckpt)
                orbax_checkpointer.save(
                    os.path.join(os.getcwd(), args.ckpt_dir, f"tokenizer_{ts}_{step}"),
                    ckpt,
                    save_args=save_args,
                )
            if step >= args.num_steps:
                break
[content] train_tokenizer.py (python)
            inputs = dict(videos=videos, rng=_rng, dropout_rng=_rng_dropout)
            train_state, loss, recon, metrics = train_step(train_state, inputs)
            if step % args.time_measurement_interval == 0:
                jax.block_until_ready(train_state)
                elapsed_time = (time.time() - start_time)
                avg_step_time = elapsed_time / args.time_measurement_interval
                print(f"Step {step}, loss: {loss}, avg step time: {avg_step_time:.2f}s")
                if args.log and jax.process_index() == 0:
                    wandb.log({"avg_step_time_s": avg_step_time, "step": step})
                start_time = time.time()
            else:
                print(f"Step {step}, loss: {loss}")