这段代码使用 PyTorch 评估神经网络模型在给定数据集上的性能,使用的评价指标是 Dice 系数。

首先,定义了一个名为 eval_net 的函数,该函数接受以下参数:

  • net: 神经网络模型
  • dataset: 数据集
  • slicetotal: 数据集中切片的总数
  • batch_size: 批量大小,默认为 12
  • gpu: 是否使用 GPU,默认为 True

函数的作用是在不使用 densecrf 的情况下计算 Dice 系数。

接下来,函数开始执行。首先将模型设置为评估模式 (net.eval()),然后初始化一些变量,包括 dice_, jac_, NEJNE,它们都是具有 14 个元素的张量,并将它们移动到 GPU 上。

然后,使用 torch.no_grad() 上下文管理器来禁用梯度计算。接下来是一个循环,用于遍历数据集中的批次。在每个批次中,首先从批次中提取图像和真实掩膜,并将它们转换为张量格式。然后,创建一个与真实掩膜相同形状的全零张量 pre_masks_evaltrue_masks_eval,用于存储预测掩膜和真实掩膜的结果。

接下来,将图像和真实掩膜移动到 GPU 上(如果使用 GPU)。然后,将图像输入到神经网络模型中,得到输出图像。对输出图像进行处理,得到预测掩膜。接下来,对每个类别进行循环,并计算该类别的 Dice 系数和 Jaccard 系数。最后,将批次中每个切片的 Dice 系数和 Jaccard 系数累加到 dice_jac_ 张量中。

在循环结束后,计算每个类别的平均 Dice 系数和 Jaccard 系数,然后将结果返回。同时,计算整个评估过程所花费的时间,并打印出来。

最后,将 dice_jac_ 张量作为函数的返回值。

def eval_net(net, dataset, slicetotal, batch_size=12, gpu=True):
    '''Evaluation without the densecrf with the dice coefficient'''

    net.eval()
    start = time.time()
    dice_ = torch.zeros(14).cuda()
    jac_ = torch.zeros(14).cuda()
    NE = torch.zeros(14).cuda()
    JNE = torch.zeros(14).cuda()
    ia = 0
    #dice_all = torch.zeros(14).cuda()
    #NE = torch.Tensor([0,3241,3241,3611,3611,2911,2911,3201,3161,3161,1110,3530,3530,3281]).cuda()
    print(1)
    with torch.no_grad():
        for i, b in enumerate(batch(dataset, batch_size)):
            #print('eval',i)


            imgs = np.array([k[0] for k in b]).astype(np.float32)
            true_masks = np.array([k[1] for k in b])

            imgs = torch.from_numpy(imgs)
            imgs = imgs.unsqueeze(1)
            true_masks = torch.from_numpy(true_masks)
            #imgs = imgs.permute(1, 2, 0)
            #imgs = imgs.unsqueeze(0).unsqueeze(0)
            pre_masks_eval = torch.zeros(true_masks.shape[0],14,256,256)
            true_masks_eval = torch.zeros(true_masks.shape[0],14,256,256)#这里考虑把256换成一个变量,这个变量等于标签的宽高
            batchshape = true_masks.shape[0]

            batch_dice = torch.zeros(14).cuda()
            if gpu:
                imgs = imgs.cuda()
                true_masks = true_masks.cuda()
                net.cuda()

            output_img = net(imgs)
            #pre_masks = output_img.cuda()
            input = output_img.cuda()
            #input = input.squeeze(0).permute(3, 0, 1, 2) #2d网络不需要这样的调整
            #pre_masks = F.softmax(pre_masks,dim=1)
            pre_masks = input.max(1)[1].float() #索引代表像素所属类别的数>字
            for ak in range(14):
                if ak == 0:
                    continue
                pre_masks_eval[:,ak] = (pre_masks==ak)
                true_masks_eval[:,ak] = (true_masks==ak)
                premasks = pre_masks_eval[:,ak].view(true_masks.shape[0],-1)
                truemasks = true_masks_eval[:,ak].view(true_masks.shape[0],-1)

                intersection = premasks * truemasks
                TP = intersection.sum(1)
                FP = premasks.sum(1) - TP
                FN = truemasks.sum(1) - TP

                for bk in range(true_masks.shape[0]):
                    if TP[bk] == 0 and FP[bk] == 0 and FN[bk] == 0:
                        NE[ak] += 1
                        JNE[ak] += 1
                    else:
                        batch_dice[ak] = batch_dice[ak] + 2*TP[bk] / (2*TP[bk] + FP[bk] + FN[bk])
                        jac_[ak] = jac_[ak] + TP[bk] / (TP[bk] + FP[bk] + FN[bk])
                #batch_dice[ak] = (2*TP.float() / (2*TP.float() + FP.float() + FN.float() + 0.01)).sum()
                #if ak>6:     
                #    print(batch_dice[ak])

            dice_ = dice_ + batch_dice
            #ia += 1
        #print(i*12)
        for knum in range(14):
            dice_[knum] = dice_[knum] / (slicetotal - NE[knum])
            jac_[knum] = jac_[knum] / (slicetotal - JNE[knum])
            #print(NE[knum])
    end = time.time()
    print('time used:',end - start)
    #print(dice_all)
    #return dice_all
    return dice_,jac_ 
PyTorch 神经网络模型评估代码详解:Dice 系数计算

原文地址: https://www.cveoy.top/t/topic/fRSg 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录