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

* 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

View File

@ -2,10 +2,14 @@ import argparse
from mmcv import Config, DictAction
from mmseg.apis import init_segmentor
def parse_args():
parser = argparse.ArgumentParser(description='Print the whole config')
parser.add_argument('config', help='config file path')
parser.add_argument(
'--graph', action='store_true', help='print the models graph')
parser.add_argument(
'--options', nargs='+', action=DictAction, help='arguments in dict')
args = parser.parse_args()
@ -22,6 +26,12 @@ def main():
print(f'Config:\n{cfg.pretty_text}')
# dump config
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__':