search

mapbox / robosat

Semantic segmentation on aerial and satellite imagery. Extracts features such as: buildings, parking lots, roads, water, clouds
MIT License
2.02k stars 382 forks source link

Bakes pretrained model into docker image #151

Closed daniel-j-h closed 5 years ago

daniel-j-h commented 5 years ago

We use a pre-trained resnet in our encoder. The pre-trained weights are getting downloaded on-demand when we use the model for the first time. Usually this is fine when working without docker. With docker the container will get shut down after each command invokation which means we have to download the model weights every time we run a new docker run .. train command.

This changset instead pre-fetches the pre-trained weights and bakes them into the docker images at

~/.cache/torch/checkpoints/resnet50-19c8e357.pth

Then docker run .. train never has to download the weights on-demand, we are no longer dependent on external services, and everything works beautifully.

The downside is the docker image size will increase since the weights are ~100 MB. But because the nvidia docker image is a few GBs anyway I don't see a problem here.

@jacquestardie merge please! Tested this on my rig!

jqtrde commented 5 years ago

Looks good @daniel-j-h, thank you!