Dockerized Version#1
Conversation
|
Hey thanks for your interest in this work, I will try to get this reviewed ASAP but before that I am curious if you have a publicly available link to wandb for a successful run of this code? |
|
Hey @btrude, thank you for your reply! Of course, please find the link to a quick training session: https://wandb.ai/ciclab-comillas/jet?nw=nwuseriderodrigo |
|
Thanks, I'm glad to see that it is robust to non-imagenet data. I will try to get the review in later today or tomorrow... |
btrude
left a comment
btrude
left a comment
There was a problem hiding this comment.
As far as I can tell, the only things making the /4090/ code Ada-specific are the use of pre-Blackwell base image and pytorch versions. I think this should just be updated to use the latest pytorch and latest CUDA base images which should allow any currently released nvidia GPU to run this code without modification. I don't see memory usage in your public wandb runs, but I'm assuming its <24gb so these changes could make for a nice MNIST configuration for those interested in running the code but not wanting to train on imagenet (which takes a few days on a 4090 iirc) or who don't have a 24gb+ GPU.
My apologies if these seem like nitpicks, I don't think any of these comments truly matter given the general lack of interest in this work, these are just my first thoughts when looking through the changes so please feel free to push back if anything seems wrong or excessive. Thanks!
| @@ -0,0 +1,34 @@ | |||
| FROM nvidia/cuda:11.8.0-base-ubuntu22.04 | |||
There was a problem hiding this comment.
Can we use a Blackwell compatible image here (ie >= 12.8)?
| gn=gn, | ||
| lr=lr, | ||
| )) | ||
| wandb.log({ |
There was a problem hiding this comment.
What happens when this is run without a wandb api key passed in? I wouldn't want to make wandb a hard requirement for this script.
| ENV WANDB_API_KEY=<your_token> | ||
| ENV HUGGINGFACE_HUB_TOKEN=<your_token> |
There was a problem hiding this comment.
Writing a dockerfile like this leaves your secrets exposed in the final built image. I would consider re-writing this with build secrets if there is any risk of a 3rd party having access to the built image.
| 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 |
There was a problem hiding this comment.
I think ideally this would be two scripts with each script accepting command line arguments that are then passed to their respective python scripts: e.g:
chmod +x docker_src/docker_test.py
python docker_src/docker_test.py \
--repo_id $0 \
--model_name $1Invoked as:
./docker_test.sh "de-Rodrigo/jet-mnist" "jet_mnist"| state_dict = orig_model.state_dict() | ||
| save_file(state_dict, "jet_mnist.safetensors") | ||
|
|
||
| upload_file( |
There was a problem hiding this comment.
Same question as above but for the huggingface api key.
|
@btrude, happy to consider your comments. I'll try to address them by the end of the week. In the meantime, are you one of the authors of the original implementation? Thank you! |
I am not. The original paper's code was written in JAX and so I put this together to enable integration with pytorch models (specifically for training "jetformers" with arbitrary backbones). |
2 participants
Run the project in a Docker container