参考资料 ultralytics/nn/tasks.py
备注
该文件可从https://github.com/ultralytics/ultralytics/blob/main/ ultralytics/nn/tasks .py 获取。如果您发现问题,请通过提交 Pull Request🛠️ 帮助修复。谢谢🙏!
ultralytics.nn.tasks.BaseModel
垒球 Module
BaseModel 类是Ultralytics YOLO 系列中所有模型的基类。
源代码 ultralytics/nn/tasks.py
74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 |
|
forward(x, *args, **kwargs)
单比例模型的前传。封装 _forward_once
方法。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
x |
Tensor | dict
|
输入图像tensor 或包含图像tensor 和 gt 标签的 dict。 |
所需 |
返回:
类型 | 说明 |
---|---|
Tensor
|
网络的输出。 |
源代码 ultralytics/nn/tasks.py
fuse(verbose=True)
保险丝 Conv2d()
和 BatchNorm2d()
层,以提高计算效率。
计算效率。
返回:
类型 | 说明 |
---|---|
Module
|
将返回融合后的模型。 |
源代码 ultralytics/nn/tasks.py
info(detailed=False, verbose=True, imgsz=640)
打印型号信息。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
detailed |
bool
|
如果为 True,则打印模型的详细信息。默认为假 |
False
|
verbose |
bool
|
如果为 True,则打印模型信息。默认为假 |
True
|
imgsz |
int
|
模型要训练的图像大小。默认为 640 |
640
|
源代码 ultralytics/nn/tasks.py
init_criterion()
is_fused(thresh=10)
检查模型的 BatchNorm 层数是否少于某个阈值。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
thresh |
int
|
BatchNorm 图层的阈值数。默认为 10。 |
10
|
返回:
类型 | 说明 |
---|---|
bool
|
如果模型中的 BatchNorm 层数少于阈值,则为 True,否则为 False。 |
源代码 ultralytics/nn/tasks.py
load(weights, verbose=True)
将权重加载到模型中。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
weights |
dict | Module
|
要加载的预训练权重。 |
所需 |
verbose |
bool
|
是否记录传输进度。默认为 True。 |
True
|
源代码 ultralytics/nn/tasks.py
loss(batch, preds=None)
计算损失。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
batch |
dict
|
计算损失的批次 |
所需 |
preds |
Tensor | List[Tensor]
|
预测 |
None
|
源代码 ultralytics/nn/tasks.py
predict(x, profile=False, visualize=False, augment=False, embed=None)
对网络进行前向传递。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
x |
Tensor
|
模型的输入tensor 。 |
所需 |
profile |
bool
|
如果为 True,则打印每一层的计算时间,默认为 False。 |
False
|
visualize |
bool
|
保存模型的特征地图(如果为 True,默认为 False)。 |
False
|
augment |
bool
|
预测时增强图像,默认为 "假"。 |
False
|
embed |
list
|
要返回的特征向量/嵌入的列表。 |
None
|
返回:
类型 | 说明 |
---|---|
Tensor
|
模型的最后输出。 |
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.DetectionModel
垒球 BaseModel
YOLOv8 检测模型。
源代码 ultralytics/nn/tasks.py
__init__(cfg='yolov8n.yaml', ch=3, nc=None, verbose=True)
使用给定的配置和参数初始化YOLOv8 检测模型。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.OBBModel
YOLOv8 定向边框(OBB)模型。
源代码 ultralytics/nn/tasks.py
__init__(cfg='yolov8n-obb.yaml', ch=3, nc=None, verbose=True)
ultralytics.nn.tasks.SegmentationModel
YOLOv8 分割模型。
源代码 ultralytics/nn/tasks.py
__init__(cfg='yolov8n-seg.yaml', ch=3, nc=None, verbose=True)
ultralytics.nn.tasks.PoseModel
YOLOv8 姿势模型。
源代码 ultralytics/nn/tasks.py
__init__(cfg='yolov8n-pose.yaml', ch=3, nc=None, data_kpt_shape=(None, None), verbose=True)
初始化YOLOv8 Pose 模型。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.ClassificationModel
垒球 BaseModel
YOLOv8 分类模式。
源代码 ultralytics/nn/tasks.py
__init__(cfg='yolov8n-cls.yaml', ch=3, nc=None, verbose=True)
使用 YAML、通道、类的数量、verbose 标志初始化 ClassificationModel。
init_criterion()
reshape_outputs(model, nc)
staticmethod
如果需要,将 TorchVision 分类模型更新为类计数'n'。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.RTDETRDetectionModel
RTDETR(使用变压器进行实时检测和跟踪)检测模型类。
该类负责构建 RTDETR 架构、定义损失函数,并促进训练和推理过程。 训练和推理过程。RTDETR 是一个对象检测和跟踪模型,从 DetectionModel 基类扩展而来。
属性
名称 | 类型 | 说明 |
---|---|---|
cfg |
str
|
配置文件路径或预设字符串。默认为 "rtdetr-l.yml"。 |
ch |
int
|
输入通道数。默认为 3(RGB)。 |
nc |
int
|
对象检测的类数。默认为 "无"。 |
verbose |
bool
|
指定初始化过程中是否显示汇总统计数据。默认为 True。 |
方法
名称 | 说明 |
---|---|
init_criterion |
初始化用于计算损失的标准。 |
loss |
计算并返回训练过程中的损失。 |
predict |
执行网络前向传递并返回输出。 |
源代码 ultralytics/nn/tasks.py
441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 |
|
__init__(cfg='rtdetr-l.yaml', ch=3, nc=None, verbose=True)
初始化 RTDETRDetectionModel。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
cfg |
str
|
配置文件名称或路径。 |
'rtdetr-l.yaml'
|
ch |
int
|
输入通道数 |
3
|
nc |
int
|
班级数量。默认为 "无"。 |
None
|
verbose |
bool
|
在初始化过程中打印附加信息。默认为 True。 |
True
|
源代码 ultralytics/nn/tasks.py
init_criterion()
loss(batch, preds=None)
计算给定批次数据的损失。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
batch |
dict
|
包含图像和标签数据的字典。 |
所需 |
preds |
Tensor
|
预先计算的模型预测值。默认为 "无"。 |
None
|
返回:
类型 | 说明 |
---|---|
tuple
|
一个元组,包含tensor 中的总损失和主要三项损失。 |
源代码 ultralytics/nn/tasks.py
predict(x, profile=False, visualize=False, batch=None, augment=False, embed=None)
对模型进行前向传递。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
x |
Tensor
|
输入tensor 。 |
所需 |
profile |
bool
|
如果为 "true",则对每一层的计算时间进行剖析。默认为 "假"。 |
False
|
visualize |
bool
|
如果为 "true",则为可视化保存特征地图。默认为 "假"。 |
False
|
batch |
dict
|
用于评估的基本真实数据。默认为 "无"。 |
None
|
augment |
bool
|
如果为 "true",则在推理过程中执行数据增强。默认为 "假"。 |
False
|
embed |
list
|
要返回的特征向量/嵌入的列表。 |
None
|
返回:
类型 | 说明 |
---|---|
Tensor
|
模型的输出tensor. |
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.WorldModel
YOLOv8 世界模式。
源代码 ultralytics/nn/tasks.py
558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 |
|
__init__(cfg='yolov8s-world.yaml', ch=3, nc=None, verbose=True)
使用给定的配置和参数初始化YOLOv8 世界模型。
源代码 ultralytics/nn/tasks.py
loss(batch, preds=None)
计算损失。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
batch |
dict
|
计算损失的批次。 |
所需 |
preds |
Tensor | List[Tensor]
|
预测 |
None
|
源代码 ultralytics/nn/tasks.py
predict(x, profile=False, visualize=False, txt_feats=None, augment=False, embed=None)
对模型进行前向传递。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
x |
Tensor
|
输入tensor 。 |
所需 |
profile |
bool
|
如果为 "true",则对每一层的计算时间进行剖析。默认为 "假"。 |
False
|
visualize |
bool
|
如果为 "true",则为可视化保存特征地图。默认为 "假"。 |
False
|
txt_feats |
Tensor
|
文本特征,如果给定,则使用它。默认为 "无"。 |
None
|
augment |
bool
|
如果为 "true",则在推理过程中执行数据增强。默认为 "假"。 |
False
|
embed |
list
|
要返回的特征向量/嵌入的列表。 |
None
|
返回:
类型 | 说明 |
---|---|
Tensor
|
模型的输出tensor. |
源代码 ultralytics/nn/tasks.py
set_classes(text, batch=80, cache_clip_model=True)
提前设置类别,以便模型可以在不剪辑模型的情况下进行离线推理。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.Ensemble
垒球 ModuleList
模型组合。
源代码 ultralytics/nn/tasks.py
__init__()
forward(x, augment=False, profile=False, visualize=False)
函数生成YOLO 网络的最后一层。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.temporary_modules(modules=None)
上下文管理器,用于临时添加或修改Python 模块缓存中的模块 (sys.modules
).
该函数可用于在运行时更改模块路径。它在重构代码时非常有用、 将模块从一个位置移到另一个位置,但为了向后兼容,仍希望支持旧的导入 路径以实现向后兼容。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
modules |
dict
|
将旧模块路径映射到新模块路径的字典。 |
None
|
示例
备注
更改仅在上下文管理器中生效,一旦上下文管理器退出,更改将被撤销。
请注意,直接操作 sys.modules
会导致不可预测的结果,尤其是在大型应用程序或库中。
尤其是在较大的应用程序或程序库中。请谨慎使用该函数。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.torch_safe_load(weight)
此函数尝试使用torch.load() 函数加载PyTorch 模型。如果出现 ModuleNotFoundError、 会捕获该错误,记录一条警告信息,并尝试通过 check_requirements() 函数安装缺失的模块。安装完成后,函数会再次尝试使用torch.load() 加载模型。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
weight |
str
|
PyTorch 模型的文件路径。 |
所需 |
返回:
类型 | 说明 |
---|---|
dict
|
装载的PyTorch 模型。 |
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.attempt_load_weights(weights, device=None, inplace=True, fuse=False)
加载一组模型 weights=[a,b,c] 或一个模型 weights=[a] 或 weights=a。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.attempt_load_one_weight(weight, device=None, inplace=True, fuse=False)
加载单个模型权重。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.parse_model(d, ch, verbose=True)
将YOLO model.yaml 字典解析为PyTorch 模型。
源代码 ultralytics/nn/tasks.py
830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 |
|
ultralytics.nn.tasks.yaml_model_load(path)
从 YAML 文件加载YOLOv8 模型。
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.guess_model_scale(model_path)
将YOLO 模型 YAML 文件的路径作为输入,并提取模型比例尺的大小字符。函数 使用正则表达式匹配来查找 YAML 文件名中的模型比例模式,该模式用 函数以字符串形式返回模型标尺的大小字符。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
model_path |
str | Path
|
YOLO 模型 YAML 文件的路径。 |
所需 |
返回:
类型 | 说明 |
---|---|
str
|
模型比例尺的大小字符,可以是 n、s、m、l 或 x。 |
源代码 ultralytics/nn/tasks.py
ultralytics.nn.tasks.guess_model_task(model)
从PyTorch 模型的架构或配置猜测其任务。
参数
名称 | 类型 | 说明 | 默认值 |
---|---|---|---|
model |
Module | dict
|
PyTorch 模型或模型配置的 YAML 格式。 |
所需 |
返回:
类型 | 说明 |
---|---|
str
|
模型的任务("检测"、"分割"、"分类"、"姿势")。 |
加薪:
类型 | 说明 |
---|---|
SyntaxError
|
如果无法确定模型的任务。 |
源代码 ultralytics/nn/tasks.py
982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 |
|