import os import sys import json import numpy as np from datetime import datetime import torch import torch.distributed import torch.multiprocessing as mp import torch.distributed as dist import logging import argparse import glob os.environ['CUDA_LAUNCH_BLOCKING'] = '1' assert torch.cuda.is_available() sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..")) # __package__ = "lungClassification3D" from BaseClassifier.Se_Cls import Classification3D from BaseClassifier.utils.gpu_utils import set_gpu from k8s_utils import CParamFiller, CK8sPathWrapper assert torch.cuda.is_available() def init_dist(backend='nccl', master_ip='127.0.0.1', port=29500): # print(backend, master_ip, port) if mp.get_start_method(allow_none=True) is None: mp.set_start_method('spawn') os.environ['MASTER_ADDR'] = master_ip os.environ['MASTER_PORT'] = str(port) rank = int(os.environ['RANK']) world_size = int(os.environ['WORLD_SIZE']) # print(rank, world_size) # gpu_ids_list = [str(i) for i in range(world_size)] # gpu_ids = ','.join(gpu_ids_list) # print(gpu_ids) # os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids num_gpus = torch.cuda.device_count() print('num_gpus in init_dist: ', num_gpus) local_rank = os.environ['LOCAL_RANK'] deviceid = eval(local_rank) % num_gpus # num_gpus # deviceid = eval(local_rank) print(fr'dist settings: local_rank {local_rank}, rank {rank}, worldsize {world_size}, gpus {num_gpus}, deviceid {deviceid}') torch.cuda.set_device(deviceid) print('set ok') dist.init_process_group(backend=backend) print('init ok') return rank, world_size def get_logger(name, task_name=None): file_handler = logging.StreamHandler(sys.stdout) logger = logging.getLogger(name) logger.setLevel(logging.INFO) logger.addHandler(file_handler) return logger def _log_msg(strmsg="\n"): if torch.distributed.get_rank() == 0: if g_logger is not None: g_logger.info(strmsg) else: print(strmsg) class MyEncoder(json.JSONEncoder): def default(self, obj): if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, datetime): return obj.__str__() else: return super(MyEncoder, self).default(obj) # set_gpu(num_gpu=1, used_percent=0.1) if __name__ == '__main__': assert torch.cuda.is_available() parser = argparse.ArgumentParser(description="full functional execute script of Detection module.") group = parser.add_mutually_exclusive_group() # ddp parser.add_argument("--local_rank", default=-1, type=int) parser.add_argument('--port', default=29500, type=int, help='port of server') parser.add_argument('--world_size', default=2, type=int) parser.add_argument('--rank', default=0, type=int) parser.add_argument('--master_ip', default='127.0.0.1', type=str) parser.add_argument('--job_data_root', default="/data/job_715/job_data_train", type=str) args = parser.parse_args() g_pathWrapper = CK8sPathWrapper(args.job_data_root) config_path = g_pathWrapper.get_input_inputconfig_filepath() if not os.path.isfile(config_path): print(fr'given config file {config_path} not exist !') else: # print(fr'using config file {config_path} to fill args !') cfg = CParamFiller.fillin_args(config_path, args) train_params = 'configs/train/input_config.json' cfg = CParamFiller.fillin_args(train_params, cfg) print('annotation_label: {}'.format(cfg.annotation_label)) print('annotation_label_list: {}'.format(cfg.annotation_label_list)) if cfg.mode == 'training': print('num_classes in training: {}'.format(cfg.training['num_class'])) else: print('num_classes in testing: {}'.format(cfg.testing['num_class'])) if cfg.pretrain_msg: cfg.pretrain_msg = os.path.join(g_pathWrapper.get_input_dirpath(), cfg.pretrain_msg) print('pretrain_msg: ', cfg.pretrain_msg) g_logger = get_logger(__name__) rank, world_size = init_dist(backend='nccl', master_ip=args.master_ip, port=args.port) cfg.rank = rank cfg.world_size = world_size cfg.local_rank = os.environ['LOCAL_RANK'] # cfg.training['train_data_path'] = os.path.join(g_pathWrapper.get_output_tmp_dirpath(), 'train_data.npy') # cfg.training['train_info_path'] = os.path.join(g_pathWrapper.get_output_tmp_dirpath(), 'train_info.npy') # cfg.training['val_data_path'] = os.path.join(g_pathWrapper.get_output_tmp_dirpath(), 'val_data.npy') # cfg.training['val_info_path'] = os.path.join(g_pathWrapper.get_output_tmp_dirpath(), 'val_info.npy') cfg.training['train_data_path'] = os.path.join(g_pathWrapper.get_output_tmp_dirpath(), 'train_split.json') cfg.training['val_data_path'] = os.path.join(g_pathWrapper.get_output_tmp_dirpath(), 'val_split.json') print('before split data') # 之后ddp是，需要在判断条件中加入： cfg.rank == 0 if (not os.path.exists(cfg.training['train_data_path']) or not os.path.exists(cfg.training['train_info_path']) or not os.path.exists(cfg.training['val_data_path']) or not os.path.exists(cfg.training['val_info_path'])) and cfg.rank == 0: # 读取所有的npy信息，并划分训练集和测试集 print('split data') npy_path = g_pathWrapper.get_output_preprocess_npy_path() train_rate = cfg.split['train'] data_list = glob.glob(os.path.join(npy_path, '*', '*_data.npy')) all_len = len(data_list) train_len = max(1, int(all_len * train_rate)) train_list = data_list[:train_len] val_list = data_list[train_len:] train_data = {'data': train_list} val_data = {'data': val_list} print(f'train len: {len(train_list)}, val len: {len(val_list)}') with open(cfg.training['train_data_path'], 'w+') as f: json.dump(train_data, f, indent=4) with open(cfg.training['val_data_path'], 'w+') as f: json.dump(val_data, f, indent=4) # data2info = {} # uids = [] # for file in os.listdir(npy_path): # uid = file[:-9] # if uid not in uids: # uids.append(uid) # for uid in uids: # data2info[uid + '_data.npy'] = uid + '_info.npy' # whole_data = [] # whole_info = [] # whole_label = [] # for key, value in data2info.items(): # key = os.path.join(npy_path, key) # value = os.path.join(npy_path, value) # whole_data.append(np.load(key)) # working_info = np.load(value, allow_pickle=True) # whole_info.append(working_info) # whole_label.append(working_info[:,-2]) # whole_data = np.concatenate(whole_data, axis=0) # whole_info = np.concatenate(whole_info, axis=0) # whole_label = np.concatenate(whole_label, axis=0) # # 根据whole_label划分训练集和测试集 # label2idx = {} # for label in range(0, cfg.training['num_class']): # label2idx[label] = [] # for idx in range(len(whole_label)): # label2idx[whole_label[idx]].append(idx) # # get train and test npy and info # train_data = [] # val_data = [] # train_info = [] # val_info = [] # for key in list(label2idx.keys()): # choice_idxs = list(np.random.choice(label2idx[key], int(len(label2idx[key]) * cfg.split['train']), replace=False)) # relax_idxs = [idx for idx in label2idx[key] if idx not in choice_idxs] # for idx in choice_idxs: # train_data.append(whole_data[idx][np.newaxis,:]) # train_info.append(whole_info[idx].reshape(-1, len(whole_info[idx]))) # for idx in relax_idxs: # val_data.append(whole_data[idx][np.newaxis,:]) # val_info.append(whole_info[idx].reshape(-1, len(whole_info[idx]))) # train_data = np.concatenate(train_data, axis=0) # train_info = np.concatenate(train_info, axis=0) # val_data = np.concatenate(val_data, axis=0) # val_info = np.concatenate(val_info, axis=0) # np.save(cfg.training['train_data_path'], train_data) # np.save(cfg.training['train_info_path'], train_info) # np.save(cfg.training['val_data_path'], val_data) # np.save(cfg.training['val_info_path'], val_info) print('split data ok') # torch.distributed.barrier() # 设置输出的文件路径 if torch.distributed.get_rank() == 0: print('before train') cfg.training['base_path'] = g_pathWrapper.get_output_train_dirpath() cfg.training['save_path'] = g_pathWrapper.get_output_train_bestmodel_dirpath() cfg.training['writer_path'] = g_pathWrapper.get_output_train_writer_path() cfg.training['eval_pmd'] = g_pathWrapper.get_output_eval_performance_md_filepath() cfg.training['eval_pjson'] = g_pathWrapper.get_output_eval_performance_filepath() cfg.training['train_metrics'] = g_pathWrapper.get_output_train_trainingmetrics_filepath() cfg.training['train_result'] = g_pathWrapper.get_output_train_trainresult_filepath() print('start train') # 开始训练, mode choice in [training, testing] train_obj = Classification3D(cfg, cfg.mode, _log_msg) train_obj()