MODNet模型部署以及一些改动

部署日记：MODNet的使用以及一些改动

文献和git

git地址：https://github.com/ZHKKKe/MODNet
论文地址：https://arxiv.org/abs/1703.03872
算法的具体内容可以看原论文或者CSDN上相关文章进行学习。本文主要介绍在进行训练和预测的时候遇到的一些事情和解决办法

论文与github的一些注意事项

官方的MODNet提供的算法中e-ASPP模块使用了se-block模块来代替。
百度Paddle的matting算是MODNet算法的复现，按照modnet的理念建立的模型，在原著基础上提供了多个不同主干网络的预训练模型如RestNet50_vd、HRNet_w18等。来满足用户在边缘端、服务端等不同场景部署的需求。
!!!
!!!
本文使用的模型是一个大佬补全的：https://github.com/actboy/MODNet
他的博客链接：link
!!!
!!!

改动

训练

训练相关文件有：
trainer.py
matting_dataset.py
setting.py

需要可以自行整理成一个文件

trainer.py

从新训练/读取预训练模型finetuning
从新训练：

modnet = torch.nn.DataParallel(MODNet())  # .cuda(）

读取预训练模型finetuning：

model = MODNet().to(device)
model_dict = model.state_dict()
pretrained_dict = torch.load('./pretrained/training/modnet_custom_portrait_matting_20_th.ckpt', map_location=device)
try:
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in pretrained_dict.items():
        k = k.replace("module.", "")
        new_state_dict[k] = v
except:
    pass
new_state_dict = {k: v for k, v in new_state_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(new_state_dict)
model.load_state_dict(model_dict)
modnet = torch.nn.DataParallel(model)  # .cuda()

关于finetuning

Torch.load(‘          ’,map_location=device)

空白字段填写模型保存的路径和名称
加上try字段缘由：
官方保存的模型可以正常使用，但自行训练的模型的dict当中，items会含有‘module’字段，读取模型时会出错，原因是自行训练的模型的dict当中含有‘module’，导致key值不对等。
保存的代码：

torch.save(modnet.state_dict(),.........)

表明了加载过程中，期望获得的key值为feature…，而不是module.features…。这是由模型保存过程中导致的，模型应该是在DataParallel模式下面，也就是采用了多GPU训练模型，然后直接保存的，所以会多出‘module.’这个。
相关链接：
https://blog.csdn.net/qq_32998593/article/details/89343507
https://blog.csdn.net/weixin_40522801/article/details/106563354

优化器

#optimizer = torch.optim.SGD(modnet.parameters(), lr=LR, momentum=0.9)
optimizer = torch.optim.Adam(modnet.parameters())

原文使用的SGD优化器，这里改用为Adam自适应动量的随机优化方法，个人认为用这个比SGD的自己去定义LR会更好。

每个epoch的训练时间

开头记得

import time

for epoch in range(0, EPOCHS):
    print(f'epoch: {epoch}/{EPOCHS - 1}')
    b_time = time.time()
    for idx, (image, trimap, gt_matte) in enumerate(dataloader):
        semantic_loss, detail_loss, matte_loss = \
            supervised_training_iter(modnet, optimizer, image, trimap, gt_matte,
                                     semantic_scale=SEMANTIC_SCALE,
                                     detail_scale=DETAIL_SCALE,
                                     matte_scale=MATTE_SCALE)
        print(f'{(idx + 1) * BS}/{len(mattingDataset)} --- '
              f'semantic_loss: {semantic_loss:f}, detail_loss: {detail_loss:f}, matte_loss: {matte_loss:f}\r',
              end='')
    e_time = time.time()
    epochtime = e_time - b_time
    print('第',epoch,'个epoch用时', epochtime)

官方的训练代码中没有记录每一代的训练时长，训练的时间跨度与显卡型号，数量，BS，数据集等有复杂的关联，记录时长可以较为清晰的知道计算的功耗与时间的关系，并且可以大致预测出一个总体的训练时间。

损失记录可视化
下包：

pip install tensorboardX

开头：

from tensorboardX import SummaryWriter

开始训练之前：

writer = SummaryWriter('./log')
result_path = './result_path/'

Writer= SummaryWriter()声明一下自己将loss写到哪个路径下面

if epoch % SAVE_EPOCH_STEP == 0 and epoch > 1:
    torch.save(modnet.state_dict(), f'pretrained/training/animal_portrait_matting_{epoch}_th.ckpt')
print(f'{len(mattingDataset)}/{len(mattingDataset)} --- '
      f'semantic_loss: {semantic_loss:f}, detail_loss: {detail_loss:f}, matte_loss: {matte_loss:f}')
niter = epoch * len(dataloader) + idx
writer.add_scalars(result_path + 'semantic_loss', {result_path + 'semantic_loss': semantic_loss.data.item()},
                   niter)
writer.add_scalars(result_path + 'detail_loss', {result_path + 'detail_loss': detail_loss.data.item()},
                   niter)
writer.add_scalars(result_path + 'matte_loss', {result_path + 'matte_loss': matte_loss.data.item()},
                   niter)

然后终端启动tensorboard可视化：

tensorboard --logdir ./log --port 8890

更详细操作参考链接：https://blog.csdn.net/u013517182/article/details/93043942

保存模型
原版：

改后：

删除了‘epoch’，‘model_state_dict’，‘optimizer_state_dict’
原因是保存了之后再finetuning的时候会读取模型失败，把这些删掉后就可以使用，并且这样能够节省一定的空间。删掉后每个模型的大小大概是原来的80%。

保存最优模型

loss_best = 3

# 保存最优模型
loss_epoch = semantic_loss.item() + detail_loss.item() + matte_loss.item()
print('本代损失=', loss_epoch)
if loss_epoch < loss_best:
    torch.save(modnet.state_dict(), f'pretrained/animalmatting_best.ckpt')
    print('本代loss更低，保存为最优模型，代数为',epoch)
    loss_best = loss_epoch

注意：
semantic_loss和detail_loss和matte_loss是tensor()格式，计算数值要在之后加上.item()

这种方式也有不科学之处，因为在测试集上损失最小的那个epoch的模型不一定就是最好的模型，不一定具有最好的泛化能力。

matting_dataset.py

读取数据集：

class MattingDataset(Dataset):    
    def __init__(self,
                 # dataset_root_dir='F:\\dataset\\MYdataset\\personal_Sundries',
                 dataset_root_dir='.\\src\\dataset\\PPM-100',
                 transform=None):
        image_path = dataset_root_dir + '\\train\\fg\\*'
        matte_path = dataset_root_dir + '\\train\\alpha\\*'
        image_file_name_list = glob(image_path)
        matte_file_name_list = glob(matte_path)

主函数：

if __name__ == '__main__':

    # test MattingDataset.gen_trimap
    matte = Image.open('D:/Project/MODNet-master/example/sample1.jpg')
    trimap1 = GenTrimap().gen_trimap(matte)
    trimap1 = np.array(trimap1) * 255
    trimap1 = np.uint8(trimap1)
    trimap1 = Image.fromarray(trimap1)
    trimap1.save('test_trimap.png')

会生成一张trimap图，以供测试用。

| |
生成的trimap：

setting.py

BS = 8  # BATCH SIZE
LR = 0.001  # LEARN RATE
EPOCHS = 100  # TOTAL EPOCHS

SEMANTIC_SCALE = 10.0
DETAIL_SCALE = 10.0
MATTE_SCALE = 1.0

SAVE_EPOCH_STEP = 20

上述的LR因为前面改成了adam优化器，可以忽略了，其他的根据自己的配置进行改动。
本人用的服务器是4张10G的RTX 3080，平时用两张卡跑，上面是我的参数。
训练的经历：
因为中途别的项目有使用服务器的需要，或者会有一些以外的事情导致训练终止，然后再次进行finetuning的时候，loss的曲线会发生震荡，大概在5-10个epoch之后恢复正常的下降的趋势，猜测是因为加载模型的时候导致精度丢失的情况，目前还没有去做解决,本人使用的数据集数量：训练集34738张,验证集2721张图片，总量为25.7G，大概一个epoch是42分钟。所以解决finetuning震荡问题是比较有必要的，不然会浪费大致4-7小时的训练时间。

预测

infer.py

读取模型路径：

if __name__ == '__main__':
    # create MODNet and load the pre-trained ckpt

    modnet = MODNet(backbone_pretrained=False)
    modnet = nn.DataParallel(modnet)

    ckp_pth = 'pretrained/matting1017.ckpt'
    if torch.cuda.is_available():
        modnet = modnet.cuda()
        weights = torch.load(ckp_pth)
    else:
        weights = torch.load(ckp_pth, map_location=torch.device('gpu'))
    modnet.load_state_dict(weights)

在ckp_pth填写预测模型的路径。
思考：训练的时候说过模型是在DataParallel模式下保存的，finetune会遇到问题，是否可以将整个模型的训练，保存，读取，都取消DataParallel的模式，这样会不会一劳永逸的解决。(我还没这么做
图片文件夹路径：

	pth = './examplepeople'
    dirlist = os.listdir(pth)

原版写的是单个目标文件，本人改动成目标文件夹里所有的图片。
原作者的代码中只给出遮罩matte，没有抠图结果，本代码添加了背景透明的前景图，以及加了统一背景颜色的图片（绿幕）

        img1 = cv2.imread(pth+ '/' +i)
        imgname_1 = os.path.basename(pth+ '/' +i)
        imgname_2 = os.path.splitext(imgname_1)
        imgname = imgname_2[0]
        prd_img.save('output/'+imgname+'.png')
        mask = cv2.imread("output/" + imgname+'.png', 0)  # 读取灰度图像

        height, width, channel = img1.shape

        b, g, r = cv2.split(img1)

    # -----------------1.获取透明的前景图像-----------
        dstt = np.zeros((4, height, width), dtype=img1.dtype)

        dstt[0][0:height, 0:width] = b
        dstt[1][0:height, 0:width] = g
        dstt[2][0:height, 0:width] = r
        dstt[3][0:height, 0:width] = mask
        cv2.imwrite("output/"+imgname+"_fore.png", cv2.merge(dstt))


    # -----------------2.与新背景图像合成-----------
        bg = np.zeros((3, height, width), dtype=img1.dtype)  # 生成背景图像
        bg[0][0:height, 0:width] = 0
        bg[1][0:height, 0:width] = 255
        bg[2][0:height, 0:width] = 0
    # 背景图像采用某一颜色

        dstt = np.zeros((3, height, width), dtype=img1.dtype)

        for i in range(3):
            dstt[i][:, :] = bg[i][:, :] * (255.0 - mask) / 255
            dstt[i][:, :] += np.array(img1[:, :, i] * (mask / 255), dtype=np.uint8)
        cv2.imwrite("output/"+imgname+"_merge.png", cv2.merge(dstt))

根据此链接提供的代码改动后的结果：
https://blog.csdn.net/qianbin3200896/article/details/87934119

后处理：
在模型不够优秀的情况下，可能会有一些图片做不到完整的抠出来，有一部分抠图的结果是目标外还有其他的区域会保留一些小的图像块，目前针对这一种情况的解决方法是使用保存最大连通域的方法去二次抠图（与第一次不同）
目前的方法：
Matting得到前景遮罩IM1----二值化----保存最大连通域IM2—把IM2作为遮罩图对IM1进行抠图得到IM3—用IM3对原图进行抠图得到前景。
注意：
此方法属于杯水车薪，治标不治本。归根结底还是缺少一个更好的模型。所以重点是在于数据集的内容。

代码

最后放上自己改动后的代码：
CSDN链接：https://download.csdn.net/download/weixin_52314265/86931556
百度云链接：https://pan.baidu.com/s/1xUsaQIgo7ce-I3LB_Fp0XQ （提取码：1015）
（本人在同一环境下去部署了很多抠图算法，所以requirements.txt按照原版安装，然后看上文补充）

数据集

关于数据集我都是找的能公开下载的、精度较高的数据集，还有一些需要私信作者才能获得，再次就不写上了。

人体	链接
PPM-100	https://github.com/ZHKKKe/PPM#download
AIM-500	https://github.com/JizhiziLi/AIM （500张，比较杂，但质量很好）
DUT	http://saliencydetection.net/duts/#org3aad434(该数据集还有人体动物物品等各种图片，下载后可以使用paddle的分类网络，添加几行分类保存的代码进行分类)
P3M-10k	https://github.com/JizhiziLi/P3M(训练集都是打了马赛克的，验证集没有马赛克)

动物	链接
AM-2k	https://github.com/JizhiziLi/GFM#am-2k
DUT	http://saliencydetection.net/duts/#org3aad434

其他	链接
IM	http://www.alphamatting.com/datasets.php
DUT	http://saliencydetection.net/duts/#org3aad434

总结

想要得到一个更好的抠图效果取决于一个更好的模型，想要一个更好的模型有几种因素：数据集，硬件设备。
数据集方面：
数量上可能稍显不足，本人使用的人像数据集虽有上万张，但实际是用几千张图片进行数据增强扩充出来的（并且大多为单人的图片，如果是背景较为干净的多人图片，训练出的模型还能抠出来，但都是少数，所以还需要多人图像的数据集进行训练。
硬件设备：
本人虽然使用的是3080，然而内存只有10G，在使用nn.parallel 进行训练时，无法提高BS，不然会爆显存，但是BS=16，显存又没有占满，三分之二都没有，算是比较尴尬的情况（不管单卡训练多卡训练都必须全部输入到同一个显卡当中然后再分配给其他显卡）。
可尝试的解决办法：
1.更高的型号，更大的显存可以更快更好地训练。
2.因为网络在反向传播的时候，计算loss的梯度默认都在主卡上计算。因此会比其他显卡多用一些显存，具体多用多少，主要还要看网络的结构。
借用下transformer-xl中用到的 BalancedDataParallel类
代码出处：https://github.com/kimiyoung/transformer-xl/blob/44781ed21dbaec88b280f74d9ae2877f52b492a5/pytorch/utils/data_parallel.py#L52
BalancedDataParallel继承了 torch.nn.DataParallel，之后通过自定义0卡batch_size的大小gpu0_bsz，即让0卡少一点数据。均衡0卡和其他卡的显存占用。
参考链接：https://blog.csdn.net/weixin_43922901/article/details/106117774
注意：
以上解决办法，本人并没有进行操作，成本和时间上的不足导致本人只能提供出意见而不能进行实践来证明方案可行性，若有大能可以尝试一下使用BalancedDataParallel改进并成功的话，希望可以也分享一下谢谢！