HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.3.59 Add ability to load any torchvision model as module (#18564)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Mohammed Yasin 2025-01-09 20:57:46 +08:00 • committed by GitHub
parent 246c3eca81
commit cc1e77138c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 104 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/nn/tasks.py
View file

@ -50,6 +50,7 @@ from ultralytics.nn.modules import (
HGBlock,
HGStem,
ImagePoolingAttn,
Index,
Pose,
RepC3,
RepConv,
@ -59,6 +60,7 @@ from ultralytics.nn.modules import (
RTDETRDecoder,
SCDown,
Segment,
TorchVision,
WorldDetect,
v10Detect,
)
@ -1052,7 +1054,7 @@ def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)
m.legacy = legacy
elif m is RTDETRDecoder: # special case, channels arg must be passed in index 1
args.insert(1, [ch[x] for x in f])
elif m is CBLinear:
elif m in {CBLinear, TorchVision, Index}:
c2 = args[0]
c1 = ch[f]
args = [c1, c2, *args[1:]]