Triton ensemble export (#2251)

ludwig/export.py

@@ -78,7 +78,8 @@ def export_triton(model_path, output_path="model_repository", model_name="ludwig
 
     model = LudwigModel.load(model_path)
     os.makedirs(output_path, exist_ok=True)
-    utils_export_triton(model, output_path, model_name, model_version)
+
+    utils_export_triton(model=model, output_path=output_path, model_name=model_name, model_version=model_version)
 
     logger.info(f"Saved to: {output_path}")
 

ludwig/utils/misc_utils.py

@@ -141,3 +141,7 @@ def set_saved_weights_in_checkpoint_flag(config):
     """
     for input_feature in config.get("input_features", []):
         input_feature["saved_weights_in_checkpoint"] = True
+
+
+def remove_empty_lines(str):
+    return "\n".join([line.rstrip() for line in str.split("\n") if line.rstrip()])

ludwig/utils/torch_utils.py

@@ -188,7 +188,8 @@ def input_dtype(self):
     @abstractmethod
     def input_shape(self) -> torch.Size:
         """Returns size of the input tensor without the batch dimension."""
-        raise NotImplementedError("Abstract class.")
+        pass
+        # raise NotImplementedError("Abstract class.")
 
     @property
     def output_shape(self) -> torch.Size:
@@ -307,3 +308,11 @@ def _set_torch_init_params(params: Optional[Tuple]):
 
 def _get_torch_init_params() -> Optional[Tuple]:
     return _TORCH_INIT_PARAMS
+
+
+def model_size(model: nn.Module):
+    """Computes PyTorch model size in bytes."""
+    size = 0
+    size += sum(param.nelement() * param.element_size() for param in model.parameters())
+    size += sum(buffer.nelement() * buffer.element_size() for buffer in model.buffers())
+    return size