| 31 | # For TF-TRT: |
| 32 | |
| 33 | class OptimizedModel(): |
| 34 | def __init__(self, saved_model_dir = None): |
| 35 | self.loaded_model_fn = None |
| 36 | |
| 37 | if not saved_model_dir is None: |
| 38 | self.load_model(saved_model_dir) |
| 39 | |
| 40 | |
| 41 | def predict(self, input_data): |
| 42 | if self.loaded_model_fn is None: |
| 43 | raise(Exception("Haven't loaded a model")) |
| 44 | x = tf.constant(input_data.astype('float32')) |
| 45 | labeling = self.loaded_model_fn(x) |
| 46 | try: |
| 47 | preds = labeling['predictions'].numpy() |
| 48 | except: |
| 49 | try: |
| 50 | preds = labeling['probs'].numpy() |
| 51 | except: |
| 52 | try: |
| 53 | preds = labeling[next(iter(labeling.keys()))] |
| 54 | except: |
| 55 | raise(Exception("Failed to get predictions from saved model object")) |
| 56 | return preds |
| 57 | |
| 58 | def load_model(self, saved_model_dir): |
| 59 | saved_model_loaded = tf.saved_model.load(saved_model_dir, tags=[tag_constants.SERVING]) |
| 60 | wrapper_fp32 = saved_model_loaded.signatures['serving_default'] |
| 61 | |
| 62 | self.loaded_model_fn = wrapper_fp32 |
| 63 | |
| 64 | class ModelOptimizer(): |
| 65 | def __init__(self, input_saved_model_dir, calibration_data=None): |