train_mn.py 文件源码

def create_result_dir(prefix='result'):
    comm = MPI.COMM_WORLD
    if comm.rank == 0:
        result_dir = 'results/{}_{}_0'.format(
            prefix, time.strftime('%Y-%m-%d_%H-%M-%S'))
        while os.path.exists(result_dir):
            i = result_dir.split('_')[-1]
            result_dir = re.sub('_[0-9]+$', result_dir, '_{}'.format(i))
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
    else:
        result_dir = None
    result_dir = comm.bcast(result_dir, root=0)
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
    return result_dir