pytorch从零开始用语义分割网络(deeplab3+)训练自己的数据集

Posted 2023-04-08 大黑山修道

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了pytorch从零开始用语义分割网络(deeplab3+)训练自己的数据集相关的知识，希望对你有一定的参考价值。

参考文档：
https://blog.csdn.net/qq_43631789/article/details/102700231
https://blog.csdn.net/qq_39056987/article/details/106455828
https://blog.csdn.net/qq_36766560/article/details/110009622
https://blog.csdn.net/gsgs1234/article/details/115267777

1. 制作自己数据集（labelme）

使用labelme对自己采集的图像进行标注。
安装和使用过程略。

然后标注的json文件默认保存在图像所在的目录内。标注完成之后一般是这样的：

2. 下载deeplab3+源码包（pytorch）

github地址：
https://github.com/jfzhang95/pytorch-deeplab-xception
下载代码包至自己的指定位置：

3. 将数据集转换为VOC格式

3.1 数据结构介绍

我们首先新建一些列文件夹，文件结构如下：

- ImageSets
	- Segmentation
		- train.txt
		-train.txt
		- val.txt
- JPEGImages
- SegmentationClass

ImageSets目录内单放一个Segmentation文件夹，然后Segmentation目录下需要制作3个txt文件：train.txt，train.txt，val.txt 用来表示训练集，验证集，测试集的划分信息。制作方式后面介绍。

JPEGImages 目录内用于存放图像数据集的原图。

SegmentationClass 目录放置的mask图像，mask是原图根据标注信息json文件生成的，生成方式后面介绍，注意mask的图像与原图的名称一一对应。

3.2 生成3个txt文件

txt的格式是每一行一个图像文件名，无后缀，不需要地址。

train,trainval,val自己按照一定比例划分
代码如下：

import os
import numpy as np
root = r"D:\\dataset\\belt\\JPEGImages"
output = r"D:\\dataset\\belt\\ImageSets\\Segmentation"
filename = []
#从存放原图的目录中遍历所有图像文件
# dirs = os.listdir(root)
for root, dir, files in os.walk(root):
    for file in files:
        print(file)
        filename.append(file[:-4])  # 去除后缀，存储


#打乱文件名列表
np.random.shuffle(filename)
#划分训练集、测试集，默认比例6:2:2
train = filename[:int(len(filename)*0.6)]
trainval = filename[int(len(filename)*0.6):int(len(filename)*0.8)]
val = filename[int(len(filename)*0.8):]

#分别写入train.txt, test.txt
with open(os.path.join(output,'train.txt'), 'w') as f1, open(os.path.join(output,'trainval.txt'), 'w') as f2,open(os.path.join(output,'val.txt'), 'w') as f3:
    for i in train:
        f1.write(i + '\\n')
    for i in trainval:
        f2.write(i + '\\n')
    for i in val:
        f3.write(i + '\\n')

print('成功！')

3.3 根据json，制作对应的mask图像

我们首先将所有的json文件存放到单独的文件夹，实例中表示为： "D:\\\\dataset\\\\json"

制作mask的图像需要用到labelme的源码。

我们首先找到labelme源码的安装位置：
用anaconda安装的话，windows一般是Users\\用户名\\.conda\\envs\\环境名\\Lib\\site-packages\\labelme。
然后找到labelme\\cli的位置，先备份一下原来的json_to_dataset.py文件，然后用下面的代码覆盖掉原来的json_to_dataset.py。

import argparse
import base64
import json
import os
import os.path as osp

import PIL.Image
import yaml

from labelme.logger import logger
from labelme import utils

path = "D:\\\\dataset\\\\json"
dirs = os.listdir(path)


def label(json_file, out_dir, label_name_to_value):
    # print("json.load(open(json_file))=", json_file)
    # json_file = os.path.join(path, json_file)
    # print("json.load(open(json_file))=", json_file)

    data = json.load(open(json_file))
    if data['imageData']:
        imageData = data['imageData']
    else:
        imagePath = os.path.join(os.path.dirname(json_file), data['imagePath'])
        with open(imagePath, 'rb') as f:
            imageData = f.read()
            imageData = base64.b64encode(imageData).decode('utf-8')
    img = utils.img_b64_to_arr(imageData)

    for shape in sorted(data['shapes'], key=lambda x: x['label']):
        label_name = shape['label']
        if label_name in label_name_to_value:
            label_value = label_name_to_value[label_name]
        else:
            label_value = len(label_name_to_value)
            label_name_to_value[label_name] = label_value
    lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)

    label_names = [None] * (max(label_name_to_value.values()) + 1)
    for name, value in label_name_to_value.items():
        label_names[value] = name
    lbl_viz = utils.draw_label(lbl, img, label_names)

    PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
    utils.lblsave(osp.join(out_dir, 'label.png'), lbl)
    PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))

    with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
        for lbl_name in label_names:
            f.write(lbl_name + '\\n')

    logger.warning('info.yaml is being replaced by label_names.txt')
    info = dict(label_names=label_names)
    with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
        yaml.safe_dump(info, f, default_flow_style=False)

    logger.info('Saved to: '.format(out_dir))


def main():
    logger.warning('This script is aimed to demonstrate how to convert the'
                   'JSON file to a single image dataset, and not to handle'
                   'multiple JSON files to generate a real-use dataset.')

    parser = argparse.ArgumentParser()
    parser.add_argument('json_file_dir')
    parser.add_argument('-o', '--out', default=None)
    args = parser.parse_args()
    label_name_to_value = '_background_': 0

    for json_file in dirs:
        # print("json_file=", json_file)
        if args.out is None:
            json_file = os.path.join(path, json_file)
            out_dir = osp.basename(json_file).replace('.', '_')
            out_dir = osp.join(osp.dirname(json_file), out_dir)
        else:
            # out_dir = args.out
            json_file = os.path.join(path, json_file)
            out_dir = osp.basename(json_file).replace('.', '_')
            out_dir = osp.join(osp.dirname(args.out), out_dir)
            # print('out_dir=',out_dir)
        if not osp.exists(out_dir):
            os.mkdir(out_dir)

        label(json_file, out_dir, label_name_to_value)


if __name__ == '__main__':
    main()

然后在cli目录输入命令行：

python .\\json_to_dataset.py json_file_dir -o D:\\dataset\\output_mask\\

json_file_dir :必须带的参数
-o ：输出目录
然后输出目录就会增加非常多的文件夹：

每个文件夹下有5个文件，而我们需要将所有的label.png文件放入到VOC格式数据集中的SegmentationClass目录，并且需要改成与原图对应的名称。

类别放在label_names.txt里，默认会有一个_background_。所以我们正常的语义分割至少要有两类。

抽离label.png的脚本：

import os
import shutil

inputdir = 'D:\\\\dataset\\\\output_mask'
outputdir = 'D:\\\\dataset\\\\belt\\\\SegmentationClass'

for dir in os.listdir(inputdir):
    # 设置旧文件名（就是路径+文件名）
    oldname = inputdir + os.sep + dir + os.sep + 'label.png'  # os.sep添加系统分隔符
    print("oldname=",oldname)
    png_id = ''
    print("dir=",oldname)
    # 之前的mask命名都是以xxx_json的目录明明，现在需要把最后的_json部分去除，还原原图的名称
    for i in range(len(dir.split('_'))) :
        if i == len(dir.split('_'))-1:
            continue
        else:
            # 原图的名称中可有可能包含一个或多个_，除了最后一个_，其余都保留。
            if png_id != '':
                png_id += '_'
            png_id += dir.split('_')[i]

    # 设置新文件名
    newname = outputdir + os.sep + png_id + '.png'

    shutil.copyfile(oldname, newname)  # 用os模块中的rename方法对文件改名
    print(oldname, '======>', newname)

执行结束后就可以发现SegmentationClass目录内放入了mask图像。

4. 修改deeplab+源码，增加自己的数据集

4.1 mypath.py 中加入自己数据集的路径

实例中增加的数据集名称为belt

belt下的文件结构就是之前提到的VOC结构：

-belt
	- ImageSets
		- Segmentation
			- train.txt
			-train.txt
			- val.txt
	- JPEGImages
	- SegmentationClass

4.2 在dataloaders/datasets目录下添加文件

复制一份pascal.py文件，并重命名为自己的数据集名称

然后打开自己的数据集py文件，修改文件内的类别数和数据集名称：

4.3 修改dateloaders目录下`utils.py`

搜素def get_cityscapes_labels()函数，然后在上方添加自己数据集的函数，例如get_belt_labels().
这个函数的主要意思就是给自己每个类设置一个掩膜颜色，有多少个类，就设置多少种颜色。

然后在decode_segmap函数内添加代码，其中n_classes是你要分割的类别数

4.4 在dataloaders目录下修改`init.py`

在第一行添加数据集名称，复制’pascal’数据集描述，把名称修改为自己数据集的名字

    if args.dataset == 'belt':
        train_set = belt.VOCSegmentation(args, split='train')
        val_set = belt.VOCSegmentation(args, split='val')

        num_class = train_set.NUM_CLASSES
        train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs)
        val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, **kwargs)
        test_loader = None

        return train_loader, val_loader, test_loader, num_class

4.5 在同级目录中修改train.py约185行添加自己数据集的名称(可以设置为默认)

5. 开始训练数据

输入指令：

python train.py --backbone mobilenet --lr 0.007 --workers 1 --epochs 50 --batch-size 8 --gpu-ids 0 --checkname deeplab-mobilenet

模型保存的路径是在代码内是设置的，在saver.py可以看到保存的路径：run/[datasetname]/[checkname]，在示例中的路径就是：run/belt/deeplab-mobilenet
保存的目录中可能会存在很多experiment_*的目录，这是每一此训练都会保存在一个experiment_的目录内，最新的训练结果保存在最后id的目录上。此外，最优的模型还会保存到run/[datasetname]/[checkname]中的model_best.pth.tar中。

若是出现报错：
AttributeError: ‘DeepLab’ object has no attribute ‘module’

解决方式：
打开train.py，找到报错的地方 'state_dict': self.model.module.state_dict(),修改为

`'state_dict': self.model.state_dict()

6. 测试

源码中没有测试代码，需要自己放入一个测试py文件。
修改–in-path为数据集的测试图片，最后的结果保存在–in-path中

#
# demo.py
#
import argparse
import os
import numpy as np
import time
 
from modeling.deeplab import *
from dataloaders import custom_transforms as tr
from PIL import Image
from torchvision import transforms
from dataloaders.utils import  *
from torchvision.utils import make_grid, save_image
 
def main():
 
    parser = argparse.ArgumentParser(description="PyTorch DeeplabV3Plus Training")
    parser.add_argument('--in-path', type=str,  default='/root/home/zyx/Seg552_VOC/test',
                        help='image to test')
    # parser.add_argument('--out-path', type=str, required=True, help='mask image to save')
    parser.add_argument('--backbone', type=str, default='resnet',
                        choices=['resnet', 'xception', 'drn', 'mobilenet'],
                        help='backbone name (default: resnet)')
    parser.add_argument('--ckpt', type=str, default='deeplab-resnet.pth',
                        help='saved model')
    parser.add_argument('--out-stride', type=int, default=16,
                        help='network output stride (default: 8)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--gpu-ids', type=str, default='0',
                        help='use which gpu to train, must be a \\
                        comma-separated list of integers only (default=0)')
    parser.add_argument('--dataset', type=str, default='belt',
                        choices=['pascal', 'coco', 'cityscapes','belt'],
                        help='dataset name (default: pascal)')
    parser.add_argument('--crop-size', type=int, default=513,
                        help='crop image size')
    parser.add_argument('--num_classes', type=int, default=2,
                        help='crop image size')
    parser.add_argument('--sync-bn', type=bool, default=None,
                        help='whether to use sync bn (default: auto)')
    parser.add_argument('--freeze-bn', type=bool, default=False,
                        help='whether to freeze bn parameters (default: False)')
 
    args = parser.parse_args()
    args.cuda = not args.no_cuda and torch.cuda.is_available()
    if args.cuda:
        try:
            args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')]
        except ValueError:
            raise ValueError('Argument --gpu_ids must be a comma-separated list of integers only')
 
    if args.sync_bn is None:
        if args.cuda and len(args.gpu_ids) > 1:
            args.sync_bn = True
        else:
            args.sync_bn = False
    model_s_time = time.time()
    model = DeepLab(num_classes=args.num_classes,
                    backbone=args.backbone,
                    output_stride=args.out_stride,
                    sync_bn=args.sync_bn,
                    freeze_bn=args.freeze_bn)
	model = nn.DataParallel(model)

    ckpt = torch.load(args.ckpt, map_location='cpu')
    model.load_state_dict(ckpt['state_dict'])
    model = model.cuda()
    model_u_time = time.time()
    model_load_time = model_u_time-model_s_time
    print("model load time is ".format(model_load_time))
 
    composed_transforms = transforms.Compose([
        tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        tr.ToTensor()])
    for name in os.listdir(args.in_path):
        s_time = time.time()
        image = Image.open(args.in_path+"/"+name).convert('RGB')
 
        # image = Image.open(args.in_path).convert('RGB')
        target = Image.open(args.in_path+"/"+name).convert('L')
        sample = 'image': image, 'label': target
        tensor_in = composed_transforms(sample)['image'].unsqueeze(0)
 
        model.eval()
        if args.cuda:
            tensor_in = tensor_in.cuda()
        with torch.no_grad():
            output = model(tensor_in)
 
        grid_image = make_grid(decode_seg_map_sequence(torch.max(output[:3], 1)[1].detach().cpu().numpy()),
                                3, normalize=False, range=(0, 255))
        save_image(grid_image,args.in_path+"/"+"_mask.png".format(name[0:-4]))
        u_time = time.time()
        img_time = u_time-s_time
        print("image: time:  ".format(name,img_time))
        # save_image(grid_image, args.out_path)
        # print("type(grid) is: ", type(grid_image))
        # print("grid_image.shape is: ", grid_image.shape)
    print("image save in in_path.")
if __name__ == "__main__":
   main()
 
# python demo.py --in-path your_file --out-path your_dst_file

注意点：
参数--dataset代码中加入自己的类：

    parser.add_argument('--dataset', type=str, default='belt',
                        choices=['pascal', 'coco', 'cityscapes','belt'],
                        help='dataset name (default: pascal)')

输入测试指令：

 python testdemo.py --dataset belt --num_classes 2 --ckpt run/Seg552/deeplab-mobilenet/checkpoint.pth.tar --backbone mobilenet

常见报错：

1. state_dict错误

RuntimeError: Error(s) in loading state_dict for DeepLab:         Missing key(s) in state_dict: "

在原版中会出现，在ckpt = torch.load(args.ckpt, map_location='cpu')之前加入 model = nn.DataParallel(model) 即可。

2. CUDA报错

报错如下：raise AssertionError("Torch not compiled with CUDA enabled") AssertionError: Torch not compiled with CUDA enabled
解决方式：

可能是在gpu上训练，用了torch-cpu，检查torch版本。
如果版本没问题，注释掉：

    model = model.cuda()

干货｜使用PyTorch从零开始构建Elman循环神经网络

摘要：循环神经网络是如何工作的？如何构建一个Elman循环神经网络？在这里，教你手把手创建一个Elman循环神经网络进行简单的序列预测。

本文以最简单的RNNs模型为例：Elman循环神经网络，讲述循环神经网络的工作原理，即便是你没有太多循环神经网络（RNNs）的基础知识，也可以很容易的理解。为了让你更好的理解RNNs，我们使用Pytorch张量包和autograd库从头开始构建Elman循环神经网络。该文中完整代码在Github上是可实现的。

在这里，假设你对前馈神经网络略有了解。Pytorch和autograd库更为详细的内容请查看我的其他教程。

干货｜使用PyTorch从零开始构建Elman循环神经网络

Elman循环神经网络

Jeff Elman首次提出了Elman循环神经网络，并发表在论文《Finding structure in time》中：它只是一个三层前馈神经网络，输入层由一个输入神经元x1和一组上下文神经元单元{c1 ... cn}组成。隐藏层前一时间步的神经元作为上下文神经元的输入，在隐藏层中每个神经元都有一个上下文神经元。由于前一时间步的状态作为输入的一部分，因此我们可以说，Elman循环神经网络拥有一定的内存——上下文神经元代表一个内存。

预测正弦波

现在，我们来训练RNNs学习正弦函数。在训练过程中，一次只为模型提供一个数据，这就是为什么我们只需要一个输入神经元x1，并且我们希望在下一时间步预测该值。输入序列x由20个数据组成，并且目标序列与输入序列相同。

干货｜使用PyTorch从零开始构建Elman循环神经网络

模型实现

首先导入包。

干货｜使用PyTorch从零开始构建Elman循环神经网络

接下来，设置模型的超参数。设置输入层的大小为7（6个上下文神经元和1个输入神经元），seq_length用来定义输入和目标序列的长度。

干货｜使用PyTorch从零开始构建Elman循环神经网络

生成训练数据：x是输入序列，y是目标序列。

干货｜使用PyTorch从零开始构建Elman循环神经网络

创建两个权重矩阵。大小为（input_size，hidden_size）的矩阵w1用于隐藏连接的输入，大小为（hidden_size，output_size）的矩阵w2用于隐藏连接的输出。用零均值的正态分布对权重矩阵进行初始化。

干货｜使用PyTorch从零开始构建Elman循环神经网络

定义forward方法，其参数为input向量、context_state向量和两个权重矩阵，连接input和context_state创建xh向量。对xh向量和权重矩阵w1执行点积运算，然后用tanh函数作为非线性函数，在RNNs中tanh比sigmoid效果要好。然后对新的context_state和权重矩阵w2再次执行点积运算。我们想要预测连续值，因此这个阶段不使用任何非线性。

请注意，context_state向量将在下一时间步填充上下文神经元。这就是为什么我们要返回context_state向量和out。

干货｜使用PyTorch从零开始构建Elman循环神经网络

训练

训练循环的结构如下：

1.外循环遍历每个epoch。epoch被定义为所有的训练数据全部通过训练网络一次。在每个epoch开始时，将context_state向量初始化为0。

2.内部循环遍历序列中的每个元素。执行forward方法进行正向传递，该方法返回pred和context_state，将用于下一个时间步。然后计算均方误差（MSE）用于预测连续值。执行backward()方法计算梯度，然后更新权重w1和w2。每次迭代中调用zero_()方法清除梯度，否则梯度将会累计起来。最后将context_state向量包装放到新变量中，以将其与历史值分离开来。