生活随笔
收集整理的這篇文章主要介紹了
CRNN-pytorch模型转libtorch模型踩坑记录
小編覺得挺不錯的,現在分享給大家,幫大家做個參考.
這段時間一直在做CRNN文字識別的問題,從pytorch中訓練好的模型然后轉到libtorch中去,但是CRNN提供的代碼沒有轉libtorch模型的部分,于是就在網上到處亂找,其中找到了這篇轉的代碼crnn模型轉換。感興趣的可以進去看一下,這個大哥也是從git上找的別人,可能他的訓練代碼的網絡和轉換代碼的網絡是一樣的結構,所以它可以轉換成功。
我在轉換的時候首先是報不能轉換成功,然后修改了部分代碼之后,能夠轉換成功,但是在libtorch上推理的時候為啥一點準確度都沒有呢,哪怕準確一點就可以啊,然后就開始各種查,檢測每個變量時候和pytorch的變量是否一樣,不管是圖片的預處理還是從網絡輸出的張量形狀,最后還是不行,就尋思是不是轉換的時候出了問題,然后就在他的基礎上進行了修改,最后可以在自己的網絡結構上轉換成功;
import numpy
as np
import time
import cv2
import torch
from torch
.autograd
import Variable
import lib
.utils
.utils
as utils
import lib
.models
.crnn
as crnn
import lib
.config
.alphabets
as alphabets
import yaml
from easydict
import EasyDict
as edict
import os
import random
def deal_cfg(path_cfg
):with open(path_cfg
, 'r') as f
:config
= yaml
.load
(f
)config
= edict
(config
)config
.DATASET
.ALPHABETS
= alphabets
.alphabetconfig
.MODEL
.NUM_CLASSES
= len(config
.DATASET
.ALPHABETS
)print("len(config.DATASET.ALPHABETS):",len(config
.DATASET
.ALPHABETS
))return config
if __name__
== '__main__':path_cfg
= "./lib/config/OWN_config.yaml"path_pth
= "/root/CRNN_Chinese_Characters_Rec-stable/output/OWN/crnn/acc_0.9766.pth"print("cuda?", torch
.cuda
.is_available
())config
= deal_cfg
(path_cfg
)device
= torch
.device
('cuda:0') if torch
.cuda
.is_available
() else torch
.device
('cpu')converter
= utils
.strLabelConverter
(config
.DATASET
.ALPHABETS
)model
= crnn
.get_crnn
(config
).to
(device
)print('loading pretrained model from {0}'.format(path_pth
))checkpoint
= torch
.load
(path_pth
)if 'state_dict' in checkpoint
.keys
():model
.load_state_dict
(checkpoint
['state_dict'])else:model
.load_state_dict
(checkpoint
)model
.eval()img
= torch
.ones
(1,1,32,320).cuda
()img
= img
.type(torch
.FloatTensor
)img
= img
.to
(device
)traced_script_module
= torch
.jit
.trace
(model
, img
)traced_script_module
.save
("./acc.9833.pt")
這樣轉成功了,我的網絡結構和上面那位大哥提供連接的結構不一樣,所以我猜是結構的原因,導致推理不準確,
另外還是有一點是我在訓練的時候添加了dropout層,在轉換的時候我把dropout層給注銷掉了,我也不知道該不該注銷,先這樣吧。接下來驗證一下精度是否可以。
總結
以上是生活随笔為你收集整理的CRNN-pytorch模型转libtorch模型踩坑记录的全部內容,希望文章能夠幫你解決所遇到的問題。
如果覺得生活随笔網站內容還不錯,歡迎將生活随笔推薦給好友。