glenn-jocher commited on
Commit
3b06225
1 Parent(s): 0825cb7

update utils.create_pretrained()

Browse files
Files changed (2) hide show
  1. Dockerfile +2 -2
  2. utils/utils.py +3 -3
Dockerfile CHANGED
@@ -1,5 +1,6 @@
1
  # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
2
  FROM nvcr.io/nvidia/pytorch:20.03-py3
 
3
 
4
  # Create working directory
5
  RUN mkdir -p /usr/src/app
@@ -10,7 +11,6 @@ COPY . /usr/src/app
10
 
11
  # Install dependencies (pip or conda)
12
  #RUN pip install -r requirements.txt
13
- RUN pip install -U gsutil
14
 
15
  # Copy weights
16
  #RUN python3 -c "from models import *; \
@@ -41,7 +41,7 @@ RUN pip install -U gsutil
41
 
42
  # Bash into running container
43
  # sudo docker container exec -it ba65811811ab bash
44
- # python -c "from utils.utils import *; create_backbone('weights/last.pt')" && gsutil cp weights/backbone.pt gs://*
45
 
46
  # Bash into stopped container
47
  # 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
 
1
  # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
2
  FROM nvcr.io/nvidia/pytorch:20.03-py3
3
+ RUN pip install -U gsutil
4
 
5
  # Create working directory
6
  RUN mkdir -p /usr/src/app
 
11
 
12
  # Install dependencies (pip or conda)
13
  #RUN pip install -r requirements.txt
 
14
 
15
  # Copy weights
16
  #RUN python3 -c "from models import *; \
 
41
 
42
  # Bash into running container
43
  # sudo docker container exec -it ba65811811ab bash
44
+ # python -c "from utils.utils import *; create_pretrained('weights/last.pt')" && gsutil cp weights/pretrained.pt gs://*
45
 
46
  # Bash into stopped container
47
  # 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
utils/utils.py CHANGED
@@ -632,8 +632,8 @@ def strip_optimizer(f='weights/best.pt'): # from utils.utils import *; strip_op
632
  print('Optimizer stripped from %s' % f)
633
 
634
 
635
- def create_backbone(f='weights/best.pt', s='weights/backbone.pt'): # from utils.utils import *; create_backbone()
636
- # create backbone 's' from 'f'
637
  device = torch.device('cpu')
638
  x = torch.load(s, map_location=device)
639
 
@@ -644,7 +644,7 @@ def create_backbone(f='weights/best.pt', s='weights/backbone.pt'): # from utils
644
  for p in x['model'].parameters():
645
  p.requires_grad = True
646
  torch.save(x, s)
647
- print('%s modified for backbone use and saved as %s' % (f, s))
648
 
649
 
650
  def coco_class_count(path='../coco/labels/train2014/'):
 
632
  print('Optimizer stripped from %s' % f)
633
 
634
 
635
+ def create_pretrained(f='weights/best.pt', s='weights/pretrained.pt'): # from utils.utils import *; create_pretrained()
636
+ # create pretrained checkpoint 's' from 'f' (create_pretrained(x, x) for x in glob.glob('./*.pt'))
637
  device = torch.device('cpu')
638
  x = torch.load(s, map_location=device)
639
 
 
644
  for p in x['model'].parameters():
645
  p.requires_grad = True
646
  torch.save(x, s)
647
+ print('%s saved as pretrained checkpoint %s' % (f, s))
648
 
649
 
650
  def coco_class_count(path='../coco/labels/train2014/'):