浏览代码

update utils.create_pretrained()

5.0
Glenn Jocher 4 年前
父节点
当前提交
3b062254a6
共有 2 个文件被更改,包括 5 次插入5 次删除
  1. +2
    -2
      Dockerfile
  2. +3
    -3
      utils/utils.py

+ 2
- 2
Dockerfile 查看文件

# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
FROM nvcr.io/nvidia/pytorch:20.03-py3 FROM nvcr.io/nvidia/pytorch:20.03-py3
RUN pip install -U gsutil


# Create working directory # Create working directory
RUN mkdir -p /usr/src/app RUN mkdir -p /usr/src/app


# Install dependencies (pip or conda) # Install dependencies (pip or conda)
#RUN pip install -r requirements.txt #RUN pip install -r requirements.txt
RUN pip install -U gsutil


# Copy weights # Copy weights
#RUN python3 -c "from models import *; \ #RUN python3 -c "from models import *; \


# Bash into running container # Bash into running container
# sudo docker container exec -it ba65811811ab bash # 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 # 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 # 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

+ 3
- 3
utils/utils.py 查看文件

print('Optimizer stripped from %s' % f) 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') device = torch.device('cpu')
x = torch.load(s, map_location=device) x = torch.load(s, map_location=device)


for p in x['model'].parameters(): for p in x['model'].parameters():
p.requires_grad = True p.requires_grad = True
torch.save(x, s) 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/'): def coco_class_count(path='../coco/labels/train2014/'):

正在加载...
取消
保存