p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/models/tokenizer.py (python)

    """
    Masked autoencoder for video reconstruction.
    """
    def __init__(self, model: TokenizerVQVAE, mask_ratio: float):
        self.model = model
        self.mask_ratio = mask_ratio

    def __call__(self, videos: jax.Array) -> jax.Array:
        return self.model(videos)
[content] jasmine/models/tokenizer.py (python)

        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.latent_dim,
            self.model_dim,
            self.ffn_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
[content] jasmine/models/tokenizer.py (python)

    def __init__(
        self,
        in_dim: int,
        model_dim: int,
        ffn_dim: int,
        latent_dim: int,
        num_latents: int,
        patch_size: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        codebook_dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
    ):
[content] jasmine/models/tokenizer.py (python)

        self.in_dim = in_dim
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.latent_dim = latent_dim
        self.num_latents = num_latents
        self.patch_size = patch_size
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.codebook_dropout = codebook_dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
[content] jasmine/models/tokenizer.py (python)

        H, W = batch["videos"].shape[2:4]
        videos_BTHWC = batch["videos"]
        outputs = self.vq_encode(videos_BTHWC, training)
        z_q_BTNL = outputs["z_q"]
        recon_BTHWC = self.decoder(z_q_BTNL)
        recon_BTHWC = recon_BTHWC.astype(jnp.float32)
        recon_BTHWC = nnx.sigmoid(recon_BTHWC)
        recon_BTHWC = recon_BTHWC.astype(self.dtype)
        recon_BTHWC = unpatchify(recon_BTHWC, self.patch_size, H, W)
        outputs["recon"] = recon_BTHWC
        return outputs
[content] jasmine/models/tokenizer.py (python)

        recon_BTNP = self.decoder(z_BTNL)
        recon_BTNP = recon_BTNP.astype(jnp.float32)
        recon_BTNP = nnx.sigmoid(recon_BTNP)
        recon_BTNP = recon_BTNP.astype(self.dtype)
        return unpatchify(recon_BTNP, self.patch_size, *video_hw)
[content] jasmine/train_tokenizer.py (python)

        tokenizer = TokenizerVQVAE(
            in_dim=args.image_channels,
            model_dim=args.model_dim,
            ffn_dim=args.ffn_dim,
            latent_dim=args.latent_dim,
            num_latents=args.num_latents,
            patch_size=args.patch_size,
            num_blocks=args.num_blocks,
            num_heads=args.num_heads,
            dropout=args.dropout,
            codebook_dropout=args.codebook_dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
            rngs=rngs,
        )
[content] jasmine/train_tokenizer.py (python)

    elif args.tokenizer_type == "mae":
        tokenizer = MaskedAutoEncoder(
            in_dim=args.image_channels,
            model_dim=args.model_dim,
            ffn_dim=args.ffn_dim,
            latent_dim=args.latent_dim,
            num_latents=args.num_latents,
            patch_size=args.patch_size,
            num_blocks=args.num_blocks,
            num_heads=args.num_heads,
            dropout=args.dropout,
            max_mask_ratio=args.max_mask_ratio,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
            rngs=rngs,
        )
[content] jasmine/train_tokenizer.py (python)
s.ffn_dim,
            latent_dim=args.latent_dim,
            num_latents=args.num_latents,
            patch_size=args.patch_size,
            num_blocks=args.num_blocks,
            num_heads=args.num_heads,
            dropout=args.dropout,
            codebook_dropout=args.codebook_dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
            rngs=rngs,
[content] jasmine/train_tokenizer.py (python)

            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
        else:
            loss = mse
[content] jasmine/train_tokenizer.py (python)

        gt_clipped = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
        recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        psnr = jnp.asarray(pix.psnr(gt_clipped, recon)).mean()
        ssim = jnp.asarray(pix.ssim(gt_clipped, recon)).mean()
        metrics = dict(
            loss=loss,
            mse=mse,
            psnr=psnr,
            ssim=ssim,
        )
[content] slurm/jobs/mihir/berlin/diffusion/breakout/train_tokenizer_mae.sh (shellscript)

#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --job-name=train_dyn_dit
[content] slurm/jobs/mihir/berlin/diffusion/breakout/train_tokenizer_mae.sh (shellscript)


CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

PYTHONUNBUFFERED=1

srun python jasmine/train_tokenizer.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --ckpt_dir $CHECKPOINT_DIR \
[tab] slurm/jobs/mihir/berlin/diffusion/breakout/train_tokenizer_mae_dev.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --job-name=train_tokenizer_mae

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

export PYTHONUNBUFFERED=1

srun python jasmine/train_tokenizer.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --tokenizer_type=mae \
    --patch_size=4 \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_tokenizer_mae_breakout_$slurm_job_id \
    --tags tokenizer mae breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[tab] slurm/jobs/mihir/berlin/diffusion/breakout/train_tokenizer_mae_dev copy.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --job-name=train_tokenizer_mae

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

export PYTHONUNBUFFERED=1

srun python jasmine/train_tokenizer.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --tokenizer_type=mae \
    --patch_size=4 \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_tokenizer_mae_breakout_$slurm_job_id \
    --tags tokenizer mae breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[content] slurm/jobs/mihir/berlin/diffusion/coinrun/train_tokenizer_mae.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --job-name=train_tokenizer_mae

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

srun python jasmine/train_tokenizer.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --tokenizer_type=mae \
    --patch_size=4 \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_tokenizer_mae_breakout_$slurm_job_id \
    --tags tokenizer mae breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \