推理时增大batch_size

opendatalab / DocLayout-YOLO

DocLayout-YOLO: Enhancing Document Layout Analysis through Diverse Synthetic Data and Global-to-Local Adaptive Perception

https://huggingface.co/spaces/opendatalab/DocLayout-YOLO

GNU Affero General Public License v3.0

537 stars 38 forks source link

推理时增大batch_size #14

Closed luciaganlulu closed 1 month ago

luciaganlulu commented 1 month ago

您好，请问是直接改这里的batch数，另外source从单个image_path改成所有的图片的PIL.Image列表嘛？https://github.com/opendatalab/DocLayout-YOLO/blob/main/doclayout_yolo/engine/model.py#L431C8-L431C98

JulioZhao97 commented 1 month ago

您好，是的

推理batch_size大小可以直接在 https://github.com/opendatalab/DocLayout-YOLO/blob/main/doclayout_yolo/engine/model.py#L431C8-L431C98 修改batch_size的值
demo.py中的input输入为列表，列表中元素可以为PIL.Image, np.array，image_path都可以

luciaganlulu commented 1 month ago

那请问这种方式的话，怎么保存每个输入图片对应的annotated_frame到指定的路径以及从det_res中获取每个输入图片对应的检测结果参数值呢（例如bbox, label, score）？predict这个函数返回的list没太看明白

luciaganlulu commented 1 month ago

您好我把det_res打印出来看到了，路径问题知道啦，但是请问如何获取图中每个检测对象的bbox和score呢？另外det_res里面的save_dir=‘runs/detect/predict’是什么意思呢？没找到这个路径在哪里诶

JulioZhao97 commented 1 month ago

您好，predict返回的list长度和输入的list相同，其中每一个元素包含了预测的结果，可以参考以下保存每一个图片的结果

image_list = [
        "assets/example/academic.jpg",
        "assets/example/exam_paper.jpg",
        "assets/example/financial.jpg",
        "assets/example/fuzzy_scan.jpg",
    ]

    det_res = model.predict(
        image_list,
        imgsz=args.imgsz,
        conf=args.conf,
        device=device,
    )
    for image_idx, image_path in enumerate(image_list):
        print(det_res[image_idx].boxes.xyxy)
        print(det_res[image_idx].boxes.cls)
        print(det_res[image_idx].boxes.conf)
        annotated_frame = det_res[image_idx].plot(pil=True, line_width=args.line_width, font_size=args.font_size)
        if not os.path.exists(args.res_path):
            os.makedirs(args.res_path)
        output_path = os.path.join(args.res_path, image_path.split("/")[-1].replace(".jpg", "_res.jpg"))
        cv2.imwrite(output_path, annotated_frame)
        print(f"Result saved to {output_path}")

JulioZhao97 commented 1 month ago

结果参数值在det_res[idx].bbox中，分类结果为det_res[idx].bbox.cls，置信度为det_res[idx].bbox.conf，检测框为det_res[idx].bbox.xyxy（或者不同格式）
save_dir只有传入save=True才会保存到runs/detect/predict这个默认路径，可以参考https://docs.ultralytics.com/modes/predict/#inference-arguments

luciaganlulu commented 1 month ago

好的谢谢，非常感谢，代码这里写得很清楚～ https://github.com/opendatalab/DocLayout-YOLO/blob/main/doclayout_yolo/engine/results.py#L444

luciaganlulu commented 1 month ago

好的非常感谢～

luciaganlulu commented 1 month ago

您好，predict返回的list长度和输入的list相同，其中每一个元素包含了预测的结果，可以参考以下保存每一个图片的结果

image_list = [
        "assets/example/academic.jpg",
        "assets/example/exam_paper.jpg",
        "assets/example/financial.jpg",
        "assets/example/fuzzy_scan.jpg",
    ]

    det_res = model.predict(
        image_list,
        imgsz=args.imgsz,
        conf=args.conf,
        device=device,
    )
    for image_idx, image_path in enumerate(image_list):
        print(det_res[image_idx].boxes.xyxy)
        print(det_res[image_idx].boxes.cls)
        print(det_res[image_idx].boxes.conf)
        annotated_frame = det_res[image_idx].plot(pil=True, line_width=args.line_width, font_size=args.font_size)
        if not os.path.exists(args.res_path):
            os.makedirs(args.res_path)
        output_path = os.path.join(args.res_path, image_path.split("/")[-1].replace(".jpg", "_res.jpg"))
        cv2.imwrite(output_path, annotated_frame)
        print(f"Result saved to {output_path}")

请问下，这种方式与os.walk逐个image推理的方式相比的话，区别就在于这种list的方式可以结合batch_size数进行批处理对吗？后者只能batch_size为1。

JulioZhao97 commented 1 month ago

您好，predict返回的list长度和输入的list相同，其中每一个元素包含了预测的结果，可以参考以下保存每一个图片的结果

image_list = [
        "assets/example/academic.jpg",
        "assets/example/exam_paper.jpg",
        "assets/example/financial.jpg",
        "assets/example/fuzzy_scan.jpg",
    ]

    det_res = model.predict(
        image_list,
        imgsz=args.imgsz,
        conf=args.conf,
        device=device,
    )
    for image_idx, image_path in enumerate(image_list):
        print(det_res[image_idx].boxes.xyxy)
        print(det_res[image_idx].boxes.cls)
        print(det_res[image_idx].boxes.conf)
        annotated_frame = det_res[image_idx].plot(pil=True, line_width=args.line_width, font_size=args.font_size)
        if not os.path.exists(args.res_path):
            os.makedirs(args.res_path)
        output_path = os.path.join(args.res_path, image_path.split("/")[-1].replace(".jpg", "_res.jpg"))
        cv2.imwrite(output_path, annotated_frame)
        print(f"Result saved to {output_path}")

请问下，这种方式与os.walk逐个image推理的方式相比的话，区别就在于这种list的方式可以结合batch_size数进行批处理对吗？后者只能batch_size为1。

根据YOLO的逻辑来说应该是的，您可以测试一下

luciaganlulu commented 1 month ago

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
det_res = model.predict(
    image_paths,
    imgsz=1024,
    conf=0.2,
    device="cuda:1"
)

您好，我用的image_list的方式，batchsize设置的8，上面这个代码的话，我想用GPU #1（显存多点），但请问为什么还是用GPU #0呢？

JulioZhao97 commented 1 month ago

您好，我这里本地测试没有问题，麻烦您试一下去掉os.environ改用CUDA_VISIBLE_DEVICES=xx python运行，将predict中device改成cuda试一试呢？

luciaganlulu commented 1 month ago

您好，按这样还是报一样的错误，而且我每0.1秒更新显存变化都看不到显存变化，然后过了好久显示这个out of memory的报错，这是什么原因呢？跟torch.no-grad有关么？

luciaganlulu commented 1 month ago

我又测了一下repo里readme的demo.py，这个是正常的可以GPU1，不知道为什么改成用image_list形式就设置不了第二块GPU了，您能看看我代码哪里错了么？运行是用的 CUDA_VISIBLE_DEVICES=1 python this.py

import os
import cv2
from pathlib import Path
from doclayout_yolo import YOLOv10
import glob
import json

FOLDER_NAME = "myfoldername"

def get_save_path(image_path):
    subfolderpath = os.path.dirname(image_path).lstrip('myprefix')
    folderpath = os.path.join('./results/annotated_imgs', subfolderpath)
    os.makedirs(folderpath, exist_ok=True)
    save_path = os.path.join(folderpath, image_path.split("/")[-1].replace(".png", "_res.png"))
    return save_path

model_path = "/doclayout_yolo_docstructbench_imgsz1024.pt"
model = YOLOv10(model_path)  # load an official model

image_paths = glob.glob('myprefix/{}/**/*.png'.format(FOLDER_NAME), recursive=True)
annotatedimage_paths = [get_save_path(e) for e in image_paths]

det_res = model.predict(
    image_paths,
    imgsz=1024,
    conf=0.2,
    device="cuda"
)

output_data = {}

for image_idx, image_path in enumerate(image_paths):
    bbox = det_res[image_idx].boxes.xyxy.tolist()
    cls = det_res[image_idx].boxes.cls.tolist()
    score = det_res[image_idx].boxes.conf.tolist()

    annotated_frame = det_res[image_idx].plot(pil=True, line_width=5, font_size=20)
    annotated_path = annotatedimage_paths[image_idx]
    os.makedirs(os.path.dirname(annotated_path), exist_ok=True)
    cv2.imwrite(annotated_path, annotated_frame)

    output_data[image_path] = {
        "bbox": bbox,
        "cls": cls,
        "score": score,
        "annotated_image_path": annotated_path
    }
with open(os.path.join('./results/info/{}.json'.format(FOLDER_NAME)), "w") as f:
    json.dump(output_data, f, indent=4)

JulioZhao97 commented 1 month ago

经过测试，加入以下两句可以指定GPU，如果不加确实都是在GPU0，具体原因暂时不明确，可能是高版本torch的影响

import torch
print(torch.cuda.is_available())

麻烦您安装低版本torch例如torch==2.0.0试一试呢？例如

pip3 install torch==2.0.0+cu117 torchvision==0.15.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html

JulioZhao97 commented 1 month ago

测试了一下貌似不行，您可以先把提到的解决方案加入，原因可能需要进一步调查

luciaganlulu commented 1 month ago

您好，加了这两句（出来是True），设置batchsize为4，过了好久我确实看到GPU1的显存明显增加，但是突然又报GPU0显存不够，跟今天一样的报错内容。我是torch2.4.1, torchvision0.19.1, cuda12.3。

另外您这个速度很快，我先用单张推理了～有个问题请教下：这个对应关系是固定的么？我看det_res出来是3.0， 5.0这种，其实就相当于是这里的3和5对吧？

JulioZhao97 commented 1 month ago

是的，det_res[image_idx].boxes.cls的值和上述对应关系是相符合的

luciaganlulu commented 1 month ago

好的非常感谢～

luciaganlulu commented 1 month ago

您好，我仿照demo.py，将device = 'cuda' if torch.cuda.is_available() else 'cpu'加入到了代码里，model.predict里面device=device，这样的话当我执行CUDA_VISIBLE_DEVICES=1 python this.py的时候确实是第二块GPU显存有变化，但最终报错还是今天发的那个out of memory。

请问下，是不是这里面的batch size不起作用呢？是不是他加载的是image_lists里面的全部图片去做推理呢？因为我发现用image_lists这种方法，我这里设置为1，也会报out of memory 的错～

JulioZhao97 commented 1 month ago

截屏2024-10-24 10 00 31 这个应该是起作用的，您可以查看推理时的log，shape的batch size会跟着变化，至于out of memory我认为可能和推理模式（predict）有关，可能存在显存泄漏，您可以尝试PDF-Extract-Kit来大量处理

luciaganlulu commented 1 month ago

嗯嗯谢谢～这个log只是print出来还是说写在哪个路径了呢？当我用单张图片推理这种形式的时候，显示的是 image 1/1 /data/0.png: 1024*376, 类别统计信息，36.8ms这种，这个的意思就是batchsize是1，当batchsize设置为多个且起作用的话，这里会变成类似于image 1/8 ......\n image 2/8......\n这种对嘛？另外我看您的示例都是1024 1024，我的有的是1024 736有的是1024 800这样，我的原图是从pdf自己存下来的，通常是3k~4k分辨率，怎么设置成出来是10241024呢？需要设置成这样么？

luciaganlulu commented 1 month ago

为什么您这个例子的batchsize是2但是出来了7个图片推理结果呢？

JulioZhao97 commented 1 month ago

这个log只是会print出来，当batch_size>1时输出的log如下截屏2024-10-24 11 06 48

至于图像尺寸，batch_size>1时会保持rect模式，例如要将两张图片padding到相同的size，才能batch推理，在batch_size=1的情况下会保持单张图像的长宽比，长边resize到imgsz，即1024，下面是我用来测试的代码，batch_size=2，您可以参考下

import os
import cv2
from pathlib import Path
from doclayout_yolo import YOLOv10
import glob
import json

FOLDER_NAME = "assets"

model_path = "../DocLayout-YOLO-git/checkpoints/doclayout_yolo_docstructbench_imgsz1024.pt"
model = YOLOv10(model_path)  # load an official model

image_paths = glob.glob('assets/example/*.jpg'.format(FOLDER_NAME), recursive=True)

# import torch
# print(torch.cuda.is_available())

while True:
    det_res = model.predict(
        image_paths,
        imgsz=1024,
        conf=0.2,
        device='cuda',
    )

关于推理时的预处理方法以及图像尺寸可以参考这里：https://github.com/ultralytics/ultralytics/issues/13494

luciaganlulu commented 1 month ago

明白了，谢谢您耐心回答，非常感谢！谢谢分享这个引用的issue，讲得蛮详细的～