p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/horeka/doom/resolution60x80/train_dyn_default_500k.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=48:00: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/doom/dynamics/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/doom/dynamics/maskgit/%x_%j.log
#SBATCH --job-name=dynamics_doom_60x80_500k
#SBATCH --exclude=hkn0735
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# --- signal trap to requeue job before timeout ---
requeue_job() {
    echo "[$(date)] caught sigusr1 (timeout warning), requeueing slurm job $SLURM_JOB_ID..."
    # optional: trigger checkpoint saving here
    # e.g., touch $checkpoint_dir/requeue_trigger
    scontrol requeue $SLURM_JOB_ID
    exit 0
}

trap requeue_job sigusr1

# set checkpoint flag based on restart count
restart_count=$(scontrol show job $SLURM_JOB_ID | grep -o 'Restarts=[0-9]*' | cut -d'=' -f2)

if [ $restart_count -eq 0 ]; then
    restore_ckpt_flag="--no-restore-ckpt"
else
    restore_ckpt_flag="--restore-ckpt"
fi



# 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_doom/doom_episodes_10m_60x80_fixed/train
array_records_dir_val=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_doom/doom_episodes_10m_60x80_fixed/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

# CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/doom/maskgit/$job_name/$slurm_job_id
CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/doom/maskgit/dynamics_doom_60x80_500k/500k-run
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/tokenizer/train_tokenizer_default_single_gpu_60x80/3547697

env | grep SLURM

srun python jasmine/train_dynamics.py \
    --save_ckpt \
    --restore_ckpt \
    --wandb_id $SLURM_JOB_ID \
    --num_actions=20 \
    --num_steps=500_000 \
    --wsd_decay_steps=50_000 \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=doom-dynamics-60x80-500k-$slurm_job_id \
    --image_height=60 \
    --image_width=80 \
    --tags doom dynamics maskgit default 60x80 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val &

child_pid=$!

wait $child_pid
[tab] slurm/jobs/mihir/horeka/doom/resolution60x80/sample.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#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/doom/dynamics/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/doom/dynamics/maskgit/%x_%j.log
#SBATCH --job-name=sample_dynamics_doom_60x80

# Unload modules that may interfere
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4

# Activate virtual environment
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_doom/doom_episodes_10m_60x80_fixed/test

CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/doom/maskgit/dynamics_doom_60x80/3551397

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=8 \
    --start_frame=4 \
    --image_height=60 \
    --image_width=80 \
    --num_actions=18 \
    --patch_size=16 \
    --output_dir=gifs/doom/action-prepend-branch/60x80/test/ $@
[content] jasmine/train_tokenizer_appendix-c.py (python)

            if dataloader_val and step % args.val_interval == 0:
                print("Calculating validation metrics...")
                val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(
                    dataloader_val, optimizer.model
                )
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                }
[content] slurm/jobs/mihir/horeka/coinrun/ablations/appendix-c.sh (shellscript)

# 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/coinrun/maskgit/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/train_tokenizer_default/3528955

env | grep SLURM

srun python jasmine/train_dynamics.py \
    --save_ckpt \
    $restore_ckpt_flag \
    --wandb_id $SLURM_JOB_ID \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=110 \
    --image_height=64 \
    --image_width=64 \
    --patch_size=16 \
    --max_lr=3e-5 \
    --max_noise_level=0 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=coinrun-dynamics-maskgit-no-noise-$slurm_job_id \
    --tags coinrun dynamics maskgit no-noise \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --eval_full_frame \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val &

child_pid=$!

wait $child_pid
[content] jasmine/train_dynamics.py (python)

            # --- Validation loss ---
            val_results = {}
            if dataloader_val and step % args.val_interval == 0:
                rng, _rng_mask_val = jax.random.split(rng, 2)
                print("Calculating validation metrics...")
                val_metrics, val_gt_batch, val_recon, val_recon_full_frame = (
                    calculate_validation_metrics(
                        dataloader_val, optimizer.model, _rng_mask_val
                    )
                )
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                    "full_frame": val_recon_full_frame,
                }
[tab] slurm/jobs/mihir/horeka/doom/resolution60x80/sample.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#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/doom/dynamics/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/doom/dynamics/maskgit/%x_%j.log
#SBATCH --job-name=sample_dynamics_doom_60x80

# Unload modules that may interfere
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4

# Activate virtual environment
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_doom/doom_episodes_10m_60x80_fixed/test

CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/doom/maskgit/dynamics_doom_60x80_500k/500k-run

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=8 \
    --start_frame=4 \
    --image_height=60 \
    --image_width=80 \
    --num_actions=18 \
    --patch_size=16 \
    --output_dir=gifs/doom/action-prepend-branch/60x80/500k-test/ $@
[content] jasmine/train_dynamics.py (python)
    @nnx.jit(donate_argnums=0)
    def train_step(
        optimizer: nnx.ModelAndOptimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            model.train()
            return dynamics_loss_fn(model, inputs)

        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            optimizer.model
        )
        optimizer.update(grads)
        if args.log_gradients:
            metrics["gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["dynamics"]
            )
        return loss, recon, metrics
[content] jasmine/train_dynamics.py (python)
    # --- TRAIN LOOP ---
    dataloader_train = (
        {
            "videos": jax.make_array_from_process_local_data(
                videos_sharding, local_data=elem["videos"]
            ),
            "actions": (
                jax.make_array_from_process_local_data(
                    actions_sharding, elem["actions"]
                )
                if args.use_gt_actions
                else None
            ),
        }
        for elem in train_iterator
    )
[content] jasmine/train_dynamics.py (python)
    if jax.process_index() == 0:
        first_batch = next(dataloader_train)
        first_batch["rng"] = rng  # type: ignore
        compiled = train_step.lower(optimizer, first_batch).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader_train = itertools.chain([first_batch], dataloader_train)
[content] jasmine/train_dynamics.py (python)
    while step < args.num_steps:
        for batch in dataloader_train:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            batch["rng"] = _rng_mask
            loss, recon, metrics = train_step(optimizer, batch)
            if step == first_step:
                print_mem_stats("After params initialized")
            step += 1
[content] jasmine/train_dynamics.py (python)
            # --- Validation loss ---
            val_results = {}
            if dataloader_val and step % args.val_interval == 0:
                rng, _rng_mask_val = jax.random.split(rng, 2)
                print("Calculating validation metrics...")
                val_metrics, val_gt_batch, val_recon, val_recon_full_frame = (
                    calculate_validation_metrics(
                        dataloader_val, optimizer.model, _rng_mask_val
                    )
                )
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                    "full_frame": val_recon_full_frame,
                }
[content] appendix_c_nodes.md (markdown)
Step 200000, validation loss: 0.7587440609931946
{'val_codebook_usage_lam': np.float64(1.0), 'val_codebook_usage_tokenizer': np.float64(0.8475988051470589), 'val_cross_entropy_loss': np.float64(0.758744008400861), 'val_entropy': np.float64(0.8734586110302046), 'val_masked_token_top16_accuracy': np.float64(0.9815394200530707), 'val_masked_token_top1_accuracy': np.float64(0.772024557870977), 'val_masked_token_top2_accuracy': np.float64(0.8777434533717585), 'val_masked_token_top5_accuracy': np.float64(0.9463883007273954), 'val_psnr': np.float64(18.16088799869313), 'val_select_logit': np.float64(15.902265399110084), 'val_select_p': np.float64(0.7548150642245424), 'val_ssim': np.float64(0.6326266468739977), 'val_total_loss': np.float64(0.758744008400861), 'val_z_loss': np.float64(310.9372079886642), 'val_loss': np.float32(0.75874406), 'val_full_frame_codebook_usage_lam': np.float64(1.0), 'val_full_frame_codebook_usage_tokenizer': np.float64(0.2629250919117647), 'val_full_frame_cross_entropy_loss': np.float64(20.045630099726658), 'val_full_frame_entropy': np.float64(1.0269374473422181), 'val_full_frame_masked_token_top16_accuracy': np.float64(0.7251559460864347), 'val_full_frame_masked_token_top1_accuracy': np.float64(0.38186273446270064), 'val_full_frame_masked_token_top2_accuracy': np.float64(0.47075533224087135), 'val_full_frame_masked_token_top5_accuracy': np.float64(0.5903074694614784), 'val_full_frame_psnr': np.float64(16.088213939292757), 'val_full_frame_select_logit': np.float64(84.10442965638404), 'val_full_frame_select_p': np.float64(0.7601301389582017), 'val_full_frame_ssim': np.float64(0.5480527725874209), 'val_full_frame_z_loss': np.float64(108920.47901348039), 'val_full_frame_loss': np.float32(20.045631)}
[content] appendix_c_nodes.md (markdown)
Step 200000, validation loss: 1.9375224113464355
{'val_codebook_usage_lam': np.float64(1.0), 'val_codebook_usage_tokenizer': np.float64(0.8475988051470589), 'val_cross_entropy_loss': np.float64(1.9375225586049698), 'val_entropy': np.float64(2.236036305334054), 'val_masked_token_top16_accuracy': np.float64(0.8798762057341781), 'val_masked_token_top1_accuracy': np.float64(0.5657584152969659), 'val_masked_token_top2_accuracy': np.float64(0.689856609877418), 'val_masked_token_top5_accuracy': np.float64(0.798516615933063), 'val_psnr': np.float64(15.530764074886546), 'val_select_logit': np.float64(10.11456164191751), 'val_select_p': np.float64(0.501306425707013), 'val_ssim': np.float64(0.5097568672077328), 'val_total_loss': np.float64(1.9375225586049698), 'val_z_loss': np.float64(149.545166015625), 'val_loss': np.float32(1.9375224), 'val_full_frame_codebook_usage_lam': np.float64(1.0), 'val_full_frame_codebook_usage_tokenizer': np.float64(0.2629250919117647), 'val_full_frame_cross_entropy_loss': np.float64(20.045630099726658), 'val_full_frame_entropy': np.float64(1.0269374473422181), 'val_full_frame_masked_token_top16_accuracy': np.float64(0.7251559460864347), 'val_full_frame_masked_token_top1_accuracy': np.float64(0.38186273446270064), 'val_full_frame_masked_token_top2_accuracy': np.float64(0.47075533224087135), 'val_full_frame_masked_token_top5_accuracy': np.float64(0.5903074694614784), 'val_full_frame_psnr': np.float64(16.088213939292757), 'val_full_frame_select_logit': np.float64(84.10442965638404), 'val_full_frame_select_p': np.float64(0.7601301389582017), 'val_full_frame_ssim': np.float64(0.5480527725874209), 'val_full_frame_z_loss': np.float64(108920.47901348039), 'val_full_frame_loss': np.float32(20.045631)}
[content] appendix_c_nodes.md (markdown)
Step 200000, validation loss: 4.311364650726318
{'val_codebook_usage_lam': np.float64(1.0), 'val_codebook_usage_tokenizer': np.float64(0.8475988051470589), 'val_cross_entropy_loss': np.float64(4.3113649312187645), 'val_entropy': np.float64(4.482920328776042), 'val_masked_token_top16_accuracy': np.float64(0.5900948678745943), 'val_masked_token_top1_accuracy': np.float64(0.24632565764819875), 'val_masked_token_top2_accuracy': np.float64(0.33755794225954544), 'val_masked_token_top5_accuracy': np.float64(0.4589417284610225), 'val_psnr': np.float64(12.573247198965035), 'val_select_logit': np.float64(6.313630515453863), 'val_select_p': np.float64(0.16665662007004606), 'val_ssim': np.float64(0.3333607175770928), 'val_total_loss': np.float64(4.3113649312187645), 'val_z_loss': np.float64(77.57380646350337), 'val_loss': np.float32(4.3113647), 'val_full_frame_codebook_usage_lam': np.float64(1.0), 'val_full_frame_codebook_usage_tokenizer': np.float64(0.2629250919117647), 'val_full_frame_cross_entropy_loss': np.float64(20.045630099726658), 'val_full_frame_entropy': np.float64(1.0269374473422181), 'val_full_frame_masked_token_top16_accuracy': np.float64(0.7251559460864347), 'val_full_frame_masked_token_top1_accuracy': np.float64(0.38186273446270064), 'val_full_frame_masked_token_top2_accuracy': np.float64(0.47075533224087135), 'val_full_frame_masked_token_top5_accuracy': np.float64(0.5903074694614784), 'val_full_frame_psnr': np.float64(16.088213939292757), 'val_full_frame_select_logit': np.float64(84.10442965638404), 'val_full_frame_select_p': np.float64(0.7601301389582017), 'val_full_frame_ssim': np.float64(0.5480527725874209), 'val_full_frame_z_loss': np.float64(108920.47901348039), 'val_full_frame_loss': np.float32(20.045631)}
[content] appendix_c_nodes.md (markdown)
Step 200000, validation loss: 6.960869312286377
{'val_codebook_usage_lam': np.float64(1.0), 'val_codebook_usage_tokenizer': np.float64(0.8475988051470589), 'val_cross_entropy_loss': np.float64(6.960869602128571), 'val_entropy': np.float64(5.414477582071342), 'val_masked_token_top16_accuracy': np.float64(0.18385726447198905), 'val_masked_token_top1_accuracy': np.float64(0.0353936998940566), 'val_masked_token_top2_accuracy': np.float64(0.05846153243499644), 'val_masked_token_top5_accuracy': np.float64(0.10274799825514064), 'val_psnr': np.float64(10.914788994134641), 'val_select_logit': np.float64(5.507573454987769), 'val_select_p': np.float64(0.07029049901985655), 'val_ssim': np.float64(0.2510850712949154), 'val_total_loss': np.float64(6.960869602128571), 'val_z_loss': np.float64(69.57815611596202), 'val_loss': np.float32(6.9608693), 'val_full_frame_codebook_usage_lam': np.float64(1.0), 'val_full_frame_codebook_usage_tokenizer': np.float64(0.2629250919117647), 'val_full_frame_cross_entropy_loss': np.float64(20.045630099726658), 'val_full_frame_entropy': np.float64(1.0269374473422181), 'val_full_frame_masked_token_top16_accuracy': np.float64(0.7251559460864347), 'val_full_frame_masked_token_top1_accuracy': np.float64(0.38186273446270064), 'val_full_frame_masked_token_top2_accuracy': np.float64(0.47075533224087135), 'val_full_frame_masked_token_top5_accuracy': np.float64(0.5903074694614784), 'val_full_frame_psnr': np.float64(16.088213939292757), 'val_full_frame_select_logit': np.float64(84.10442965638404), 'val_full_frame_select_p': np.float64(0.7601301389582017), 'val_full_frame_ssim': np.float64(0.5480527725874209), 'val_full_frame_z_loss': np.float64(108920.47901348039), 'val_full_frame_loss': np.float32(20.045631)}
[content] jasmine/train_dynamics_appendix-c.py (python)
<                 batch=inputs,
<                 seq_len=args.seq_len,
<                 noise_level=0.0,
<                 temperature=args.val_temperature,
<                 sample_argmax=args.val_sample_argmax,
<                 maskgit_steps=args.val_maskgit_steps,
[tab] slurm/jobs/mihir/horeka/coinrun/ablations/appendix-c.sh (shellscript)

# 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/coinrun/maskgit/$job_name/$slurm_job_id

# main
CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/coinrun/maskgit/train_dynamics_maskgit/3533426

mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/train_tokenizer_default/3528955

env | grep SLURM

srun python jasmine/train_dynamics_appendix-c.py \
    --save_ckpt \
    --restore_ckpt \
    --wandb_id $SLURM_JOB_ID \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=110 \
    --image_height=64 \
    --image_width=64 \
    --patch_size=16 \
    --max_lr=3e-5 \
    --no-log \
    --no-eval-full-frame \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --eval_full_frame \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val
[tab] slurm/jobs/mihir/horeka/doom/resolution120x160/train_tokenizer_default_1node_requeue.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/doom/tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/doom/tokenizer/%x_%j.log
#SBATCH --exclude=hkn0735
#SBATCH --job-name=train_tokenizer_default_single_node_120x160

# 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_doom/doom_episodes_10m_120x160_fixed/train
array_records_dir_val=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_doom/doom_episodes_10m_120x160_fixed/val


job_name=$SLURM_JOB_NAME
slurm_job_id=3546530

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

env | grep SLURM

srun python jasmine/train_tokenizer_full_precision.py \
    --save_ckpt \
    --restore_ckpt \
    --image_height=120 \
    --image_width=160 \
    --no-use-flash-attention \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --name=doom-tokenizer-default-120x160-$slurm_job_id-requeue \
    --tags tokenizer doom default 120x160 \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val
[content] appendix_c_nodes.md (markdown)
Step 200000, validation loss: 0.7374703884124756
{'val_codebook_usage_lam': np.float64(1.0), 'val_codebook_usage_tokenizer': np.float64(0.8475988051470589), 'val_cross_entropy_loss': np.float64(0.7374703171206456), 'val_entropy': np.float64(0.8575162127906201), 'val_masked_token_top16_accuracy': np.float64(0.9825415529456794), 'val_masked_token_top1_accuracy': np.float64(0.776269306154812), 'val_masked_token_top2_accuracy': np.float64(0.8820458533717137), 'val_masked_token_top5_accuracy': np.float64(0.949582414299834), 'val_psnr': np.float64(22.455391640756645), 'val_select_logit': np.float64(15.983110203462488), 'val_select_p': np.float64(0.7558641293469597), 'val_ssim': np.float64(0.6962924330842262), 'val_total_loss': np.float64(0.7374703171206456), 'val_z_loss': np.float64(313.3120548023897), 'val_loss': np.float32(0.7374704), 'val_full_frame_codebook_usage_lam': np.float64(1.0), 'val_full_frame_codebook_usage_tokenizer': np.float64(0.2629250919117647), 'val_full_frame_cross_entropy_loss': np.float64(1.7930428303924262), 'val_full_frame_entropy': np.float64(0.252041938550332), 'val_full_frame_masked_token_top16_accuracy': np.float64(0.9929032676360187), 'val_full_frame_masked_token_top1_accuracy': np.float64(0.8329210959228814), 'val_full_frame_masked_token_top2_accuracy': np.float64(0.9222704674683365), 'val_full_frame_masked_token_top5_accuracy': np.float64(0.9736630788036421), 'val_full_frame_psnr': np.float64(28.751719717885933), 'val_full_frame_select_logit': np.float64(100.21908793729894), 'val_full_frame_select_p': np.float64(0.912218968073527), 'val_full_frame_ssim': np.float64(0.8437704747798396), 'val_full_frame_z_loss': np.float64(51192.01493566176), 'val_full_frame_loss': np.float32(1.7930428)}
[content] appendix_c_nodes.md (markdown)
Step 200000, validation loss: 3.6271512508392334
{'val_codebook_usage_lam': np.float64(1.0), 'val_codebook_usage_tokenizer': np.float64(0.8475988051470589), 'val_cross_entropy_loss': np.float64(3.6271509796965358), 'val_entropy': np.float64(4.294355560751522), 'val_masked_token_top16_accuracy': np.float64(0.7079328228445614), 'val_masked_token_top1_accuracy': np.float64(0.29852578511425093), 'val_masked_token_top2_accuracy': np.float64(0.41297777552230686), 'val_masked_token_top5_accuracy': np.float64(0.5610486105376599), 'val_psnr': np.float64(14.66877135108499), 'val_select_logit': np.float64(5.881039871889002), 'val_select_p': np.float64(0.18452119885706433), 'val_ssim': np.float64(0.4027574781109305), 'val_total_loss': np.float64(3.6271509796965358), 'val_z_loss': np.float64(67.83754580628639), 'val_loss': np.float32(3.6271513), 'val_full_frame_codebook_usage_lam': np.float64(1.0), 'val_full_frame_codebook_usage_tokenizer': np.float64(0.2629250919117647), 'val_full_frame_cross_entropy_loss': np.float64(1.7930428303924262), 'val_full_frame_entropy': np.float64(0.252041938550332), 'val_full_frame_masked_token_top16_accuracy': np.float64(0.9929032676360187), 'val_full_frame_masked_token_top1_accuracy': np.float64(0.8329210959228814), 'val_full_frame_masked_token_top2_accuracy': np.float64(0.9222704674683365), 'val_full_frame_masked_token_top5_accuracy': np.float64(0.9736630788036421), 'val_full_frame_psnr': np.float64(28.751719717885933), 'val_full_frame_select_logit': np.float64(100.21908793729894), 'val_full_frame_select_p': np.float64(0.912218968073527), 'val_full_frame_ssim': np.float64(0.8437704747798396), 'val_full_frame_z_loss': np.float64(51192.01493566176), 'val_full_frame_loss': np.float32(1.7930428)}
[content] appendix_c_nodes.md (markdown)
Step 200000, validation loss: 6.788418769836426
{'val_codebook_usage_lam': np.float64(1.0), 'val_codebook_usage_tokenizer': np.float64(0.8475988051470589), 'val_cross_entropy_loss': np.float64(6.788418115354052), 'val_entropy': np.float64(5.4373855964810245), 'val_masked_token_top16_accuracy': np.float64(0.19955201680753745), 'val_masked_token_top1_accuracy': np.float64(0.032450725634892784), 'val_masked_token_top2_accuracy': np.float64(0.05669430161223692), 'val_masked_token_top5_accuracy': np.float64(0.10574950832946628), 'val_psnr': np.float64(10.75548319723092), 'val_select_logit': np.float64(4.335007929334454), 'val_select_p': np.float64(0.06822160222366744), 'val_ssim': np.float64(0.25764136659164055), 'val_total_loss': np.float64(6.788418115354052), 'val_z_loss': np.float64(54.11139491960114), 'val_loss': np.float32(6.788419), 'val_full_frame_codebook_usage_lam': np.float64(1.0), 'val_full_frame_codebook_usage_tokenizer': np.float64(0.2629250919117647), 'val_full_frame_cross_entropy_loss': np.float64(1.7930428303924262), 'val_full_frame_entropy': np.float64(0.252041938550332), 'val_full_frame_masked_token_top16_accuracy': np.float64(0.9929032676360187), 'val_full_frame_masked_token_top1_accuracy': np.float64(0.8329210959228814), 'val_full_frame_masked_token_top2_accuracy': np.float64(0.9222704674683365), 'val_full_frame_masked_token_top5_accuracy': np.float64(0.9736630788036421), 'val_full_frame_psnr': np.float64(28.751719717885933), 'val_full_frame_select_logit': np.float64(100.21908793729894), 'val_full_frame_select_p': np.float64(0.912218968073527), 'val_full_frame_ssim': np.float64(0.8437704747798396), 'val_full_frame_z_loss': np.float64(51192.01493566176), 'val_full_frame_loss': np.float32(1.7930428)}