提交 a6b0f49c 编写于 作者: littletomatodonkey's avatar littletomatodonkey

add pretrain list show

上级 eef1e879
...@@ -24,7 +24,6 @@ import tqdm ...@@ -24,7 +24,6 @@ import tqdm
import zipfile import zipfile
from ppcls.modeling import similar_architectures from ppcls.modeling import similar_architectures
from ppcls.utils.check import check_architecture
from ppcls.utils import logger from ppcls.utils import logger
__all__ = ['get'] __all__ = ['get']
...@@ -168,11 +167,16 @@ def _decompress(fname): ...@@ -168,11 +167,16 @@ def _decompress(fname):
os.remove(fname) os.remove(fname)
def _get_pretrained():
with open('./ppcls/utils/pretrained.list') as flist:
pretrained = [line.strip() for line in flist]
return pretrained
def _check_pretrained_name(architecture): def _check_pretrained_name(architecture):
assert isinstance(architecture, str), \ assert isinstance(architecture, str), \
("the type of architecture({}) should be str". format(architecture)) ("the type of architecture({}) should be str". format(architecture))
with open('./configs/pretrained.list') as flist: pretrained = _get_pretrained()
pretrained = [line.strip() for line in flist]
similar_names = similar_architectures(architecture, pretrained) similar_names = similar_architectures(architecture, pretrained)
model_list = ', '.join(similar_names) model_list = ', '.join(similar_names)
err = "{} is not exist! Maybe you want: [{}]" \ err = "{} is not exist! Maybe you want: [{}]" \
...@@ -181,6 +185,14 @@ def _check_pretrained_name(architecture): ...@@ -181,6 +185,14 @@ def _check_pretrained_name(architecture):
raise ModelNameError(err) raise ModelNameError(err)
def list_models():
pretrained = _get_pretrained()
msg = "All avialable pretrained models are as follows: {}".format(
pretrained)
logger.info(msg)
return
def get(architecture, path, decompress=True): def get(architecture, path, decompress=True):
""" """
Get the pretrained model. Get the pretrained model.
...@@ -188,5 +200,6 @@ def get(architecture, path, decompress=True): ...@@ -188,5 +200,6 @@ def get(architecture, path, decompress=True):
_check_pretrained_name(architecture) _check_pretrained_name(architecture)
url = _get_url(architecture) url = _get_url(architecture)
fname = _download(url, path) fname = _download(url, path)
if decompress: _decompress(fname) if decompress:
_decompress(fname)
logger.info("download {} finished ".format(fname)) logger.info("download {} finished ".format(fname))
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
# #
#Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
#you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
#You may obtain a copy of the License at # You may obtain a copy of the License at
# #
# http://www.apache.org/licenses/LICENSE-2.0 # http://www.apache.org/licenses/LICENSE-2.0
# #
#Unless required by applicable law or agreed to in writing, software # Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS, # distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
#limitations under the License. # limitations under the License.
import sys
import argparse import argparse
from ppcls import model_zoo from ppcls import model_zoo
...@@ -26,6 +25,7 @@ def parse_args(): ...@@ -26,6 +25,7 @@ def parse_args():
parser.add_argument('-a', '--architecture', type=str, default='ResNet50') parser.add_argument('-a', '--architecture', type=str, default='ResNet50')
parser.add_argument('-p', '--path', type=str, default='./pretrained/') parser.add_argument('-p', '--path', type=str, default='./pretrained/')
parser.add_argument('-d', '--decompress', type=str2bool, default=True) parser.add_argument('-d', '--decompress', type=str2bool, default=True)
parser.add_argument('-l', '--list', type=str2bool, default=False)
args = parser.parse_args() args = parser.parse_args()
return args return args
...@@ -33,7 +33,10 @@ def parse_args(): ...@@ -33,7 +33,10 @@ def parse_args():
def main(): def main():
args = parse_args() args = parse_args()
model_zoo.get(args.architecture, args.path, args.decompress) if args.list:
model_zoo.list_models()
else:
model_zoo.get(args.architecture, args.path, args.decompress)
if __name__ == '__main__': if __name__ == '__main__':
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册