deeplabv3+——训练自己的数据集 torch1.12.0 cuda11.3

目录

前言

环境

源码

参考博客

一、制作自己的数据集

二、训练

三、可视化

---

前言

环境

torch==1.12.0+cu113

cuda==11.3

显卡为 RTX3070ti

tips：30系显卡最好还是用配套的cuda，博主原本使用的cuda10.0版本，一直报错很坑

源码

pytorch-deeplab-xception

参考博客

windows10】使用pytorch版本deeplabv3+训练自己数据集

DeeplabV3+ 训练自己的数据集。pytorch

一、制作自己的数据集

参考博客二中讲的很清楚，重点是

1、mask为单通道，所以要将标签图二值化

2、label的像素为1、2、3.... 例如你要做一个二分类的任务，那么背景的像素值为0，标签的像素 值为1；如果要做三分类的任务，那么背景的像素值为0，标签1的像素值为1，标签2的像素值为2，以此类推

import os
import cv2
from PIL import Image

def white_label_to_mask():
    # 获取目录下所有图片名
    filename = os.listdir(r"E:\..")
    print(filename)
    base_dir = r"E:\.." # input
    new_dir  = r"E:\.."  # output
    for img in filename:
        name = img
        path1 = os.path.join(base_dir ,img)
        img = cv2.imread(path1)
        #二值化
        Grayimg = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, thresh = cv2.threshold(Grayimg,170, 1 ,cv2.THRESH_BINARY)
        cv2.imwrite('img.png', thresh)
        image = Image.open('img.png')
        path = os.path.join(new_dir ,name)
        image.save(path)
white_label_to_mask()

这里贴一个将黑白（三通道）的标签图（目标为白色）批量转化为mask的脚本 可根据需要自行更改

二、训练

python train.py 

三、可视化

运行结果会存放在run中，这里我运用的mobilnet，所以结果存放在了../pytorch-deeplab-xception-master\run\mydata\deeplab-mobilenet下

在当前目录下cmd，并调用tensorboard

tensorboard --logdir=='experiment_0'

然后出现如下结果

....
TensorFlow installation not found - running with reduced feature set.
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.10.0 at http://localhost:6006/ (Press CTRL+C to quit)

复制http://localhost:6006/到网页框  即可查看结果