保姆级教程：从零在Windows上用PyCharm复现TransUNet（含数据集处理完整代码）

Windows下用PyCharm复现TransUNet的完整实战指南医学图像分割是计算机视觉在医疗领域的重要应用而TransUNet作为结合Transformer与U-Net优势的模型正在成为该领域的新标杆。但对于Windows用户和深度学习新手来说从零开始复现论文模型往往充满挑战——环境配置复杂、路径报错频发、数据集处理繁琐等问题让许多人望而却步。本文将彻底解决这些痛点提供一套真正适合Windows平台的保姆级解决方案。与常见教程不同我们特别针对PyCharm IDE进行了优化所有操作都基于图形界面完成无需记忆复杂命令行。从数据集预处理到模型训练测试每个步骤都配有详细截图和常见错误解决方案即使是刚接触医学图像分析的新手也能顺利完成复现。下面让我们从最基础的环境搭建开始逐步攻克这个项目。1. 环境配置与准备工作在开始之前我们需要确保开发环境正确配置。TransUNet作为基于PyTorch的模型对硬件和软件都有一定要求。以下是经过实测的推荐配置硬件要求显卡NVIDIA GTX 1060 6GB或更高需支持CUDA内存16GB以上处理3D医学图像时内存消耗较大存储至少50GB可用空间原始数据集和预处理文件会占用大量空间软件准备清单PyCharm Professional 2023.3社区版也可用但缺少部分专业功能Python 3.8.10这是与PyTorch各版本兼容性最好的Python版本Git for Windows用于克隆原始仓库7-Zip或WinRAR用于解压数据集首先在PyCharm中创建新项目建议使用虚拟环境而非系统Python环境。创建时勾选New environment using VirtualenvPython版本选择3.8.10。虚拟环境创建完成后我们需要安装核心依赖包。打开PyCharm的TerminalAltF12逐条执行以下命令pip install torch1.12.1cu113 torchvision0.13.1cu113 --extra-index-url https://download.pytorch.org/whl/cu113 pip install nibabel4.0.2 pip install opencv-python4.5.5.64 pip install tqdm4.64.1 pip install scikit-image0.19.3注意torch的CUDA版本需要与本地安装的CUDA工具包版本匹配。可通过nvidia-smi命令查看支持的CUDA版本。如果遇到兼容性问题可以尝试torch 1.8.0cu111这个更稳定的组合。常见问题解决报错Unable to find vcvarsall.bat安装Visual Studio Build Tools勾选C桌面开发组件报错NVIDIA driver is too old更新显卡驱动到最新版本PyCharm无法识别新建的虚拟环境在Settings → Project → Python Interpreter中手动添加解释器路径2. 数据集预处理全流程详解医学图像通常以NIfTI格式(.nii.gz)存储而TransUNet需要PNG图像和NPZ格式的组合输入。我们将分两步完成这个转换过程确保即使没有Linux经验的用户也能轻松操作。2.1 NIfTI到PNG的转换实战在项目根目录下创建如下文件夹结构TransUNet_Project/ ├── predata/ # 存放原始.nii.gz文件 ├── 2Ddata/ # 存放切片后的PNG图像 ├── data/ │ ├── train_npz/ # 最终训练用的npz文件 │ └── lists/ # 存放训练集/测试集划分文件将提供的process_nii_to_png.py脚本放入项目根目录这个改良版脚本特别处理了Windows路径问题# coding:utf-8 import numpy as np import nibabel as nib import os from PIL import Image from tqdm import tqdm data_path ./predata output_dir ./2Ddata def safe_mkdir(path): if not os.path.exists(path): os.makedirs(path) def process_file(file_path): img nib.load(file_path) label_path file_path.replace(_gt.nii.gz, _label.nii.gz) label nib.load(label_path) img_data img.get_fdata() label_data label.get_fdata() # 窗宽窗位调整 img_clipped np.clip(img_data, -125, 275) img_normalised (img_clipped - (-125)) / (275 - (-125)) * 255 for i in range(img_clipped.shape[2]): slice_num i 1 case_name os.path.splitext(os.path.basename(file_path))[0].replace(_gt.nii, ) # Windows路径处理 img_filename f{case_name}_{slice_num:03d}.png label_filename f{case_name}_{slice_num:03d}_label.png img_slice Image.fromarray(img_normalised[:, :, i].astype(np.uint8)) label_slice Image.fromarray(label_data[:, :, i].astype(np.uint8)) img_slice.save(os.path.join(output_dir, img_filename)) label_slice.save(os.path.join(output_dir, label_filename)) if __name__ __main__: safe_mkdir(output_dir) for root, _, files in os.walk(data_path): for file in tqdm(files, descProcessing NIfTI files): if file.endswith(_gt.nii.gz): process_file(os.path.join(root, file))运行此脚本前请确保原始数据命名符合{case_id}_gt.nii.gz和{case_id}_label.nii.gz格式所有.nii.gz文件已放入predata文件夹2Ddata文件夹已创建常见错误处理如果遇到Permission denied错误请以管理员身份运行PyCharm如果遇到内存不足可以分批处理文件或增加虚拟内存。2.2 生成NPZ文件与数据集划分转换完成后我们需要将配对的图像-标签组合保存为NPZ格式这是PyTorch高效加载数据的理想格式。创建generate_npz.py文件import glob import cv2 import numpy as np from tqdm import tqdm import os import random def safe_mkdir(path): if not os.path.exists(path): os.makedirs(path) def split_dataset(npz_files, train_ratio0.8): random.shuffle(npz_files) split_idx int(len(npz_files) * train_ratio) return npz_files[:split_idx], npz_files[split_idx:] def generate_npz(): png_dir ./2Ddata output_dir ./data/train_npz list_dir ./data/lists safe_mkdir(output_dir) safe_mkdir(list_dir) image_files [f for f in glob.glob(f{png_dir}/*.png) if not f.endswith(_label.png)] # 生成NPZ文件 for img_path in tqdm(image_files, descGenerating NPZ files): label_path img_path.replace(.png, _label.png) image cv2.imread(img_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) label cv2.imread(label_path, flags0) case_id os.path.basename(img_path).split(_)[0] slice_num os.path.basename(img_path).split(_)[1].split(.)[0] npz_filename f{case_id}_{slice_num}.npz np.savez(os.path.join(output_dir, npz_filename), imageimage, labellabel) # 数据集划分 npz_files [os.path.basename(f) for f in glob.glob(f{output_dir}/*.npz)] train_files, test_files split_dataset(npz_files) with open(f{list_dir}/train.txt, w) as f: f.write(\n.join(train_files)) with open(f{list_dir}/test.txt, w) as f: f.write(\n.join(test_files)) if __name__ __main__: generate_npz()这个脚本完成了三个关键任务将PNG图像对转换为NPZ格式自动划分训练集和测试集默认8:2比例生成训练和测试用的文件列表关键参数调整建议对于小数据集1000张图像建议增加train_ratio到0.9如果遇到内存问题可以分批处理文件对于3D数据集建议保持原始病例级别的划分而非切片级别3. TransUNet模型训练技巧3.1 模配置与参数优化下载官方代码后我们需要针对Windows和本地环境进行几处关键修改。首先在train.py中做如下调整# 修改数据加载方式 dataset Dataset( base_dirargs.data_dir, splittrain, list_diros.path.join(args.data_dir, lists), # 确保路径正确 transformtransforms.Compose([ transforms.ToTensor(), ]) ) # 添加Windows特定的多进程处理设置 if os.name nt: # Windows系统 torch.multiprocessing.set_start_method(spawn, forceTrue) num_workers 0 # Windows下建议设为0或1 else: num_workers 4 train_loader DataLoader( dataset, batch_sizeargs.batch_size, shuffleTrue, num_workersnum_workers, # 根据系统自动调整 pin_memoryTrue )推荐训练参数配置RTX 3060 12GB显存参数推荐值说明batch_size8可根据显存调整但不少于4max_epochs200医学图像通常需要更多epochlr3e-4使用AdamW优化器时可适当增大img_size224原始论文尺寸不建议修改save_freq10每10个epoch保存一次模型创建train_transunet.bat批处理文件简化训练启动echo off set PYTHONPATH. python train.py --data_dir ./data --dataset Synapse --batch_size 8 --max_epochs 200 --lr 3e-4 pause3.2 训练监控与问题排查使用TensorBoard监控训练过程tensorboard --logdir ./logs --port 6006常见训练问题及解决方案Loss不下降检查数据归一化是否正确尝试减小学习率如1e-5确认标签是否为单通道且像素值正确GPU内存不足减小batch_size使用--gradient_accumulation_steps 2参数尝试混合精度训练添加--amp参数验证指标波动大增加验证集大小检查数据增强是否过于激进尝试更小的学习率配合warmup在PyCharm中配置TensorBoard非常简单点击Run → Edit Configurations → → Python设置Script path: 选择Python解释器路径下的tensorboard/main.pyParameters:--logdir ./logs --port 60064. 模型测试与结果可视化4.1 测试脚本配置创建test.py并添加以下关键修改# 在测试脚本开头添加Windows特定设置 if os.name nt: import warnings warnings.filterwarnings(ignore, categoryUserWarning, messageLazy modules are a new feature.*) # 修改结果保存路径 result_save_dir os.path.join(args.save_dir, test_results) os.makedirs(result_save_dir, exist_okTrue) # 添加可视化函数 def save_visualization(image, label, pred, save_path): plt.figure(figsize(18, 6)) plt.subplot(1, 3, 1) plt.imshow(image, cmapgray) plt.title(Input Image) plt.subplot(1, 3, 2) plt.imshow(label, cmapjet) plt.title(Ground Truth) plt.subplot(1, 3, 3) plt.imshow(pred, cmapjet) plt.title(Prediction) plt.savefig(save_path) plt.close()4.2 性能评估与指标解读TransUNet常用的评估指标包括Dice系数衡量分割重叠度范围0-1越接近1越好def dice_coef(y_true, y_pred): intersection np.sum(y_true * y_pred) return (2. * intersection) / (np.sum(y_true) np.sum(y_pred))Hausdorff距离衡量边界匹配程度单位像素越小越好from scipy.spatial.distance import directed_hausdorff def hausdorff_distance(y_true, y_pred): return max(directed_hausdorff(y_true, y_pred)[0], directed_hausdorff(y_pred, y_true)[0])灵敏度(Sensitivity)衡量正样本识别能力def sensitivity(y_true, y_pred): tp np.sum(y_true * y_pred) fn np.sum(y_true * (1 - y_pred)) return tp / (tp fn 1e-7)在Synapse多器官分割数据集上的预期表现器官Dice(%)HD95(mm)脾脏92.58.7右肾88.312.4左肾90.110.8肝脏94.26.5提示实际结果可能因数据预处理差异而略有不同。如果Dice系数低于预期5个百分点以上建议检查标签是否正确对齐或数据增强是否合理。