医学图像的分割任务的模型训练train代码

梦里梦见梦不见的497

已于 2025-02-02 23:43:08 修改

阅读量127

点赞数 1

文章标签： pytorch 人工智能 python

于 2025-02-02 22:29:32 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_73101319/article/details/145423379

版权

使用了五折交叉验证来训练模型，验证集使用AUC来筛选出最优模型。

import argparse
import os
import nibabel as nib
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import KFold
from sklearn.metrics import roc_auc_score

from utils import NiiDataset
from model.UNet import UNet
# 自定义数据集类

# 定义图像和掩码路径
image_paths = [
    r"D:\Data\DegmentData\OriginalNii\DCE\CAO_ZHAN_GUO.nii",
    r"D:\Data\DegmentData\OriginalNii\DCE\CHAI_GUI_LAN.nii",
    # 添加其他路径...
]

mask_paths = [
    r"D:\Data\DegmentData\ROI\CAO_ZHAN_GUO-label.nii",
    r"D:\Data\DegmentData\ROI\CHAI_GUI_LAN-label.nii",
    # 添加其他路径...
]

parser = argparse.ArgumentParser()

parser.add_argument('-b', type=int, default=4, help='batch size for dataloader')

parser.add_argument('-lr', type=flo