p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_lamap.py (python)
    handler_registry.add(
        "train_dataloader_state",
        grain.checkpoint.CheckpointSave,
        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
    )
    handler_registry.add(
        "train_dataloader_state",
        grain.checkpoint.CheckpointRestore,
        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
    )
    if args.val_data_dir:
        handler_registry.add(
            "val_dataloader_state",
            grain.checkpoint.CheckpointSave,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )
        handler_registry.add(
            "val_dataloader_state",
            grain.checkpoint.CheckpointRestore,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )
[content] train_lamap.py (python)
    step = 0
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add(
        "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
    )
    handler_registry.add(
        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
    )
[content] train_lamap.py (python)
    image_shape = (args.image_height, args.image_width, args.image_channels)
    train_iterator = create_dataloader_iterator(args.data_dir, image_shape, args.seq_len, args.batch_size, args.seed)
    if args.val_data_dir:
        val_iterator = create_dataloader_iterator(args.val_data_dir, image_shape, args.seq_len, args.batch_size, args.seed)
[tab] slurm/jobs/mihir/horeka/coinrun/train_tokenizer_single_gpu.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:20:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_1e-4
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_gt_actions_split/train
array_records_dir_val=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_gt_actions_split/val


job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=125 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=250 \
    --log_checkpoint_interval=250 \
    --log \
    --name=coinrun-tokenizer-dev-$slurm_job_id \
    --tags tokenizer coinrun dev \
    --entity instant-uv \
    --project jafar \
    --warmup_steps 0 \
    --wsd_decay_steps 0 \
    --num_steps 250 \
    --data_dir $array_records_dir_train
[content] train_lam.py (python)
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        args.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
[content] train_lamap.py (python)

@nnx.jit
def val_step(lam: LatentActionModel, inputs: dict) -> tuple[jax.Array, jax.Array, dict]:
    lam.eval()
    (loss, (recon, _, metrics)) = lam_loss_fn(lam, inputs, training=False)
    return loss, recon, metrics

def calculate_validation_metrics(val_dataloader, lam):
    step = 0
    loss_per_step = []
    metrics_per_step = []
    for videos in val_dataloader:
        inputs = dict(videos=videos)
        loss, recon, metrics = val_step(lam, inputs)
        loss_per_step.append(loss)
        metrics_per_step.append(metrics)
        step += 1
        if step > args.val_steps:
            break

    if step < args.val_steps:
        print(f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}")

    val_loss = np.mean(loss_per_step)
    val_metrics = {
        f"val_{key}": np.mean([m[key] for m in metrics_per_step])
        for key in metrics_per_step[0].keys()
    }
    return val_loss, val_metrics, inputs, recon
[content] train_lamap.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
    train_state = TrainState.create(apply_fn=lamap.apply, params=init_params, tx=tx)
[content] train_lamap.py (python)

    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
    inputs["videos"] = gt.astype(args.dtype)
    outputs = model(inputs, training=training)
    outputs["recon"] = outputs["recon"].astype(jnp.float32)
    gt_future_frames = gt[:, 1:]
    mse = jnp.square(gt_future_frames - outputs["recon"]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
    commitment_loss = jnp.square(
        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss

    # --- Compute validation metrics ---
    gt = gt_future_frames.clip(0, 1).reshape(-1, *gt_future_frames.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt, recon)).mean()
    count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
    index_counts = count_fn(jnp.arange(args.num_latents))
    metrics = dict(
        loss=loss,
        mse=mse,
        q_loss=q_loss,
        commitment_loss=commitment_loss,
        psnr=psnr,
        ssim=ssim,
        codebook_usage=(index_counts != 0).mean(),
    )
    return loss, (outputs["recon"], index_counts, metrics)
[content] train_lamap.py (python)

    def loss_fn(
        model: LatentActionModel,
    ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
        model.train()
        return lam_loss_fn(model, inputs, training=True)

    # --- Update model ---
    (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
        loss_fn, has_aux=True
    )(lam)
    optimizer.update(grads)