diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..77320b3 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +output/* diff --git a/README.md b/README.md index 0abb187..52841dc 100644 --- a/README.md +++ b/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 +``` \ No newline at end of file diff --git a/docker_src/docker_start.sh b/docker_src/docker_start.sh new file mode 100644 index 0000000..c80fa32 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/docker_src/docker_test.py b/docker_src/docker_test.py new file mode 100644 index 0000000..d144303 --- /dev/null +++ b/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) diff --git a/docker_src/docker_train.py b/docker_src/docker_train.py new file mode 100644 index 0000000..a58dc1b --- /dev/null +++ b/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() diff --git a/dockerfiles/rtx4090/Dockerfile b/dockerfiles/rtx4090/Dockerfile new file mode 100644 index 0000000..6d8d24e --- /dev/null +++ b/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= +ENV HUGGINGFACE_HUB_TOKEN= + +COPY . . + +ENV PYTHONPATH=/app + +RUN chmod +x docker_src/docker_start.sh +CMD ["bash", "docker_src/docker_start.sh"] diff --git a/dockerfiles/rtx4090/requirements.txt b/dockerfiles/rtx4090/requirements.txt new file mode 100644 index 0000000..5f80325 --- /dev/null +++ b/dockerfiles/rtx4090/requirements.txt @@ -0,0 +1,6 @@ +einops +fire +safetensors +huggingface-hub==0.14.0 +datasets +wandb \ No newline at end of file diff --git a/jet_pytorch/train.py b/jet_pytorch/train.py index 40cc96f..f9f0168 100644 --- a/jet_pytorch/train.py +++ b/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() diff --git a/jet_pytorch/util.py b/jet_pytorch/util.py index 90d0333..b0d6c37 100644 --- a/jet_pytorch/util.py +++ b/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)