BraTS2025(TransUNet)
安装
pip install numpy==1.23 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install monai -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install matplotlib batchgenerators pandas SimpleITK medpy tqdm -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install segmentation_models_pytorch monai einops SimpleITK -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install pyyaml einops adamp gco-wrapper medpy nibabel tensorboardX tqdm ml_collections -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install fvcore -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install nnunet -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install git+https://github.com/facebookresearch/detectron2.git -i https://pypi.tuna.tsinghua.edu.cn/simple #要用加速网站
数据转换
目录格式
nnUNet_raw_data_base/nnUNet_raw_data/Task001_BrainTumour/
├── dataset.json
├── imagesTr
│ ├── BRATS_001_0000.nii.gz
│ ├── BRATS_001_0001.nii.gz
│ ├── BRATS_001_0002.nii.gz
│ ├── BRATS_001_0003.nii.gz
│ ├── BRATS_002_0000.nii.gz
│ ├── BRATS_002_0001.nii.gz
│ ├── BRATS_002_0002.nii.gz
│ ├── BRATS_002_0003.nii.gz
│ ├── BRATS_003_0000.nii.gz
│ ├── BRATS_003_0001.nii.gz
│ ├── BRATS_003_0002.nii.gz
│ ├── BRATS_003_0003.nii.gz
│ ├── BRATS_004_0000.nii.gz
│ ├── BRATS_004_0001.nii.gz
│ ├── BRATS_004_0002.nii.gz
│ ├── BRATS_004_0003.nii.gz
│ ├── ...
├── imagesTs
│ ├── BRATS_485_0000.nii.gz
│ ├── BRATS_485_0001.nii.gz
│ ├── BRATS_485_0002.nii.gz
│ ├── BRATS_485_0003.nii.gz
│ ├── BRATS_486_0000.nii.gz
│ ├── BRATS_486_0001.nii.gz
│ ├── BRATS_486_0002.nii.gz
│ ├── BRATS_486_0003.nii.gz
│ ├── BRATS_487_0000.nii.gz
│ ├── BRATS_487_0001.nii.gz
│ ├── BRATS_487_0002.nii.gz
│ ├── BRATS_487_0003.nii.gz
│ ├── BRATS_488_0000.nii.gz
│ ├── BRATS_488_0001.nii.gz
│ ├── BRATS_488_0002.nii.gz
│ ├── BRATS_488_0003.nii.gz
│ ├── BRATS_489_0000.nii.gz
│ ├── BRATS_489_0001.nii.gz
│ ├── BRATS_489_0002.nii.gz
│ ├── BRATS_489_0003.nii.gz
│ ├── ...
└── labelsTr
├── BRATS_001.nii.gz
├── BRATS_002.nii.gz
├── BRATS_003.nii.gz
├── BRATS_004.nii.gz
├── ...
Json 格式
{
"description": "BraTS2025 dataset converted to BraTS2020 format",
"labels": {
"0": "background",
"1": "edema",
"2": "non-enhancing",
"3": "enhancing",
"4": "other"
},
"licence": "BraTS2025 license",
"modality": {
"0": "T1c",
"1": "T1n",
"2": "T2f",
"3": "T2w"
},
"name": "BraTS2025",
"numTest": 0,
"numTraining": 1296,
"reference": "BraTS2025",
"release": "0.0",
"tensorImageSize": "4D",
"test": [],
"training": [
{
"image": "./imagesTr/BraTS-MET-00001-000.nii.gz",
"label": "./labelsTr/BraTS-MET-00001-000.nii.gz"
},
注意 json 里面的 image 路径不能带模态标识符, 代码会自动添加
预处理
export nnUNet_codebase="/data/coding/3D-TransUNet-main" # 代码路径
export nnUNet_raw_data_base="/data/coding/nnUNet_raw" # 原始数据路径
export nnUNet_preprocessed="/data/coding/nnUNet_preprocessed" # 预处理数据路径
export RESULTS_FOLDER="/data/coding/nnUNet_results" # 结果路径
nnUNet_plan_and_preprocess -t 082 --verify_dataset_integrity
训练
train.sh
#!/bin/bash
# 切换到项目根目录
cd /data/coding/3D-TransUNet-main
export nnUNet_N_proc_DA=36
export nnUNet_codebase="/data/coding/3D-TransUNet-main" # 你的代码路径
export nnUNet_raw_data_base="/data/coding/nnUNet_raw" # 你的原始数据路径
export nnUNet_preprocessed="/data/coding/nnUNet_preprocessed" # 你的预处理数据路径
export RESULTS_FOLDER="/data/coding/nnUNet_results" # 你的结果路径
CONFIG="/data/coding/3D-TransUNet-main/configs/Brats/encoder_plus_decoder.yaml"
echo $CONFIG
### unit test
fold=0
echo "run on fold: ${fold}"
nnunet_use_progress_bar=1 CUDA_VISIBLE_DEVICES=0,1 \
python3 -m torch.distributed.launch --master_port=4322 --nproc_per_node=2 \
/data/coding/3D-TransUNet-main/train.py --fold=${fold} --config=$CONFIG --resume=''
train.py
# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
# 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 argparse
import os
import pickle
import torch
import torch.distributed
import yaml
from nnunet.paths import default_plans_identifier
from nnunet.run.load_pretrained_weights import load_pretrained_weights
from nn_transunet.default_configuration import get_default_configuration
def main():
parser = argparse.ArgumentParser()
# change batch_size in nnUNetTrainer.py self.batch_size = stage_plans['batch_size']; can change batch_size=1 if world_size=16
parser.add_argument("--network", default="3d_fullres", type=str)
parser.add_argument("--network_trainer", default="nnUNetTrainerV2_DDP")
parser.add_argument("--task", default="Task801_WORD", help="can be task name or task id")
parser.add_argument("--task_pretrained", default="Task801_WORD", help="option Task801_WORD, Task850_ABD1K")
parser.add_argument("--fold", help='0, 1, ..., 5 or \'all\'')
parser.add_argument("--model", default="Generic_UNet", type=str)
parser.add_argument("--disable_ds", default=False, type=bool)
parser.add_argument("--resume", default='local_latest', type=str) # auto
parser.add_argument("-val", "--validation_only", help="use this if you want to only run the validation",
action="store_true")
parser.add_argument("-c", "--continue_training", help="use this if you want to continue a training",
action="store_true")
parser.add_argument("-p", help="plans identifier. Only change this if you created a custom experiment planner",
default=default_plans_identifier, required=False)
parser.add_argument("--use_compressed_data", default=False, action="store_true",
help="If you set use_compressed_data, the training cases will not be decompressed. Reading compressed data "
"is much more CPU and RAM intensive and should only be used if you know what you are "
"doing", required=False)
parser.add_argument("--deterministic",
help="Makes training deterministic, but reduces training speed substantially. I (Fabian) think "
"this is not necessary. Deterministic training will make you overfit to some random seed. "
"Don't use that.",
required=False, default=False, action="store_true")
parser.add_argument("--fp32", required=False, default=False, action="store_true",
help="disable mixed precision training and run old school fp32")
parser.add_argument("--dbs", required=False, default=False, action="store_true", help="distribute batch size. If "
"True then whatever "
"batch_size is in plans will "
"be distributed over DDP "
"models, if False then each "
"model will have batch_size "
"for a total of "
"GPUs*batch_size")
parser.add_argument("--npz", required=False, default=False, action="store_true", help="if set then nnUNet will "
"export npz files of "
"predicted segmentations "
"in the vlaidation as well. "
"This is needed to run the "
"ensembling step so unless "
"you are developing nnUNet "
"you should enable this")
parser.add_argument("--valbest", required=False, default=False, action="store_true", help="")
parser.add_argument("--vallatest", required=False, default=False, action="store_true", help="")
parser.add_argument("--find_lr", required=False, default=False, action="store_true", help="")
parser.add_argument("--val_folder", required=False, default="validation_raw",
help="name of the validation folder. No need to use this for most people")
parser.add_argument("--disable_saving", required=False, action='store_true',
help="If set nnU-Net will not save any parameter files. Useful for development when you are "
"only interested in the results and want to save some disk space")
parser.add_argument("--disable_postprocessing_on_folds", required=False, action='store_true',
help="Running postprocessing on each fold only makes sense when developing with nnU-Net and "
"closely observing the model performance on specific configurations. You do not need it "
"when applying nnU-Net because the postprocessing for this will be determined only once "
"all five folds have been trained and nnUNet_find_best_configuration is called. Usually "
"running postprocessing on each fold is computationally cheap, but some users have "
"reported issues with very large images. If your images are large (>600x600x600 voxels) "
"you should consider setting this flag.")
parser.add_argument('-pretrained_weights', type=str, required=False, default=None,
help='path to nnU-Net checkpoint file to be used as pretrained model (use .model '
'file, for example model_final_checkpoint.model). Will only be used when actually training. '
'Optional. Beta. Use with caution.')
parser.add_argument('--config', default='', type=str, metavar='FILE',
help='YAML config file specifying default arguments')
parser.add_argument('--batch_size', default=2, type=int)
parser.add_argument('--max_num_epochs', default=1000, type=int)
parser.add_argument('--initial_lr', default=0.1, type=float)
parser.add_argument('--min_lr', default=0, type=float)
parser.add_argument('--opt_eps', default=1e-8, type=float, metavar='EPSILON',
help='Optimizer Epsilon (default: 1e-8), from MAE ft')
parser.add_argument('--opt_betas', default=None, type=float, nargs='+', metavar='BETA',
help='Optimizer Betas (default: None, use opt default), from MAE ft')
parser.add_argument('--weight_decay', default=3e-5, type=float)
parser.add_argument("--local-rank", type=int) # must pass
parser.add_argument('--world-size', default=1, type=int, help='number of nodes for distributed training')
parser.add_argument('--rank', default=0, type=int, help='node rank for distributed training')
parser.add_argument('--total_batch_size', default=None, type=int, help='node rank for distributed training')
parser.add_argument('--hdfs_base', default='', type=str)
parser.add_argument('--optim_name', default='', type=str) # sgd as default, otherwise will have effect in nnUNetTrainerV2_DDP
parser.add_argument('--lrschedule', default='', type=str) # polylr as default
parser.add_argument('--warmup_epochs', default=None, type=int)
parser.add_argument("--val_final", default=False, action="store_true", help="")
parser.add_argument("--is_ssl", default=False, action="store_true", help="SSL pretraining")
parser.add_argument("--is_spatial_aug_only", default=False, action="store_true", help="SSL pretraining")
parser.add_argument('--mask_ratio', default=0.75, type=float)
parser.add_argument('--loss_name', default='', type=str)
parser.add_argument('--plan_update', default='', type=str)
parser.add_argument('--crop_size', nargs='+', type=int, default=None,
help='input to network')
parser.add_argument('--reclip', nargs='+', type=int)
parser.add_argument("--pretrained", default=False, action="store_true", help="")
parser.add_argument("--disable_decoder", default=False, action="store_true", help="disable decoder of mae network")
parser.add_argument("--model_params", default={})
parser.add_argument('--layer_decay', default=1.0, type=float, help="layer-wise dacay for lr")
parser.add_argument('--drop_path', type=float, default=0.0, metavar='PCT',
help='Drop path rate (default: 0.1), drop_path=0 for MAE pretrain')
parser.add_argument("--find_zero_weight_decay", default=False, action="store_true", help="")
parser.add_argument('--n_class', default=17, type=int, help="17 for WORD including background")
parser.add_argument('--deep_supervision_scales', nargs='+', type=int, default=[], help='remember to align with pat_emb_stride for z')
parser.add_argument("--fix_ds_net_numpool", default=False, action="store_true", help="")
parser.add_argument("--skip_grad_nan", default=False, action="store_true", help="skip_grad_nan in nnUNetTrainerV2_DDP")
parser.add_argument("--merge_femur", default=False, action="store_true", help="merge class-15 and class-16 (head of femur) during training")
parser.add_argument("--is_sigmoid", default=False, action="store_true", help="is_sigmoid for output instead of softmax")
parser.add_argument('--max_loss_cal', default='', type=str, help="v0, v1")
# debug
args_config, _ = parser.parse_known_args() # expect return 'remaining' standing for the namspace from launch? but not...
# if args_config.config:
with open(args_config.config, 'r') as f:
cfg = yaml.safe_load(f)
parser.set_defaults(**cfg)
model_params = cfg.get("model_params", {})
args = parser.parse_args() # update args from yaml
task = args.task
fold = args.fold
network = args.network
network_trainer = args.network_trainer
validation_only = args.validation_only
plans_identifier = args.p
use_compressed_data = args.use_compressed_data
decompress_data = not use_compressed_data
deterministic = args.deterministic
valbest = args.valbest
vallatest = args.vallatest
find_lr = args.find_lr
val_folder = args.val_folder
fp32 = args.fp32
disable_postprocessing_on_folds = args.disable_postprocessing_on_folds
# init DDP, in favor of multi-node training
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
torch.distributed.barrier()
if fold.startswith('all'):
pass
else:
fold = int(fold)
if not args.hdfs_base:
args.hdfs_base = network + '_' + args.model
plans_file, output_folder_name, dataset_directory, batch_dice, stage, trainer_class = get_default_configuration(network, task, network_trainer, plans_identifier, hdfs_base=args.hdfs_base, plan_update=args.plan_update)
resolution_index = 1
if args.config.find('500Region') != -1:
batch_dice = True
resolution_index = 0
if '005' in plans_file or '004' in plans_file or '002' in plans_file or '001' in plans_file or 'BraTS' in plans_file:
resolution_index = 0
info = pickle.load(open(plans_file, "rb"))
plan_data = {}
plan_data["plans"] = info
patch_size = plan_data['plans']['plans_per_stage'][resolution_index]['patch_size']
if args.crop_size is None:
args.crop_size = patch_size
if trainer_class is None:
raise RuntimeError("Could not find trainer class in meddec.model_training")
if args.pretrained:
fold_name = 'all' if isinstance(fold, str) and fold.startswith('all') else 'fold_'+str(fold)
init_ckpt_base = model_params['init_ckpt']
pretrained_output_folder = output_folder_name.replace(args.hdfs_base, init_ckpt_base) + '/' + fold_name
pretrained_ckpt_path = pretrained_output_folder + "/model_latest.model" # check network_trainer.load_latest_checkpoint()
if args.task_pretrained!= args.task:
pretrained_output_folder = pretrained_output_folder.replace(args.task, args.task_pretrained)
pretrained_ckpt_path = pretrained_ckpt_path.replace(args.task, args.task_pretrained)
os.makedirs(pretrained_output_folder, exist_ok=True)
if args.local_rank==0:
downloaded = pretrained_ckpt_path if os.path.exists(pretrained_ckpt_path) else False
if not downloaded:
print("pretrained weights not existed in both local and remote")
else:
print("pretrained weights downloaded to remote")
torch.distributed.barrier() # make sure each rank has updated model_params
model_params['init_ckpt'] = pretrained_ckpt_path
print("###########update model_params['init_ckpt']: ", model_params['init_ckpt'])
trainer = trainer_class(plans_file, fold, local_rank=args.local_rank, output_folder=output_folder_name,
dataset_directory=dataset_directory, batch_dice=batch_dice, stage=stage,
unpack_data=decompress_data, deterministic=deterministic, fp16=not fp32,
distribute_batch_size=args.dbs,
# model=args.model, disable_ds=args.disable_ds, resume=args.resume,
input_size=args.crop_size,
args=args) # for V2
if args.disable_saving:
trainer.save_latest_only = False # if false it will not store/overwrite _latest but separate files each
trainer.save_intermediate_checkpoints = False # whether or not to save checkpoint_latest
trainer.save_best_checkpoint = False # whether or not to save the best checkpoint according to self.best_val_eval_criterion_MA
trainer.save_final_checkpoint = False # whether or not to save the final checkpoint
trainer.initialize(not validation_only)
resume_epoch = 0
if find_lr:
trainer.find_lr()
else:
if not validation_only:
if args.continue_training:
# -c was set, continue a previous training and ignore pretrained weights
trainer.load_latest_checkpoint()
# trainer.load_checkpoint_ram()
elif (not args.continue_training) and (args.pretrained_weights is not None):
# we start a new training. If pretrained_weights are set, use them
load_pretrained_weights(trainer.network, args.pretrained_weights)
else:
# new training without pretraine weights, do nothing
pass
if args.resume == 'auto':
fold_name = fold if isinstance(fold, str) and fold.startswith('all') else 'fold_'+str(fold)
output_folder = output_folder_name + '/' + fold_name
assert trainer.output_folder == output_folder, "output_folder path are not consistent!" # check if consistent!
if args.local_rank == 0: os.makedirs(output_folder, exist_ok=True)
ckpt_path = output_folder + "/model_latest.model" # check network_trainer.load_latest_checkpoint()
if args.local_rank == 0: # downloaded for each node
resume = ckpt_path if os.path.exists(ckpt_path) else False
torch.distributed.barrier()
resume = ckpt_path if os.path.exists(ckpt_path) else False # set resume flag for every process
if resume: # will find ckpt_path (find 1. best 2. final 3. latest) in network_trainer
print("### resume, load_latest_checkpoint")
trainer.load_latest_checkpoint() # load ckpt, opt, amp, epoch, plot.... check network_trainer.load_latest_checkpoint(), which will call nnUNetTrainerV2_DDP.load_checkpoint_ram()
resume_epoch = trainer.epoch
elif args.resume == 'local_latest':
fold_name = fold if isinstance(fold, str) and fold.startswith('all') else 'fold_'+str(fold)
output_folder = output_folder_name + '/' + fold_name
assert trainer.output_folder == output_folder, "output_folder path are not consistent!" # check if consistent!
if args.local_rank == 0: os.makedirs(output_folder, exist_ok=True)
torch.distributed.barrier()
ckpt_path = output_folder + "/model_latest.model" # check network_trainer.load_latest_checkpoint()
resume = ckpt_path if os.path.exists(ckpt_path) else False # set resume flag for every process
if resume: # will find ckpt_path (find 1. best 2. final 3. latest) in network_trainer
print("### resume, load_latest_checkpoint")
trainer.load_latest_checkpoint() # load ckpt, opt, amp, epoch, plot.... check network_trainer.load_latest_checkpoint(), which will call nnUNetTrainerV2_DDP.load_checkpoint_ram()
resume_epoch = trainer.epoch
trainer.run_training()
else:
if valbest:
trainer.load_best_checkpoint(train=False)
elif vallatest:
trainer.load_latest_checkpoint(train=False)
else:
trainer.load_final_checkpoint(train=False)
trainer.network.eval()
# predict validation !!!!!!
if args.val_final or vallatest:
trainer.validate(save_softmax=args.npz, validation_folder_name=val_folder,
run_postprocessing_on_folds=not disable_postprocessing_on_folds)
if network == '3d_lowres':
raise NotImplementedError
# torch.distributed.barrier()
print("######### run_training_DDP done!")
# torch.distributed.destroy_process_group()
if __name__ == "__main__":
main()
验证和推理
Validation和 inference 都需要调用**predict_preprocessed_data_return_seg_and_softmax, 这个函数完全用不了. 最后还是尝试改通了, 但是又出现了两个问题: validation流程不符合 plan. Pkl 和 dataset. Json 约束, 代码是封装到 nnunetv 1 包里面的, v 1 会尝试自己创建一个单独的 5 折分割而且强制分割不能手动分, 这样把测试数据放到验证集测数据的法子就用不了. Inference 流程则有封装在 nnunetv 1 里面不可名状的测试集读取逻辑, 文件里代码看似是直接遍历测试文件夹, 实则在不可见的流程中间在乱拼凑文件名导致访问不存在的路径, 这个问题无论如何都无法解决
目前只有训练过程中对整个验证集测出来的数据, 没有推理文件和每个测试集 case 的数据.