I've written pytorch model,than converted it to .onnx format, buit a tf representation, tried to convert it to tflite and while exporting graph error occured this is how I saved the model and converted it to .onnx: (LeNet5 is the name of the model class)
torch.save(model.state_dict(), '/Plastic_dataset/MyDrive/model.pth')
trained_model = LeNet5(activation='relu', conv_size=3, pooling='avg', use_batch_norm=True)
trained_model.load_state_dict(torch.load('/Plastic_dataset/MyDrive/model.pth'), strict = False)
dummy_input = torch.randn(1, 1, 64, 64)
torch.onnx.export(trained_model, dummy_input, '/Plastic_dataset/MyDrive/model.onnx')onnx_model = onnx.load('/Plastic_dataset/MyDrive/model.onnx')
tf_rep = prepare(onnx_model)
tf_rep.export_graph("/Plastic_dataset/MyDrive/model.pb")
converter = tf.lite.TFLiteConverter.from_frozen_graph(
"/Plastic_dataset/MyDrive/model.pb", tf_rep.inputs, tf_rep.outputs)
tflite_model = converter.convert()
open("/Plastic_dataset/MyDrive/model.tflite", "wb").write(tflite_model)
