Run yolov5 on Xavier and solve the pytorch version compatibility problem

This article mainly describes how I r u n yolov5:v4.0 on Jetson AGX Xavier 0

My Xavier configuration:

  • JetPack: 4.3
  • cuda:10.0
  • python:3.6.9

YOLOv5

The download training reasoning here runs on the server.

Download yolov5:

git clone https://github.com/Oswells/yolov5.git

Note: This is fork ultralytics/**yolov5 **Yolo V5, version 4.0 of the great God, has launched version 5.0. To download the latest version, you still need to download it on the github of the great God.

Here I will mainly talk about several important documents.

  • models/yolo5s.yaml

    There are four similar files here, namely s, m, l and X. The model increases in turn. These four files define four network structures. If you don't care about the network structure, you only need to pay attention to nc. You can change it into several classes. I have only one class, so I set it to 1

  • data/hyp.scratch.yaml

    This file defines some super parameters. Change this file to select different super parameters before training. From hsv_h starts with super parameters related to data enhancement. When it is set to 0, it means it is not enabled.

Training yolov5

Main reference This blog , yes, the enthusiastic netizen mentioned by the author is me.

Choose python version 3.7 or above, and then install the corresponding library. The yolov5 author wrote requirements Txt can help us deploy the environment quickly.

pip3 install -r requirements.txt

The author also gave Dockerfile. If there is docker, you can choose to use Dockerfile to generate container deployment environment. I tried this method. It can be used, but it is not necessary hh.

yolov5 uses the argparse library, which seems to have become the standard configuration of neural network programs. It can flexibly transfer parameters from the terminal to the program. yolov5 in train Py and detect Py defines many parameters that can be passed in, but some parameters are generally fixed with a certain value, such as data (where the dataset is located). We can directly modify the default value in the program, so that we don't have to type the same thing in the terminal every time.

yolov5 reasoning

Generally, you can reason by running the following code.

python detect.py --weights 'weights/best.pt' --source '1' --conf-thres 0.4 --device 'cpu'

Here, weights is the trained weight. source when it is a number, it means reading an image from the camera. When it is a folder path, it will infer the pictures / videos in the whole folder. When it is a file, it will infer the file. Conf thres is the confidence threshold, which is easy to recognize other objects as targets when set low. Device selects the device on which the program runs. When it is cpu, it runs slowly. The 640x480 image has only 2FPS. When it is digital, it runs on the graphics card (gpu). When selecting multiple graphics cards, refer to this format '0,1'.

Reasoning does not need so many files. We can delete some files. Here I list the files needed for reasoning.

Folder: models, utils, data

Files: weight files, detect py,

Xavier pytoch runs yolov5

Basic environment installation

miniforge is recommended. It is a substitute for anaconda. Anaconda cannot be installed on Xavier. Installation method reference This blog . miniforge is used in exactly the same way as anaconda.

After installation, create a new environment, select 3.6 for Python version, and open the requirements in yolov5 file Txt, except opencv python, torch and torch vision, they are all installed with conda install. conda can no longer use pip3. The opencv Python library comes with xavier and does not need to be installed.

Torch and torch vision installation

Because Xavier is an arm architecture, torch and torchvision installed directly from pip or conda cannot be used. You need to install them from This website Download the corresponding versions of torch and torch vision. Because the Jetpack version of my Xavier is 4.3, I can't install a higher version of torch. Here I download and install version 1.4.0 of torch and version 0.5.0 of torch vision.

torch installation

pip3 install torch-1.4.0-cp36-cp36m-linux_aarch64.whl	#Here, my default whl file is in the directory opened by the terminal

torchvision installation

https://github.com/pytorch/vision/tags?after=v0.7.0-rc3 this interface downloads the corresponding version of torch vision, and then decompresses the cd into this folder.

sudo python3 setup.py install 

Resolve compatibility issues

Because we use a higher version of torch when training and 1.4 version of torch when actually using, we will inevitably encounter compatibility problems.

  • Change torch model save format

    First, ensure that the torch version of your computer is ≥ 1.7.

    Create a new py file on your computer. The contents are as follows, and modify it yourself_ File and load_path,load_ path.

import torch
import os
load_file = "5s_multi_scale_gaussian"
load_path = os.path.join(load_file,"weights","best.pt")
load_path = os.path.join("weights",load_file+".pt")

ckpt = torch.load(load_path)
print(ckpt.get('model'))
ckpt = {'epoch': ckpt.get('epoch'),
        'best_fitness': ckpt.get('best_fitness'),
        'training_results': ckpt.get('training_results'),
        'model': ckpt.get('model'),
        'ema': ckpt.get('ema'),
        'updates': ckpt.get('updates'),
        'optimizer': ckpt.get('optimizer'),
        'wandb_id': ckpt.get('wandb_id')}
torch.save(ckpt,save_path,_use_new_zipfile_serialization=False)

Run the program and copy the new weight file to Xavier.

Note that since Pickle is used to save the loading model in torch, the weight must be successfully loaded in a directory structure similar to the original yolov5 directory. Here, there must be models and utils folders in the same directory as this py file.

  • Change torch underlying code

    Since yolov5 calls the classes only available in the new version of torch, we need to supplement these functions and classes when using the old version of torch.

    First, find the directory of torch in Xavier and enter it in the terminal

    pip3 show torch
    

    You can find its storage Location in Location, open this directory, and then open / NN / modules / activation Py, add the following

    class SiLU(Module):
        def forward(self,x):
            return x * torch.sigmoid(x)
    
    
    class Hardswish(Module):
        def forward(self,x):
            return x * F.hardtanh(x + 3, 0., 6.) / 6.
    
  • Modify yolov5 / Models / experimental py

    Start with from torch nn. modules. activation import SiLU,Hardswish

    Locate line 117 and modify as follows

    #original
    if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
    #Change to
    if type(m) in [Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, SiLU]:
    

    I think this is because I don't know the bottom of python. I don't know how to make NN Silu and NN Hardwish can also use it. My method may still be clumsy.

OK, now you can run detect in the terminal Py.

tensorRT acceleration

Using the above method, I found that yolov5s was only about 19FPS in actual use, which could not achieve real-time effect. I heard that tensorRT could accelerate reasoning, so I tried to do it, and the effect was good. Yolov5s finally reached about 40FPS.

yolo to tensorRT

yolo's pt file can't be accelerated directly by tensorRT and needs to be converted. A great God on GitHub has developed a one-step format conversion. Here I post the conversion yolov5:: v4 Links to 0: https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v4.0/yolov5

The Config first selection model in his readme can be ignored, and other contents can be set according to his instructions.

After setting, follow How to Run, yolov5s as example to achieve sudo of the second point/ yolov5 -s yolov5s. wts yolov5s. Engine s is enough. Here, because my computer doesn't have gpu, I perform the above steps on xavier. Here, we have constructed the weight file required for operation engine.

Call camera

I mainly refer to it here This blog The fourth point of video detection. Since he doesn't use torch, I use engine directly in the base environment. Install pycuda first

pip3 install pycuda

Then create a new py file, copy the author's code, and modify the key parameters. However, I think there are some areas that can be improved in the process of use. I write the modified function below.

def draw_boxes(image_raw, result_boxes,color, result_scores, result_classid):
    for i in range(len(result_boxes)):
        box = result_boxes[i]
        plot_one_box(
            box,
            image_raw,
            color,
            label="{}:{:.2f}".format(
                categories[int(result_classid[i])], result_scores[i]
            ),
        )
    return image_raw

def detect_camera(camera, yolov5_wrapper):
    ##Start loop detection and write the result to result In MP4
    color = [random.randint(0, 255) for _ in range(3)]
    while True:
        #ret,img = camera.read()  # usb camera with this
        t1 = time.time()
        img = camera.read()
        img, result_boxes, result_scores, result_classid = yolov5_wrapper.infer(img)
        img = draw_boxes(img, result_boxes,color, result_scores, result_classid)
        cv2.namedWindow("result", 0)
        cv2.resizeWindow("result", 640, 480)
        cv2.imshow("result", img)
        print("%d FPS"%(1/(time.time()-t1)))
        if cv2.waitKey(1) == ord('q'):
            break

Pay attention to modifying main_ Path in camera() and usb camera.

# If the csiamera () method is not called, there is no need to import
def main_camera():
    camera = cv2.VideoCapture(0)  # usb camera can use this, or use the usb interface in jetcam
    # camera = CSICamera(capture_device=0, width=224, height=224)
    # load custom plugins
    PLUGIN_LIBRARY = "build/libmyplugins.so"
    ctypes.CDLL(PLUGIN_LIBRARY)
    engine_file_path = "build/yolov5s.engine"

    # YoLov5TRT instance
    yolov5_wrapper = YoLov5TRT(engine_file_path)
    print("start detection!")
    detect_camera(camera, yolov5_wrapper)
    camera.release()  # Use cv method to open the camera
    cv2.destroyAllWindows()
    print("\nfinish!")
    
main_camera()

Finally, run the py file.

Keywords: Python OpenCV Pytorch Object Detection

Added by jstgermain on Sat, 22 Jan 2022 02:22:59 +0200