如何用PyTorch加载CLRNet预训练模型？新手技术求助

1. 预处理预训练权重文件 #

下载culane_r18.pth.zip后，先解压得到culane_r18.pth，将文件放在CLRNet项目目录下（比如新建pretrained_weights文件夹统一存放）。

2. 补充代码关键部分 #

可以新建单独的测试文件（比如test_pretrained.py），也可以在项目现有脚本中修改。以下是完整可运行的代码示例：

import torch
from models.clrnet import CLRNet
from utils.config import Config

# 1. 加载预训练模型对应的配置文件（culane_r18对应项目内的configs/culane/clr_res18.yaml）
cfg = Config.fromfile('configs/culane/clr_res18.yaml')

# 2. 初始化匹配结构的模型
model = CLRNet(cfg.model)

# 3. 加载预训练权重
FILE = "pretrained_weights/culane_r18.pth"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load(FILE, map_location=device))
model.to(device)
model.eval()  # 切换到推理模式

# 4. 可选：用单张图片测试推理流程
from PIL import Image
from utils.transforms import get_transform

transform = get_transform(cfg.val_transform)
img = Image.open("test_image.jpg").convert('RGB')
img_tensor = transform(img).unsqueeze(0).to(device)

with torch.no_grad():
    outputs = model(img_tensor)
# 后续可根据项目文档解析输出的车道线结果
print(outputs)

关键注意事项 #

• 模型结构必须和预训练权重完全匹配：culane_r18.pth是基于ResNet18骨干、针对CULane数据集训练的，因此必须加载项目中对应的clr_res18.yaml配置文件，避免手动指定参数出错。
• 确保已安装CLRNet要求的所有依赖（如torch、mmcv等），环境配置符合项目文档要求。

内容的提问来源于stack exchange，提问作者Yiwei Miao