【已解决】YOLO11模型转wts时报错:PytorchStreamReader failed reading zip archive

问题:在把训练好的新YOLO11s模型转wts文件时报错,具体信息如下图(PytorchStreamReader failed reading zip archive: failed finding central directory)

解决:新老版本pytorch之间的兼容问题,改动一下生成wts文件即可。代码帖在下面。

python 复制代码
import sys  # noqa: F401
import argparse
import os
import struct
import torch


def parse_args():
    parser = argparse.ArgumentParser(description='Convert .pt file to .wts')
    parser.add_argument('-w', '--weights', required=True,
                        help='Input weights (.pt) file path (required)')
    parser.add_argument(
        '-o', '--output', help='Output (.wts) file path (optional)')
    parser.add_argument(
        '-t', '--type', type=str, default='detect', choices=['detect', 'cls', 'seg', 'pose', 'obb'],
        help='determines the model is detection/classification')
    args = parser.parse_args()
    if not os.path.isfile(args.weights):
        raise SystemExit('Invalid input file')
    if not args.output:
        args.output = os.path.splitext(args.weights)[0] + '.wts'
    elif os.path.isdir(args.output):
        args.output = os.path.join(
            args.output,
            os.path.splitext(os.path.basename(args.weights))[0] + '.wts')
    return args.weights, args.output, args.type


pt_file, wts_file, m_type = parse_args()

print(f'Generating .wts for {m_type} model')

# Load model
print(f'Loading {pt_file}')

# Initialize
device = 'cpu'

# Load model
model = torch.load(pt_file, map_location=device, weights_only=False)  # Load FP32 weights
model = model['ema' if model.get('ema') else 'model'].float()

if m_type in ['detect', 'seg', 'pose', 'obb']:
    anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]

    delattr(model.model[-1], 'anchors')

model.to(device).eval()

with open(wts_file, 'w') as f:
    f.write('{}\n'.format(len(model.state_dict().keys())))
    for k, v in model.state_dict().items():
        vr = v.reshape(-1).cpu().numpy()
        f.write('{} {} '.format(k, len(vr)))
        for vv in vr:
            f.write(' ')
            f.write(struct.pack('>f', float(vv)).hex())
        f.write('\n')
相关推荐
小龙报2 分钟前
【数据结构与算法】单链表核心精讲:从概念到实战,吃透指针与动态内存操作
c语言·开发语言·数据结构·c++·人工智能·算法·链表
测试工程师成长之路2 分钟前
AI视觉模型如何重塑UI自动化测试:告别DOM依赖的新时代
人工智能·ui
Code Slacker4 分钟前
第八届传智杯AI虚实共振实拍创作大赛练习题库
人工智能
慧都小项4 分钟前
金融文档的“自主可控”:Python下实现Word到ODT的转换
python·金融·word
格林威5 分钟前
Baumer相机碳纤维布纹方向识别:用于复合材料铺层校验的 5 个核心技巧,附 OpenCV+Halcon 实战代码!
人工智能·数码相机·opencv·算法·计算机视觉·视觉检测
人工智能培训5 分钟前
如何将模拟器中的技能有效迁移到物理世界?
人工智能·大模型·知识图谱·具身智能·人工智能 培训·企业人工智能培训
AI有元力6 分钟前
解锁AI营销新密码,GEO优化助力品牌连接精准AI买家
人工智能
拓云者也7 分钟前
常用的生物信息学数据库以及处理工具
数据库·python·oracle·r语言·bash
GISer_Jing8 分钟前
原生HTML项目重构:Vue/React双框架实战
vue.js·人工智能·arcgis·重构·html
ZCXZ12385296a9 分钟前
YOLOv11-C3k2-wConv改进脐橙目标检测与分级模型研究
人工智能·yolo·目标检测