вот мой код:< /p>
model.py
Код: Выделить всё
import torch
import torch.nn as nn
from torchcrf import CRF
from ASR.BiLstm.config import CFG
class BiLSTM(nn.Module):
def __init__(self):
super().__init__()
self.cfg = CFG()
self.embedding_layer = nn.Embedding(
self.cfg.token_size, self.cfg.hidden_size)
self.lstm = nn.LSTM(
input_size=self.cfg.hidden_size,
hidden_size=self.cfg.hidden_size,
batch_first=True,
bidirectional=True)
self.fc = nn.Linear(self.cfg.hidden_size * 2, self.cfg.entity_type_num)
self.crf_layer = CRFLayer(num_tags=self.cfg.entity_type_num)
self.drop = nn.Dropout(0.5)
def forward(
self,
x: torch.Tensor,
senquence_length: torch.Tensor,
label: torch.Tensor = None):
x = self.embedding_layer(x)
output_sequence, (_, _) = self.lstm(x) # [batch,seq_len,hidden_size*2]
lstm_features = self.drop(output_sequence)
lstm_logits = self.fc(lstm_features)
mask = self.crf_layer.create_mask(
senquence_length, max_length=lstm_logits.shape[1])
if label is not None:
loss = -self.crf_layer.crf(emissions=lstm_logits,
tags=label, mask=mask)
return loss
else:
output = self.crf_layer.crf.decode(
emissions=lstm_logits, mask=mask)
# 将列表转换为张量
output_tensor = torch.tensor(output, dtype=torch.long)
return output_tensor
class CRFLayer(nn.Module):
def __init__(self, num_tags):
super().__init__()
self.cfg = CFG()
self.crf = CRF(num_tags=num_tags, batch_first=True)
def create_mask(self, sequence_length: torch.Tensor, max_length: int):
batch_size = sequence_length.shape[0]
arange = torch.arange(max_length).unsqueeze(0).expand(batch_size, -1)
# 如果你的模型是在GPU上运行，确保arange也在相同的设备上
if self.cfg.device == 'cuda':
arange = arange.to(sequence_length.device)
# 将范围张量扩展成 [batch_size, max_sequence_length] 的形状，
# 并与每个序列的实际长度进行比较，小于实际长度的位置为True，否则为False
mask = arange < sequence_length.unsqueeze(1)
return mask
Код: Выделить всё
def export_onnx(model, dataset):
'''
导出onnx
'''
cfg = CFG()
model.eval()
input_sequence, sequence_length = dataset.transfer_query('画个本环')
inputs = (input_sequence, sequence_length)
# 使用torch.jit.trace进行追踪处理
traced_model = torch.jit.trace(model, inputs)
torch.onnx.export(
model=traced_model,
args=inputs,
f='BiLstm_model.onnx',
input_names=['input_sequence', 'sequence_length'],
output_names=['decode_result'],
dynamic_axes={
'input_sequence': {0: 'batch_size'},
'sequence_length': {0: 'batch_size'},
'decode_result': {0: 'batch_size'}
},
do_constant_folding=True,
opset_version=15
)
try:
onnx.checker.check_model(r'BiLstm_model.onnx')
netron.start(r'BiLstm_model.onnx')
except Exception as e:
print(e)
if __name__ == '__main__':
# train()
model = BiLSTM()
checkpoint = torch.load(
r'C:\Users\admin\Desktop\项目\ASR\interaction\ASR\BiLstm\checkpoints\checkpoint_4.pth')
# 加载模型的state_dict
model.load_state_dict(checkpoint['model_state_dict'])
dataset = MyData()
export_onnx(model, dataset)
и форма "input_sequence" — факел. size([1,20]), форма последовательности_длины — torch.size([1]).
А вот график модели, сгенерированный 'netron.start':
введите здесь описание изображения
Моя версия Python — 3.11. версия факела — 2.3.1+cu121. версия onnx — 1.16.1.
Я искал в Интернете свою проблему, но не нашел такой же проблемы.
Я хочу знать, в чем Я могу правильно преобразовать модель в формат onnx
Подробнее здесь: https://stackoverflow.com/questions/790 ... nnx-format
