本文仅适用于使用英文标签训练结束之后,将英文标签的权重替换成中文标签的权重
一、使用数据字典进行映射
通过数据字典将权重中对应的英文标签一一替换成中文标签
1、加载模型并显示英文标签
w=torch.load('/path/to/your/weight')
使用torch导入英文训练的权重
print(w.get('model').names)
依次打印出该权重下对应的标签名称
2、字典映射
word_map = {
'ENG1': 'CHN1',
'ENG2': 'CHN2',
'ENG3': 'CHN3',
}
3、遍历权重,将对应的英文替换成中文
for i in range(len(w.get('model').names)):
if w.get('model').names[i] in word_map:
w.get('model').names[i] = word_map[w.get('model').names[i]]
4、打印替换之后的权重标签并保存
# 打印替换后的列表
print('替换后')
print(w.get('model').names)
#保存替换后的模型
torch.save(w,'/path/to/your/save/folder')
二、完整代码
import torch
# 加载模型
w=torch.load('/path/to/your/weight')
#打印所有name
print(w.get('model').names)
# 定义一个将英文单词映射到中文单词的字典
word_map = {
'ENG1': 'CHN1',
'ENG2': 'CHN2',
'ENG3': 'CHN3',
}
# 遍历列表,将每个英文单词替换为其中文对应词
for i in range(len(w.get('model').names)):
if w.get('model').names[i] in word_map:
w.get('model').names[i] = word_map[w.get('model').names[i]]
# 打印替换后的列表
print('替换后')
print(w.get('model').names)
#保存替换后的模型
torch.save(w,'/path/to/your/save/folder')