mirrors/mmsegmentation
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmsegmentation.git synced 2025-06-03 22:03:48 +08:00
Code Issues Projects Releases Wiki Activity

add print model graph args for tools/print_config.py (#451)

Browse Source 
* add print model graph for print_config tool

* add print model graph for print_config tool

* fix double quoted

* fix iosort
This commit is contained in:
sshuair 2021-04-07 06:14:54 +08:00 committed by GitHub
parent 4345ee3d55
commit 9524f37fe6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tools/print_config.py
View File

@ -2,10 +2,14 @@ import argparse
from mmcv import Config, DictAction from mmcv import Config, DictAction
from mmseg.apis import init_segmentor
def parse_args(): def parse_args():
parser = argparse.ArgumentParser(description='Print the whole config') parser = argparse.ArgumentParser(description='Print the whole config')
parser.add_argument('config', help='config file path') parser.add_argument('config', help='config file path')
parser.add_argument(
'--graph', action='store_true', help='print the models graph')
parser.add_argument( parser.add_argument(
'--options', nargs='+', action=DictAction, help='arguments in dict') '--options', nargs='+', action=DictAction, help='arguments in dict')
args = parser.parse_args() args = parser.parse_args()
@ -22,6 +26,12 @@ def main():
print(f'Config:\n{cfg.pretty_text}') print(f'Config:\n{cfg.pretty_text}')
# dump config # dump config
cfg.dump('example.py') cfg.dump('example.py')
# dump models graph
if args.graph:
model = init_segmentor(args.config, device='cpu')
print(f'Model graph:\n{str(model)}')
with open('example-graph.txt', 'w') as f:
f.writelines(str(model))
if __name__ == '__main__': if __name__ == '__main__':