File size: 1,754 Bytes
22fe4c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
from transformers import PretrainedConfig, PreTrainedModel
import torch
import torch.nn as nn
class ONNXBaseConfig(PretrainedConfig):
model_type = "onnx-base"
def __init__(self, model_path=None, **kwargs):
self.model_path = model_path
super().__init__(**kwargs)
model_directory = './new_model'
config = ONNXBaseConfig(model_path='model.onnx')
config.save_pretrained(save_directory=model_directory)
class ONNXBaseModel(PreTrainedModel):
config_class = ONNXBaseConfig
def __init__(self, config):
super().__init__(config)
self.dummy_param = nn.Parameter(torch.zeros(0))
def forward(self, inputs):
return torch.zeros_like(inputs)
def save_pretrained(self, save_directory: str, **kwargs):
super().save_pretrained(save_directory=save_directory, **kwargs)
onnx_file_path = save_directory + '/model.onnx'
dummy_input = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
torch.onnx.export(self, dummy_input, onnx_file_path,
input_names=['input'], output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
# Initialize model
model = ONNXBaseModel(config)
# Save model
model.save_pretrained(save_directory=model_directory)
model = model.from_pretrained(model_directory)
# Test model
dummy_input = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
output_tensor = model(dummy_input)
print(output_tensor)
# Test the onnx model
onnx_file_path = model_directory + '/model.onnx'
import onnx
import onnxruntime as ort
ort_session = ort.InferenceSession(onnx_file_path)
outputs = ort_session.run(None, {'input': dummy_input.numpy()})
print("Model output:", outputs)
|