【YoLo】从 PyTorch Hub 加载 YOLOv5并对网络进行修改

一、从pytorc hub 上下载YOLOv5网络结构

http://t.csdn.cn/Wf8IX这篇博客很详细里介绍pytorch如何下载YOLOv5s网络结构以及YOLOv5s.pt权重（在运行程序时如果没有，则会自动下载）。

在运行代码时，可能会出现报错，遇到有两个模块没有的情况：

1.No module named 'psutil'

2.ModuleNotFoundError: No module named 'ultralytics'

 

遇到的问题可能会因个人的环境不一样，会出现不同的问题。当时环境我用的是mmdetection的环境。

下面是官网的配置环境

 

运行代码下载一些东西时，会出现一些错误，可能是由于没有用梯子的缘故！！！

import torch
import torch.nn as nn
from torchvision import models
# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Image
im = 'https://ultralytics.com/images/zidane.jpg'

# Inference
results = model(im)
results.show()
print(results.pandas().xyxy[0])
print(model)

代码成功运行之后，展示效果

 

如果下载了YOLOv5s.pt权重文件的话，是可以直接处理图片，并展示出效果的。

results = model(im)
results.show()

 

print(results.pandas().xyxy[0])
打印出YOLO返回的信息

 

print(model)   展示下载的YOLOv5s网络结构，下载的是最新v7.0版本的YOLOv5s，是自动下载的最新版本的权重文件。

 

二、修改网络结构

修改下载下来的YOLOv5网络结构，可以参考这篇博客：http://t.csdn.cn/Gjz7X

print(model.model)   与上面的model模型对比可知，第一个model是导入模型的变量，后续的三个model分别表示

(model): DetectMultiBackend
(model): DetectionModel
(model): Sequential。

例如，

print(model.model)

print(model.model.model.model[0])

print(model.model.model.model[0].conv)

 个人觉得，利用pytorch下载YOLO网络结构并修改网络结构是有一定好处：可以根据自己修改网络结构并且训练好权重文件、权重文件、权重文件！！！

在将权重文件放到指定位置，加载到模型中，就可以实现调用。这时网络模型就可以作为一个模块，即插即用。前提是需要训练权重文件。