File size: 1,999 Bytes
a20151e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | from modal import App, Image, Volume
import modal
import json
model_name = "Learnable"
app = App(f"Generalization Model {model_name} with ImageNet")
# Build image with all local dependencies added directly
image = (
Image.from_registry("nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04",add_python="3.10") # <--- This is the required 'tag'
.pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib",
"scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"])
.run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
.add_local_file("pyproject.toml", "/root/pyproject.toml")
.add_local_dir("src", "/root/src")
)
# Shared volume for saving outputs or checkpoints
volume1 = Volume.from_name("weights", create_if_missing=True)
volume2 = Volume.from_name("datasets", create_if_missing=True)
volume3 = Volume.from_name("plots", create_if_missing=True)
volume4 = Volume.from_name("results", create_if_missing=True)
@app.function(
image=image,
gpu="A100",
timeout=3600 * 24,
volumes={
"/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4,
},
)
def run_command():
import os
os.system("""
WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/imagenet/train_model.py \
--input-size 32 --data-set CIFAR10 --patch-size 4 --hidden-size 128 --num-hidden-layers 6 --warmup-epochs 5\
--num-attention-heads 4 --intermediate-size 512 --position-embeddings "rope" --num-labels 10\
--lr 5e-3 --epochs 50 --batch-size 128 --seed 0 --num-shared-experts 1 --num-routed-experts 0 --topk 0\
--wandb-project "LMC-Attention" --wandb-group "ViT-CIFAR10-FFN" --wandb-entity "fpt-team"\
--save-dir /root/weights/lmc/cifar10 --data-path /root/datasets/cifar10
""")
if __name__ == "__main__":
with app.run():
run_command.remote() |