跳到内容

Swin Unet 代码复现

2025年2月27日

参考博文:

环境搭建

conda create -n swinunet python=3.8
conda activate swinunet
pip install -r requirements.txt

如果运行时提示缺少某些包,可以根据提示安装缺少的依赖。

下载预训练模型和数据集

  1. 下载预训练模型 访问 https://drive.google.com/drive/folders/1UC3XOoezeum0uck4KBVGa8osahs6rKUY 下载 swin_tiny_patch4_window7_224.pth 预训练模型,并将其放入 pretrained_ckpt 文件夹。
  2. 下载官方数据集 访问 https://drive.google.com/drive/folders/1ACJEoTp-uqfFJ73qS3eUObQh52nGuzCd 下载官方数据集 project_TransUNet.zip,并将解压后的 data 文件夹放到根目录。
  3. 生成val.txt文件 将 lists 文件夹中的 list_Synapse 重命名为 Synapse,然后在 Synapse 文件夹内创建 val.txt 文件,内容如下:
case0001
case0002
case0003
case0004
case0008
case0022
case0025
case0029
case0032
case0035
case0036
case0038

修改部分代码

import argparse
import os
import random
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from networks.vision_transformer import SwinUnet as ViT_seg
from trainer import trainer_synapse
from config import get_config

parser = argparse.ArgumentParser()
parser.add_argument('--root_path', type=str, default='./data/Synapse', help='root dir for data')
parser.add_argument('--dataset', type=str, default='Synapse', help='experiment_name')
parser.add_argument('--list_dir', type=str, default='./lists/lists_Synapse', help='list dir')
parser.add_argument('--num_classes', type=int, default=9, help='output channel of network')
parser.add_argument('--output_dir', type=str, default='./output', help='output dir')
parser.add_argument('--max_iterations', type=int, default=30000, help='maximum epoch number to train')
parser.add_argument('--max_epochs', type=int, default=150, help='maximum epoch number to train')
parser.add_argument('--batch_size', type=int, default=24, help='batch_size per gpu')
parser.add_argument('--n_gpu', type=int, default=1, help='total gpu')
parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training')
parser.add_argument('--base_lr', type=float, default=0.01, help='segmentation network learning rate')
parser.add_argument('--img_size', type=int, default=224, help='input patch size of network input')
parser.add_argument('--seed', type=int, default=1234, help='random seed')
parser.add_argument('--cfg', type=str, default='./configs/swin_tiny_patch4_window7_224_lite.yaml', metavar="FILE", help='path to config file', )
parser.add_argument("--opts", default=None, nargs='+', help="Modify config options by adding 'KEY VALUE' pairs. ",)
parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
                    help='no: no cache, '
                         'full: cache all data, '
                         'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
parser.add_argument('--resume', help='resume from checkpoint')
parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps")
parser.add_argument('--use-checkpoint', action='store_true', help="whether to use gradient checkpointing to save memory")
parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'], help='mixed precision opt level, if O0, no amp is used')
parser.add_argument('--tag', help='tag of experiment')
parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
parser.add_argument('--throughput', action='store_true', help='Test throughput only')
# parser.add_argument("--dataset_name", default="datasets")
parser.add_argument("--n_class", default=4, type=int)
parser.add_argument("--num_workers", default=0, type=int)
parser.add_argument("--eval_interval", default=1, type=int)

args = parser.parse_args()
if args.dataset == "Synapse":
    args.root_path = os.path.join(args.root_path, "train_npz")
config = get_config(args)

if __name__ == "__main__":
    if not args.deterministic:
        cudnn.benchmark = True
        cudnn.deterministic = False
    else:
        cudnn.benchmark = False
        cudnn.deterministic = True

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    dataset_name = args.dataset
    dataset_config = {
        args.dataset: {
            'root_path': args.root_path,
            'list_dir': f'./lists/{args.dataset}',
            'num_classes': args.n_class,
        },
    }

    if args.batch_size != 24 and args.batch_size % 6 == 0:
        args.base_lr *= args.batch_size / 24
    args.num_classes = dataset_config[dataset_name]['num_classes']
    args.root_path = dataset_config[dataset_name]['root_path']
    args.list_dir = dataset_config[dataset_name]['list_dir']

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    net = ViT_seg(config, img_size=args.img_size, num_classes=args.num_classes).cuda()
    net.load_from(config)

    # trainer = {'Synapse': trainer_synapse}
    trainer_synapse(args, net, args.output_dir)

# python train.py --output_dir ./model_out/datasets --dataset datasets --img_size 224 --batch_size 32 --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --root_path /media/aicvi/11111bdb-a0c7-4342-9791-36af7eb70fc0/NNUNET_OUTPUT/nnunet_preprocessed/Dataset001_mm/nnUNetPlans_2d_split

运行代码

在配置好环境后,可以运行以下命令启动训练。根据你的 GPU 配置调整 --batch_size 的大小:

(swinunet) > python train.py --max_epochs 20 --base_lr 0.05 --batch_size 12

启动后,程序会显示类似以下的信息:

=> merge config from ./configs/swin_tiny_patch4_window7_224_lite.yaml
SwinTransformerSys expand initial----depths:[2, 2, 2, 2];depths_decoder:[1, 2, 2, 2];drop_path_rate:0.2;num_classes:4

---
final upsample expand_first---
pretrained_path:./pretrained_ckpt/swin_tiny_patch4_window7_224.pth
---

Namespace(accumulation_steps=None, amp_opt_level='O1', base_lr=0.05, batch_size=8, cache_mode='part', cfg='./configs/swin_tiny_patch4_window7_224_lite.yaml', dataset='Synapse', deterministic=1, eval=False, eval_interval=1, img_size=224, list_dir='./lists/Synapse', max_epochs=20, max_iterations=30000, n_class=4, n_gpu=1, num_classes=4, num_workers=0, opts=None, output_dir='./output', resume=None, root_path='./data/Synapse\\train_npz', seed=1234, tag=None, throughput=False, use_checkpoint=False, zip=False)
The length of train set is: 2211
277 iterations per epoch. 5540 max iterations
0%|                              | 0/20 [00:00<?, ?it/s]
Train: 0:   7%|█████▋                                 | 19/277 [00:08<01:31,  2.80it/s]

训练完成后,模型将保存在 output 目录下,生成 best_model.pth 和 last_model.pth 文件。

上一篇
Nuxt UI Pro 编译记录
下一篇
邮件服务理论基础