丹摩征文活动 | 深度学习实战:UNet模型的训练与测试详解


🍑个人主页:Jupiter. 🚀 所属专栏:Linux从入门到进阶 欢迎大家点赞收藏评论😊

目录


1、云实例:配置选型与启动

1.1 登录注册

首先进入登录界面注册并登录账号

1.2 配置 SSH 密钥对

配置 SSH 密钥对的作用是后续远程登录服务器不需要密码验证,更加方便。

首先创建本地公钥,进入本地.ssh目录输入ssh-keygen -o命令,这里文件名可以设置为id_dsa,也可以是其他任意名字

之后我们可以在.ssh目录看到刚刚创建的两个文件

id_dsa id_dsa.pub

其中id_dsa.pub就是需要的公钥文件

进入密钥对配置,创建密钥对,将id_dsa.pub的内容复制到这里就可以

1.3 创建实例

进入GPU 云实例,点击创建实例。如下图所示,按需选择需要的 GPU 型号和镜像

1.4 登录云实例

等待实例创建完成后,点击复制"访问链接"。

接着来到任意一个 SSH 连接终端进行云实例登录,我这里选择的是 VSCode,如下所示

成功后,输入:

bash 复制代码
nvidia-smi
torch.cuda.is_available()

简单验证一下功能即可,如下所示即为成功

2、云存储:数据集上传与下载

文件存储为网络共享存储,可挂载至的不同实例中。相比本地数据盘,其优势是实例间共享,可以多点读写,不受实例释放的影响;此外存储后端有多冗余副本,数据可靠性非常高;但缺陷是 IO 性能一般

考虑到以上优劣,推荐使用方式:将重要数据或代码存放于文件存储中,所有实例共享,便利的同时数据可靠性也有保障;在训练时,需要高 IO 性能的数据(如训练数据),先拷贝到实例本地数据盘,从本地盘读数据获得更好的 IO 性能。如此兼顾便利、安全和性能。

接下来,我们将训练数据上传到云实例数据盘中。使用scp工具如下

bash 复制代码
scp -rP 35740 ./DRIVE-SEG-DATA root@cn-north-b.ssh.damodel.com:/root/workspace

具体地:

35740与cn-north-b.ssh.damodel.com分别为端口号和远程地址,请参考 1.4 节替换为自己的参数
./DRIVE-SEG-DATA是本地数据集路径
/root/workspace是远程实例数据集路径

数据的下载也是类似的命令

bash 复制代码
scp -rP 35740 root@cn-north-b.ssh.damodel.com:/root/workspace ./DRIVE-SEG-DATA

本文提到的数据集可以在DRIVE 数据集中下载:链接:https://drive.grand-challenge.org/Download/

3、云开发:眼底血管分割案例

3.1 案例背景

眼底,作为眼球的内膜,其结构复杂且关键,包括黄斑、视网膜以及视网膜中央动静脉等重要组成部分。在眼科医学领域,眼底图像是医生诊断眼疾病不可或缺的重要依据。近年来,随着深度学习技术的迅猛发展,医学影像分割领域迎来了革命性的变化,眼底图像分割技术也随之取得了显著进步。

深度学习模型,如AlexNet、VGGNet、GoogLeNet和ResNet等,通过训练大规模数据集,能够学习到更加抽象和高级的特征表示。这些特征表示对于实现精确的眼底图像分割至关重要。相较于传统方法,深度学习模型在分割精度和泛化能力上均表现出色。它们能够对未见过的眼底图像数据做出相对准确的预测,从而提高了分割结果的可靠性和稳定性。

此外,深度学习技术还支持端到端的学习方式。这意味着从原始眼底图像输入到最终的分割结果输出,整个过程无需手工设计复杂的特征提取和预处理流程。这种端到端的学习方式不仅简化了分割算法的开发流程,还提高了分割效率和准确性。

值得注意的是,医学影像数据往往包含多种模态,如CT、MRI以及眼底图像等。深度学习技术凭借其强大的数据处理能力,能够更好地处理这些多模态数据。通过实现不同模态之间的信息融合,深度学习模型能够捕捉到更加全面和丰富的医学影像信息,从而进一步提高医学影像分割的准确性和全面性。

在眼底图像分割领域,深度学习技术的应用不仅提高了分割精度和效率,还为眼科医生提供了更加可靠和全面的诊断依据。随着技术的不断进步和应用的深入拓展,深度学习技术有望在眼科医学领域发挥更加重要的作用,为患者带来更加精准和有效的治疗方案。

本次实践,我们采用 UNet 进行眼底血管医学图像分割任务。UNet 是一种被广泛应用于语义分割任务的网络结构,其编码器-解码器结构以及跳跃连接的设计,使其能够有效地捕获图像中不同尺度的特征信息,从而在眼底血管分割任务中取得较好的效果。同时,在推理阶段,UNet 采用全卷积网络结构,能够快速对新的眼底图像进行血管分割,为临床应用提供了实时性支持。

3.2 网络搭建

选用 U-Net 网络结构作为基础分割模型的原因在于其通过编解码器架构,有效地结合局部信息和全局信息,提高分割准确性;同时,U-Net 的跳跃连接结构有助于保留和恢复图像中的细节和边缘信息,且在小样本情况下表现优异,能够充分利用有限数据进行有效训练,广泛应用于医学图像分割任务中。网络架构如下

python 复制代码
class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        self.up1 = Up(1024, 256, bilinear)
        self.up2 = Up(512, 128, bilinear)
        self.up3 = Up(256, 64, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

3.3 网络训练

基于 PyTorch 的神经网络训练流程可以分为以下步骤(不考虑前期数据准备和模型结构):

定义损失函数 根据任务类型选择合适的损失函数(loss function),如分类任务常用的交叉熵损失(Cross-Entropy Loss)或回归任务中的均方误差(Mean Square Error)。

选择优化器 选择合适的优化器(optimizer),如随机梯度下降(SGD)、Adam 或 RMSprop,并设置初始学习率及其它优化参数。

训练模型 在训练过程中,通过迭代训练数据集来调整模型参数。每个迭代周期称为一个 epoch。对于每个 epoch,数据会被分成多个 batch,每个 batch 被输入到模型中进行前向传播、计算损失、反向传播更新梯度,并最终优化模型参数。

保存模型 当满足需求时,可以将训练好的模型保存下来,以便后续部署和使用。

根据这个步骤编写以下代码

python 复制代码
def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
    dataset = Dateset_Loader(data_path)
    per_epoch_num = len(dataset) / batch_size
    train_loader = torch.utils.data.DataLoader(dataset=dataset,
                                               batch_size=batch_size,
                                               shuffle=True)
    optimizer = optim.Adam(net.parameters(),lr=lr,betas=(0.9, 0.999),eps=1e-08, weight_decay=1e-08,amsgrad=False)
    criterion = nn.BCEWithLogitsLoss()
    best_loss = float('inf')
    loss_record = []
    with tqdm(total=epochs*per_epoch_num) as pbar:
        for epoch in range(epochs):
            net.train()
            for image, label in train_loader:
                optimizer.zero_grad()
                image = image.to(device=device, dtype=torch.float32)
                label = label.to(device=device, dtype=torch.float32)
                pred = net(image)
                loss = criterion(pred, label)
                pbar.set_description("Processing Epoch: {} Loss: {}".format(epoch+1, loss))
                if loss < best_loss:
                    best_loss = loss
                    torch.save(net.state_dict(), 'best_model.pth')
                loss.backward()
                optimizer.step()
                pbar.update(1)
            loss_record.append(loss.item())

    plt.figure()
    plt.plot([i+1 for i in range(0, len(loss_record))], loss_record)
    plt.title('Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.savefig('/root/shared-storage/results/training_loss.png')

运行这个脚本,可以在终端看到进度

训练损失函数如下,可以看到已经收敛

3.4 模型测试

测试逻辑如下所示,主要是计算 IoU 指标

python 复制代码
def cal_miou(test_dir="/root/workspace/DRIVE-SEG-DATA/Test_Images",
             pred_dir="/root/workspace/DRIVE-SEG-DATA/results", gt_dir="/root/workspace/DRIVE-SEG-DATA/Test_Labels",
             model_path='best_model_drive.pth'):
    name_classes = ["background", "vein"]
    num_classes = len(name_classes)
    if not os.path.exists(pred_dir):
        os.makedirs(pred_dir)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    net = UNet(n_channels=1, n_classes=1)
    net.to(device=device)
    net.load_state_dict(torch.load(model_path, map_location=device))
    net.eval()

    img_names = os.listdir(test_dir)
    image_ids = [image_name.split(".")[0] for image_name in img_names]

    time.sleep(1)
    for image_id in tqdm(image_ids):
        image_path = os.path.join(test_dir, image_id + ".png")
        img = cv2.imread(image_path)
        origin_shape = img.shape
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img = cv2.resize(img, (512, 512))
        img = img.reshape(1, 1, img.shape[0], img.shape[1])
        img_tensor = torch.from_numpy(img)
        img_tensor = img_tensor.to(device=device, dtype=torch.float32)
        pred = net(img_tensor)
        pred = np.array(pred.data.cpu()[0])[0]
        pred[pred >= 0.5] = 255
        pred[pred < 0.5] = 0
        pred = cv2.resize(pred, (origin_shape[1], origin_shape[0]), interpolation=cv2.INTER_NEAREST)
        cv2.imwrite(os.path.join(pred_dir, image_id + ".png"), pred)

    hist, IoUs, PA_Recall, Precision = compute_mIoU_gray(gt_dir, pred_dir, image_ids, num_classes, name_classes)
    miou_out_path = "/root/shared-storage/results/"
    show_results(miou_out_path, hist, IoUs, PA_Recall, Precision, name_classes)

模型保存的时候保存到共享存储路径/root/shared-storage,其他实例可以直接从共享存储中获取训练后的模型


相关推荐
碣石潇湘无限路35 分钟前
【AI篇】当Transformer模型开始学习《孙子兵法》
人工智能·学习
看到我,请让我去学习1 小时前
OpenCV开发-初始概念
人工智能·opencv·计算机视觉
汀沿河1 小时前
8.1 prefix Tunning与Prompt Tunning模型微调方法
linux·运维·服务器·人工智能
陈敬雷-充电了么-CEO兼CTO1 小时前
大模型技术原理 - 基于Transformer的预训练语言模型
人工智能·深度学习·语言模型·自然语言处理·chatgpt·aigc·transformer
学术 学术 Fun1 小时前
✨ OpenAudio S1:影视级文本转语音与语音克隆Mac整合包
人工智能·语音识别
风铃喵游2 小时前
让大模型调用MCP服务变得超级简单
前端·人工智能
旷世奇才李先生2 小时前
Pillow 安装使用教程
深度学习·microsoft·pillow
booooooty2 小时前
基于Spring AI Alibaba的多智能体RAG应用
java·人工智能·spring·多智能体·rag·spring ai·ai alibaba
PyAIExplorer2 小时前
基于 OpenCV 的图像 ROI 切割实现
人工智能·opencv·计算机视觉
风口猪炒股指标3 小时前
技术分析、超短线打板模式与情绪周期理论,在市场共识的形成、分歧、瓦解过程中缘起性空的理解
人工智能·博弈论·群体博弈·人生哲学·自我引导觉醒