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 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 # os.environ['CUDA_VISIBLE_DEVICES'] = "0,1,2,3,4,5,6,7" def init_dist(backend='nccl', master_ip='127.0.0.1', port=29500): 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']) num_gpus = torch.cuda.device_count() local_rank = os.environ['LOCAL_RANK'] deviceid = eval(local_rank) % num_gpus torch.cuda.set_device(deviceid) print(fr'dist settings: local_rank {local_rank}, rank {rank}, worldsize {world_size}, gpus {num_gpus}, deviceid {deviceid}') dist.init_process_group(backend=backend) 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__': 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="/fileser51/baoyc/Work/lung_tasks/GroupLung/Datasets/cls_data/job_data_cls_train_2", 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) g_logger = get_logger(__name__) if cfg.mode == 'training': 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'] if cfg.mode == 'training': cfg.training['train_data_path'] = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'train_data.npy') cfg.training['train_info_path'] = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'train_info.npy') cfg.training['val_data_path'] = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'val_data.npy') cfg.training['val_info_path'] = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'val_info.npy') # 之后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信息，并划分训练集和测试集 npy_path = g_pathWrapper.get_output_preprocess_npy_path() 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(1, cfg.training['num_class'] + 1): 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) torch.distributed.barrier() # 设置输出的文件路径 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() # 开始训练, mode choice in [training, testing] train_obj = Classification3D(cfg, cfg.mode, _log_msg) train_obj() elif cfg.mode == 'testing': cfg.testing['test_data_path'] = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'val_data.npy') cfg.testing['test_info_path'] = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'val_info.npy') npy_path = g_pathWrapper.get_output_preprocess_npy_path() 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(1, cfg.training['num_class'] + 1): # 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) all_num = whole_data.shape[0] inter = all_num // 15 for i in range(15): if i == 14: data = whole_data[i*inter: , :] info = whole_info[i*inter: , :] else: data = whole_data[i*inter: (i+1)*inter, :] info = whole_info[i*inter: (i+1)*inter, :] print(data.shape, info.shape) save_data_path = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'val_data_' + str(i) + '.npy') save_info_path = os.path.join(g_pathWrapper.get_output_preprocess_preprocessfiledetails_dirpath(), 'val_info_' + str(i) + '.npy') np.save(save_data_path, data) np.save(save_info_path, info) cfg.testing['test_data_path'] = save_data_path cfg.testing['test_info_path'] = save_info_path # 设置输出的文件路径 cfg.testing['eval_pmd'] = g_pathWrapper.get_output_eval_performance_md_filepath() cfg.testing['eval_pjson'] = g_pathWrapper.get_output_eval_performance_filepath() cfg.testing['eval_result'] = g_pathWrapper.get_output_eval_evalresult_filepath() cfg.testing['save_dir'] = g_pathWrapper.get_output_eval_dirpath() # 开始训练, mode choice in [training, testing] train_obj = Classification3D(cfg, cfg.mode, _log_msg) train_obj()