main.py

Go to the documentation of this file.

import sys
sys_path = ['/workspace/smlab', '/workspace']
for path in sys_path:
    if path not in sys.path:
        sys.path.append(path)
        
import argparse
 
from mmengine.config import Config
 
from smrunner.train import train
from smrunner.test import test
 
 
parser = argparse.ArgumentParser()
parser.add_argument('--mode', default='train', help='train or test')
parser.add_argument('--cfg', default='', help='cfg file name')
parser.add_argument('--device', type=int, default=None, nargs='+', help='select device')
 
if __name__ == '__main__':    
    
    args = parser.parse_args()
    mode = args.mode
    cfg_path = args.cfg
    cfg = Config.fromfile(cfg_path)
    device_ids = args.device
    
    if mode == 'train':
        train(cfg, device_ids)
    elif mode == 'test':
        test(cfg, device_ids)
    else:
        pass