def caffe2_global_init(logging_print_net_summary=0, num_threads=None): if num_threads is None: if get_num_processes_per_machine() > 1: # by default use single thread when DDP with multiple processes num_threads = 1 else: # GlobalInit will clean PyTorch's num_threads and set it to 1, # thus keep PyTorch's default value to make it truly default. num_threads = torch.get_num_threads() if not get_local_rank() == 0: logging_print_net_summary = 0 # only enable for local main process from caffe2.python import workspace workspace.GlobalInit([ "caffe2", "--caffe2_log_level=2", "--caffe2_logging_print_net_summary={}".format( logging_print_net_summary), "--caffe2_omp_num_threads={}".format(num_threads), "--caffe2_mkl_num_threads={}".format(num_threads), ]) logger.info("Using {} threads after GlobalInit".format( torch.get_num_threads()))
def create_dir_on_local_main_process(dir): if get_local_rank() == 0 and dir: PathManager.mkdirs(dir)