模型训练
环境构建
1. 创建虚拟环境
conda create -n hrnet python=3.7
conda activate hrnet2. 安装cuda和cudnn
conda install cudatoolkit=10.2
conda install cudnn3. 安装pytorch
pip install torch==1.7.0
pip install torchvision==0.8.04. 下载项目代码
git clone https://github.com/HRNet/HRNet-Semantic-Segmentation.git
cd HRNet-Semantic-Segmentation-HRNet-OCR或者直接从 https://github.com/HRNet/HRNet-Semantic-Segmentation/tree/HRNet-OCR 处下载zip包
5. 安装项目依赖的其他包
pip install -r requirements.txt数据整理
1. 标注
使用labelme或其他标注工具做语义分割标注，网上关于labelme标注的教程很多，且操作并不复杂，这里不进行赘述
2. 数据集格式整理
按照cityscapes数据集格式，整理自己的数据，格式如下：
项目目录/data
├── cityscapes
| ├── gtFine
| | ├── train # 训练集标签，即label图片，由类别索引构成
| | └── val # 验证集标签，即label图片，由类别索引构成
| └── leftImg8bit
| ├── train # 训练集原始图像
| └── val # 验证集原始图像
└── list
└── cityscapes
├── train.lst # 训练集 原始图像路径\t标签路径
└── val.lst # 验证集 原始图像路径 \t标签路径train.lst 或 val.lst 中的内容示例如下：
cityscapes/leftImg8bit/train/005094.jpg cityscapes/gtFine/train/005094.png
cityscapes/leftImg8bit/train/001748.jpg cityscapes/gtFine/train/001748.png
cityscapes/leftImg8bit/train/005328.jpg cityscapes/gtFine/train/005328.png
cityscapes/leftImg8bit/train/000750.jpg cityscapes/gtFine/train/000750.png
cityscapes/leftImg8bit/train/002818.jpg cityscapes/gtFine/train/002818.png代码修改
1. 修改 lib/datasets/cityscapes.py
62行左右，原代码：
self.class_weights = torch.FloatTensor([0.8373, 0.918, 0.866, 1.0345,
1.0166, 0.9969, 0.9754, 1.0489,
0.8786, 1.0023, 0.9539, 0.9843,
1.1116, 0.9037, 1.0865, 1.0955,
1.0865, 1.1529, 1.0507]).cuda()改为：
self.class_weights = None
或者
self.class_weights = torch.FloatTensor([0.6, 0.9, 0.9]).cuda() # 数据个位等于类别数，数值为计算loss时权重，自己定117行左右，原代码：
label = cv2.imread(os.path.join(self.root, item["label"]),
cv2.IMREAD_GRAYSCALE)
label = self.convert_label(label)改为：
label = np.array(
Image.open(os.path.join(self.root, item['label'])).convert('P')
)在改完后的代码下方，再添加以下代码：
if 'val' in self.list_path:
image, label = self.resize_short_length(
image,
label=label,
short_length=self.base_size,
fit_stride=8
)
image, label = self.rand_crop(image, label)
image = self.input_transform(image)
image = image.transpose((2, 0, 1))
return image.copy(), label.copy(), np.array(size), name
image, label = self.resize_short_length(image, label, short_length=self.base_size)2. 修改 lib/models/bn_helper.py
第10行，原代码：
BatchNorm2d_class = BatchNorm2d = torch.nn.SyncBatchNorm改为：
BatchNorm2d_class = BatchNorm2d = torch.nn.BatchNorm2d3. 修改 experiments/cityscapes/seg_hrnet_ocr_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml
GPUS: (0,1,2,3) 改为 GPUS: (0,)
NUM_CLASSES: 19 处改为自己数据集类别数
PRETRAINED: "pretrained_models/hrnetv2_w48_imagenet_pretrained.pth" 改为 PRETRAINED: ""
IMAGE_SIZE，BASE_SIZE 和 BATCH_SIZE_PER_GPU 训练时自己根据显存使用情况修改
启动训练
1. 启动
python tools/train.py --cfg experiments/cityscapes/seg_hrnet_ocr_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml2. 出现如下日志，则训练正常
2023-03-17 22:53:55,854 Epoch: [0/484] Iter:[0/1235], Time: 3.59, lr: [0.01], Loss: 1.664520
2023-03-17 22:54:02,300 Epoch: [0/484] Iter:[10/1235], Time: 0.91, lr: [0.0099998494], Loss: 1.615116
2023-03-17 22:54:08,746 Epoch: [0/484] Iter:[20/1235], Time: 0.78, lr: [0.0099996988], Loss: 1.671794
2023-03-17 22:54:15,251 Epoch: [0/484] Iter:[30/1235], Time: 0.74, lr: [0.0099995482], Loss: 1.695571
2023-03-17 22:54:21,768 Epoch: [0/484] Iter:[40/1235], Time: 0.72, lr: [0.0099993977], Loss: 1.677186
2023-03-17 22:54:28,254 Epoch: [0/484] Iter:[50/1235], Time: 0.71, lr: [0.0099992471], Loss: 1.654732
2023-03-17 22:54:34,774 Epoch: [0/484] Iter:[60/1235], Time: 0.70, lr: [0.0099990965], Loss: 1.631594
2023-03-17 22:54:41,357 Epoch: [0/484] Iter:[70/1235], Time: 0.69, lr: [0.0099989464], Loss: 1.593724
2023-03-17 22:54:47,891 Epoch: [0/484] Iter:[80/1235], Time: 0.69, lr: [0.0099987954], Loss: 1.5581063. 可以通过 tensorboard 查看 loss 及 mIOU 变化情况
tensorboard --logdir log/4. 模型保存
训练完成后，到output目录的子目录下寻找 best.pth
TensorRT转换
需要宿主机安装有docker
环境构建
1. 拉取 nvidia 官方 docker 镜像
$ docker pull nvcr.io/nvidia/tensorrt:21.03-py32. 在该镜像中安装OpenCV
启动容器，将 /home 目录挂载到容器的 /workspace 目录（目录根据自己喜好选择，非强制）：
$ docker run -it -v /home:/workspace nvcr.io/nvidia/tensorrt:21.03-py3 bash下载OpenCV-4.5.0源码，下载时注意选择 4.5.0 版本（其他有效的亦可）：
1）OpenCV 源码链接：https://github.com/opencv/opencv
2）下载后解压，放入宿主机 /home 下，即容器的 /workspace 目录下
开始在容器中操作，安装依赖：
apt install build-essential
apt install libgtk2.0-dev pkg-config libavcodec-dev libavformat-dev libswscale-dev
apt install libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff-dev libdc1394-22-dev
apt install libgl1 # 后续安装opencv-python所需开始安装 OpenCV，依次执行以下指令：
cd /workspace/opencv-4.5.0
mkdir build
cd build
cmake -D CMAKE_BUILD_TYPE=Release -D CMAKE_INSTALL_PREFIX=/usr/local ..
make -j$nproc
make install（可选）安装pytorch和opencv-python
pip install torch==1.7.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install torchvision==0.8.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install torchaudio==0.7.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install opencv-python==4.5.3.56 -i https://pypi.tuna.tsinghua.edu.cn/simple模型转换
以下步骤主要参考 tensorrtx 项目hrnet目录下的README
1. 下载 tensorrtx 项目
项目链接：https://github.com/wang-xinyu/tensorrtx，下载zip包并解压
进入到项目的hrnet/hrnet-semantic-segmentation目录下
2. 生成 wts 模型文件
将gen_wts.py文件拷贝到模型训练项目的 tools/目录中
在模型训练项目目录下运行：
python tools/gen_wts.py --cfg experiments/cityscapes/seg_hrnet_ocr_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml --ckpt_path output/best.pth --save_path hrnet_ocr_w48.wts运行之后得到 hrnet_ocr_w48.wts文件
将 hrnet_ocr_w48.wts拷贝到 tensorrtx 项目的hrnet/hrnet-semantic-segmentation目录下
3. 生成 TensorRT 模型文件
修改hrnet/hrnet-semantic-segmentation目录下的 hrnet_ocr.cpp文件
将第 17、18 行的图像输入的 INPUT_H、INPUT_W 改为和训练时一致；
将第 19 行的 NUM_CLASSES 改为自己数据集的类别数
复制 tensorrtx 项目hrnet目录下的 hrnet-semantic-segmentation 目录，拷贝到宿主机/home目录下，也即上述环境构建时的容器的 /workspace 目录下
开始在容器中操作：
编译c++代码
cd /workspace/hrnet-semantic-segmentation
mkdir build
cd build
cmake ..
make生成 TensorRT 的 engine 模型文件
./hrnet_ocr -s ../hrnet_ocr_w48.wts ./hrnet_ocr_w48.engine 48运行之后得到 hrnet_ocr_w48.engine 文件
测试 TensorRT 模型
mkdir ../samples # 向该目录中放入一些测试集图片
./hrnet_ocr -d ./hrnet_ocr_w48.engine ../samples # 使用 hrnet_ocr_w48.engine 推理Triton部署
需要宿主机安装有docker
Server端
1. 构建目录
目录结构构建如下：
/home # 根目录也可以是其他的
└── models
└── hrnet_ocr # 模型名称，命名随意
├── 1 # 模型版本
│ └── model.plan # 即TensorRT编译运行后生成的hrnet_ocr_w48.engine，重命名为model.plan
└── config.pbtxt # 自己创建的文件，内容如下config.pbtxt文件内容
platform: "tensorrt_plan"
max_batch_size: 1
input [
{
name: "data"
data_type: TYPE_FP32
dims: [512, 512, 3]
}
]
output [
{
name: "output"
data_type: TYPE_INT32
dims: [1, 512, 512]
}
]
# default_model_filename: "model.plan"
instance_group [
{
count: 1
kind: KIND_GPU
gpus: [ 0 ]
}
]
model_warmup [
{
name: "zero_input"
batch_size: 1
inputs: {
key: "data"
value: {
data_type: TYPE_FP32
dims: [512, 512, 3]
zero_data: true
}
}
}
]2. 拉取镜像
docker pull nvcr.io/nvidia/tritonserver:21.03-py3 # 和上面TensorRT编译镜像环境保持一致3. 启动服务
docker run -d --gpus all --shm-size=16G -p 8000:8000 -p 8001:8001 -p 8002:8002 -v /home/models:/models nvcr.io/nvidia/tritonserver:21.03-py3 tritonserver --model-repository=/models --grpc-infer-allocation-pool-size=16 --log-verbose 1可通过 docker logs <container id> 查看日志，显示以下内容，启动成功
+--------------------+---------+--------+
| Model | Version | Status |
+--------------------+---------+--------+
| hrnet_ocr | 1 | READY |
+--------------------+---------+--------+ Client端
1. 构建环境
创建虚拟环境
conda create -n triton-client python=3.7
conda activate triton-client安装依赖
pip install tritonclient[all]
pip install numpy
pip install opencv-python2. 客户端代码
使用python编写客户端程序client.py，内容如下：
import argparse
import numpy as np
import cv2
import tritonclient.grpc as grpcclient
input_h, input_w = 512, 512
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--model', type=str, default='hrnet_ocr')
parser.add_argument('-v', '--video', type=str, default=r'./videos/demo.mp4')
FLAGS = parser.parse_args()
# Create server context
ip = '127.0.0.1'
port = 8001
url = "%s:%d" % (ip, port)
triton_client = grpcclient.InferenceServerClient(
url=url,
verbose=False,
ssl=False,
root_certificates=None,
private_key=None,
certificate_chain=None)
def preprocess(image):
image = cv2.resize(image, (input_h, input_w))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = image.astype(np.float32)
return image
def postprocess(result, height, width):
color_dict = {
0: (255, 0, 0),
1: (0, 255, 0),
2: (255, 0, 255)
} # 根据自己数据集自行修改
result = result.reshape((input_h, input_w)).astype(np.uint8)
result = cv2.resize(result, (width, height), interpolation=cv2.INTER_NEAREST)
canvas = np.zeros((height, width, 3), dtype=np.uint8)
for cls in color_dict:
canvas[result == cls] = color_dict[cls]
return canvas
def client(image_input):
height, width = image_input.shape[:2]
inputs = []
outputs = []
inputs.append(grpcclient.InferInput('data', [1, input_h, input_w, 3], "FP32"))
outputs.append(grpcclient.InferRequestedOutput('output'))
input_image_buffer = preprocess(image_input).astype(np.float32)
input_image_buffer = np.expand_dims(input_image_buffer, axis=0)
inputs[0].set_data_from_numpy(input_image_buffer)
results = triton_client.infer(model_name=FLAGS.model,
inputs=inputs,
outputs=outputs)
result = results.as_numpy('output')
seg_image = postprocess(result, height, width)
return seg_image
if __name__ == '__main__':
cap = cv2.VideoCapture(FLAGS.video)
while True:
ret, frame = cap.read()
if ret:
seg = client(frame)
cv2.imshow('seg result', seg)
cv2.waitKey(1)
else:
break
cap.release()3. 启动客户端
启动命令格式：
python client.py -m <模型名称，与server端一致> -v <要推理的视频路径>示例：文章来源：https://www.toymoban.com/news/detail-723446.html
python client.py -m hrnet_ocr -v videos/demo.mp4至此，HRNet从训练到加速，再到部署整个流程完成。文章来源地址https://www.toymoban.com/news/detail-723446.html
到了这里，关于HRNet语义分割训练及TensorRT部署的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
