mirror of
https://github.com/open-mmlab/mmsegmentation.git
synced 2025-06-03 22:03:48 +08:00
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:
parent
4345ee3d55
commit
9524f37fe6
@ -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__':
|
||||
|
Loading…
x
Reference in New Issue
Block a user