diff --git a/Dockerfile b/Dockerfile index d73affc..bda78c0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,6 @@ # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch FROM nvcr.io/nvidia/pytorch:20.03-py3 +RUN pip install -U gsutil # Create working directory RUN mkdir -p /usr/src/app @@ -10,7 +11,6 @@ COPY . /usr/src/app # Install dependencies (pip or conda) #RUN pip install -r requirements.txt -RUN pip install -U gsutil # Copy weights #RUN python3 -c "from models import *; \ @@ -41,7 +41,7 @@ RUN pip install -U gsutil # Bash into running container # sudo docker container exec -it ba65811811ab bash -# python -c "from utils.utils import *; create_backbone('weights/last.pt')" && gsutil cp weights/backbone.pt gs://* +# python -c "from utils.utils import *; create_pretrained('weights/last.pt')" && gsutil cp weights/pretrained.pt gs://* # Bash into stopped container # sudo docker commit 6d525e299258 user/test_image && sudo docker run -it --gpus all --ipc=host -v "$(pwd)"/coco:/usr/src/coco --entrypoint=sh user/test_image diff --git a/utils/utils.py b/utils/utils.py index cf819c5..c33f41f 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -632,8 +632,8 @@ def strip_optimizer(f='weights/best.pt'): # from utils.utils import *; strip_op print('Optimizer stripped from %s' % f) -def create_backbone(f='weights/best.pt', s='weights/backbone.pt'): # from utils.utils import *; create_backbone() - # create backbone 's' from 'f' +def create_pretrained(f='weights/best.pt', s='weights/pretrained.pt'): # from utils.utils import *; create_pretrained() + # create pretrained checkpoint 's' from 'f' (create_pretrained(x, x) for x in glob.glob('./*.pt')) device = torch.device('cpu') x = torch.load(s, map_location=device) @@ -644,7 +644,7 @@ def create_backbone(f='weights/best.pt', s='weights/backbone.pt'): # from utils for p in x['model'].parameters(): p.requires_grad = True torch.save(x, s) - print('%s modified for backbone use and saved as %s' % (f, s)) + print('%s saved as pretrained checkpoint %s' % (f, s)) def coco_class_count(path='../coco/labels/train2014/'):