Dockerized Version

.gitignore

@@ -0,0 +1 @@
+output/*

README.md

@@ -109,3 +109,19 @@ sample(
 )
 ```
 Visualization results are stored at `output/jet.png`
+
+### Run docker
+
+Execute from ./jet-pytorch/
+
+Available cards: `RTX4090`
+
+#### Create the docker
+```bash
+docker build -f dockerfiles/your_card/Dockerfile -t jet .
+```
+
+#### Run the docker
+```bash
+docker run -it --gpus '"device=0"' -v $(pwd)/output:/app/output --ipc=host jet
+```

docker_src/docker_start.sh

@@ -0,0 +1,13 @@
+chmod +x docker_src/docker_test.py
+chmod +x docker_src/docker_train.py
+
+python docker_src/docker_test.py \
+    --repo_id "de-Rodrigo/jet-mnist" \
+    --model_name "jet_mnist"
+
+# python docker_src/docker_train.py \
+#     --dataset_name "ylecun/mnist" \
+#     --wandb_entity "ciclab-comillas" \
+#     --wandb_project "jet" \
+#     --wandb_run_name "mnist" \
+#     --hf_repo_id "de-Rodrigo/jet-mnist"

docker_src/docker_test.py

@@ -0,0 +1,70 @@
+from torch.distributions import Normal
+from jet_pytorch import Jet
+from torchvision.utils import save_image
+import os
+from jet_pytorch.util import get_pretrained
+import argparse
+from huggingface_hub import login
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo_id", type=str, default=None)
+    parser.add_argument("--model_name", type=str, default=None)
+    args = parser.parse_args()
+
+    repo_id = args.repo_id
+    model_name = args.model_name
+
+    login(token=os.getenv("HUGGINGFACE_HUB_TOKEN"))
+
+    # Config Jet model
+    jet_config = dict(
+        patch_size=4,
+        patch_dim=48,
+        n_patches=256,
+        coupling_layers=32,
+        block_depth=2,
+        block_width=512,
+        num_heads=8,
+        scale_factor=2.0,
+        coupling_types=(
+            "channels", "channels",
+            "channels", "channels",
+            "spatial",
+        ),
+        spatial_coupling_projs=(
+            "checkerboard", "checkerboard-inv",
+            "vstripes", "vstripes-inv",
+            "hstripes", "hstripes-inv",
+        )
+    )
+
+    model = Jet(**jet_config)
+    weights = get_pretrained(repo_id, model_name)
+    model.load_state_dict(weights)
+
+    # Generate latent vectors z ~ N(0, 1)
+    batch_size = 16
+    n_patches = 256
+    patch_dim = 48
+    z = Normal(0, 1).sample((batch_size, n_patches, patch_dim))
+
+    # Rebuild imgs from z
+    img, _ = model.inverse(z)  # img: (B, H, W, C)
+
+    # Reformat (B, C, H, W)
+    img = img.permute(0, 3, 1, 2)
+
+    # Normalize imgs [0,1]
+    img_min = img.amin(dim=(1, 2, 3), keepdim=True)
+    img_max = img.amax(dim=(1, 2, 3), keepdim=True)
+    img = (img - img_min) / (img_max - img_min + 1e-6)
+    img = img.clamp(0, 1)
+
+    # Save imgs
+    os.makedirs("output", exist_ok=True)
+    for i in range(img.size(0)):
+        save_image(img[i], f"output/sample_{i:03d}.png")
+    save_image(img, "output/grid.png", nrow=4)

docker_src/docker_train.py

@@ -0,0 +1,116 @@
+import torchvision.io
+from torchvision.io import read_image as original_read_image
+
+# TODO Fix this properly
+def safe_read_image(path):
+    return original_read_image(str(path))
+
+# Monkey patch global
+import torchvision
+torchvision.io.read_image = safe_read_image
+
+import argparse
+import os
+from tqdm import tqdm
+from jet_pytorch.train import train
+from datasets import load_dataset
+import wandb
+from huggingface_hub import login
+
+
+def save_hf_dataset_to_disk(hf_dataset, output_dir, percentage=0.01):
+    os.makedirs(output_dir, exist_ok=True)
+    subset = hf_dataset.select(range(int(len(hf_dataset) * percentage)))
+    for i, sample in tqdm(enumerate(subset), total=len(subset)):    
+        img = sample["image"]
+        img = sample["image"].convert("RGB").resize((64, 64))
+        img.save(os.path.join(output_dir, f"{i}.png"))
+
+
+def get_hf_dataset():
+
+    dataset = load_dataset(dataset_name)
+    save_hf_dataset_to_disk(dataset["train"], f"./{dataset_name}_train")
+    save_hf_dataset_to_disk(dataset["test"], f"./{dataset_name}_valid")
+
+
+def train_jet():
+
+    wandb.login(key=os.getenv("WANDB_API_KEY"))
+    wandb.init(
+        project=wandb_project,
+        name=wandb_run_name,
+        entity=wandb_entity,
+        config={
+            "batch_size": 64,
+            "accumulate_steps": 16,
+            "epochs": 50,
+            "learning_rate": 3e-4,
+            "patch_size": 4,
+            "patch_dim": 48,
+            "n_patches": 256,
+            "coupling_layers": 32,
+        },
+    )
+
+    login(token=os.getenv("HUGGINGFACE_HUB_TOKEN"))
+
+    jet_config = dict(
+        patch_size=4,
+        patch_dim=48,
+        n_patches=256,
+        coupling_layers=32,
+        block_depth=2,
+        block_width=512,
+        num_heads=8,
+        scale_factor=2.0,
+        coupling_types=(
+            "channels", "channels",
+            "channels", "channels",
+            "spatial",
+        ),
+        spatial_coupling_projs=(
+            "checkerboard", "checkerboard-inv",
+            "vstripes", "vstripes-inv",
+            "hstripes", "hstripes-inv",
+        )
+    )
+
+
+    train(
+        jet_config=jet_config,
+        batch_size=32,
+        accumulate_steps=16,
+        device="cuda:0",
+        epochs=250,
+        warmup_percentage=0.1,
+        max_grad_norm=1.0,
+        learning_rate=3e-4,
+        weight_decay=1e-5,
+        adam_betas=(0.9, 0.95),
+        images_path_train=f"./{dataset_name}_train",
+        images_path_valid= f"./{dataset_name}_valid",
+        num_workers=8,
+        checkpoint_path="jet.pt",
+        hf_repo_id=hf_repo_id
+    )
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_name", type=str)
+    parser.add_argument("--wandb_entity", type=str)
+    parser.add_argument("--wandb_project", type=str)
+    parser.add_argument("--wandb_run_name", type=str)
+    parser.add_argument("--hf_repo_id", type=str, default=None)
+    args = parser.parse_args()
+
+    dataset_name = args.dataset_name
+    wandb_entity = args.wandb_entity
+    wandb_project = args.wandb_project
+    wandb_run_name = args.wandb_run_name
+    hf_repo_id = args.hf_repo_id
+
+    get_hf_dataset()
+    train_jet()

dockerfiles/rtx4090/Dockerfile

@@ -0,0 +1,34 @@
+FROM nvidia/cuda:11.8.0-base-ubuntu22.04
+
+RUN apt-get update && apt-get install -y \
+        python3.10 python3-pip libjpeg-dev zlib1g-dev git \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1
+WORKDIR /app
+
+RUN pip install --upgrade pip
+
+RUN pip install --no-cache-dir "numpy<2"
+
+RUN pip install --no-cache-dir \
+        --extra-index-url https://download.pytorch.org/whl/cu118 \
+        torch==2.1.2+cu118 \
+        torchvision==0.16.2+cu118
+
+COPY dockerfiles/rtx4090/requirements.txt /tmp/requirements.txt
+RUN pip install --no-cache-dir -r /tmp/requirements.txt
+
+RUN pip check
+
+RUN ln -sf /usr/bin/python3.10 /usr/bin/python
+
+ENV WANDB_API_KEY=<your_token>
+ENV HUGGINGFACE_HUB_TOKEN=<your_token>
+
+COPY . .
+
+ENV PYTHONPATH=/app
+
+RUN chmod +x docker_src/docker_start.sh
+CMD ["bash", "docker_src/docker_start.sh"]

dockerfiles/rtx4090/requirements.txt

@@ -0,0 +1,6 @@
+einops
+fire
+safetensors
+huggingface-hub==0.14.0
+datasets
+wandb

jet_pytorch/train.py

@@ -7,6 +7,9 @@
 from torch.utils.data import DataLoader
 from torchvision.transforms import v2
 from tqdm import tqdm
+import wandb
+from huggingface_hub import upload_file
+from safetensors.torch import save_file
 
 from jet_pytorch import Jet
 from jet_pytorch.util import bits_per_dim
@@ -30,6 +33,7 @@ def train(
     num_workers=16,
     device="cuda:0",
     checkpoint_path="jet_imagenet.pt",
+    hf_repo_id = None
 ):
     t.set_float32_matmul_precision("medium")
 
@@ -87,6 +91,8 @@ def train(
     log_nll = 0.0
     log_logdet = 0.0
 
+    best_val_loss = float("inf")
+
     for epoch in range(epochs):
         pbar = tqdm(dataloader, total=len(dataloader) // accumulate_steps)
 
@@ -121,6 +127,14 @@ def train(
                     gn=gn,
                     lr=lr,
                 ))
+                wandb.log({
+                    "train/bpd": log_bpd,
+                    "train/nll": log_nll,
+                    "train/logdet": log_logdet,
+                    "train/grad_norm": gn,
+                    "train/lr": lr,
+                    "epoch": epoch,
+                })
                 steps = 0
                 log_bpd = 0.0
                 log_nll = 0.0
@@ -159,6 +173,28 @@ def train(
         val_nll /= len(val_dataloader)
         val_logdet /= len(val_dataloader)
         print(f"validation metrics - bpd: {val_bpd:.2f}, nll: {val_nll:.2f}, logdet: {val_logdet:.2f}")
+        wandb.log({
+            "val/bpd": val_bpd,
+            "val/nll": val_nll,
+            "val/logdet": val_logdet,
+            "epoch": epoch,
+        })
+
+        # Save model if it's the best so far
+        if val_bpd < best_val_loss:
+            print(f"New best validation loss: {val_bpd:.4f} (prev {best_val_loss:.4f})")
+            best_val_loss = val_bpd
+
+            state_dict = orig_model.state_dict()
+            save_file(state_dict, "jet_mnist.safetensors")
+
+            upload_file(
+                path_or_fileobj="jet_mnist.safetensors",
+                path_in_repo="jet_mnist.safetensors",
+                repo_id=hf_repo_id,
+                repo_type="model"
+            )
+
         model.train()
 
 

jet_pytorch/util.py

@@ -12,16 +12,20 @@
 from torchvision.io import read_image
 
 
-def get_pretrained():
-    pretrained_path = "models/jet_imagenet64x64_200m/jet_imagenet64x64_200m.safetensors"
+def get_pretrained(repo_id: str = None, model_name:str = None):
+    if model_name is None:
+        pretrained_path = "models/jet_imagenet64x64_200m/jet_imagenet64x64_200m.safetensors"
+    else:
+        pretrained_path = f"models/{model_name}/{model_name}.safetensors"
     os.makedirs("models", exist_ok=True)
 
     if not os.path.exists(pretrained_path):
         _ = hf_hub_download(
-            repo_id="btrude/jet_imagenet64x64_200m",
-            filename="jet_imagenet64x64_200m.safetensors",
-            local_dir="models/jet_imagenet64x64_200m",
+            repo_id=repo_id,
+            filename=f"{model_name}.safetensors",
+            local_dir=f"models/{model_name}",
         )
+
     return load_file(pretrained_path)