|
|
|
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
|
|
#
|
|
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
# you may not use this file except in compliance with the License.
|
|
|
|
# You may obtain a copy of the License at
|
|
|
|
#
|
|
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
#
|
|
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
# See the License for the specific language governing permissions and
|
|
|
|
# limitations under the License.
|
|
|
|
|
|
|
|
import os
|
|
|
|
import platform
|
|
|
|
import random
|
|
|
|
import multiprocessing as mp
|
|
|
|
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
|
|
|
|
|
|
|
|
|
|
def get_environ_info():
|
|
|
|
"""
|
|
|
|
Collect environment information.
|
|
|
|
"""
|
|
|
|
|
|
|
|
env_info = dict()
|
|
|
|
# TODO is_compiled_with_cuda() has not been moved
|
|
|
|
compiled_with_cuda = paddle.is_compiled_with_cuda()
|
|
|
|
if compiled_with_cuda:
|
|
|
|
if 'gpu' in paddle.get_device():
|
|
|
|
gpu_nums = paddle.distributed.get_world_size()
|
|
|
|
else:
|
|
|
|
gpu_nums = 0
|
|
|
|
if gpu_nums == 0:
|
|
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
|
|
place = 'gpu' if compiled_with_cuda and gpu_nums else 'cpu'
|
|
|
|
env_info['place'] = place
|
|
|
|
env_info['num'] = int(os.environ.get('CPU_NUM', 1))
|
|
|
|
if place == 'gpu':
|
|
|
|
env_info['num'] = gpu_nums
|
|
|
|
|
|
|
|
return env_info
|
|
|
|
|
|
|
|
|
|
|
|
def get_num_workers(num_workers):
|
|
|
|
if not platform.system() == 'Linux':
|
|
|
|
# Dataloader with multi-process model is not supported
|
|
|
|
# on MacOS and Windows currently.
|
|
|
|
return 0
|
|
|
|
if num_workers == 'auto':
|
|
|
|
num_workers = mp.cpu_count() // 2 if mp.cpu_count() // 2 < 2 else 2
|
|
|
|
return num_workers
|
|
|
|
|
|
|
|
|
|
|
|
def init_parallel_env():
|
|
|
|
env = os.environ
|
|
|
|
if 'FLAGS_allocator_strategy' not in os.environ:
|
|
|
|
os.environ['FLAGS_allocator_strategy'] = 'auto_growth'
|
|
|
|
dist = 'PADDLE_TRAINER_ID' in env and 'PADDLE_TRAINERS_NUM' in env
|
|
|
|
if dist:
|
|
|
|
trainer_id = int(env['PADDLE_TRAINER_ID'])
|
|
|
|
local_seed = (99 + trainer_id)
|
|
|
|
random.seed(local_seed)
|
|
|
|
np.random.seed(local_seed)
|
|
|
|
|
|
|
|
if paddle.distributed.get_world_size() > 1:
|
|
|
|
paddle.distributed.init_parallel_env()
|