Upload folder using huggingface_hub

README.md

@@ -1,85 +1,88 @@
----
-license: bsd-3-clause
-pipeline_tag: robotics
-tags:
-  - ballnet
-  - metaball
-  - multimodal
-  - onnx
-  - pytorch
-library_name: transformers
-datasets:
-  - han-xudong/ballnet-100k
----
-
-<p align="center">
-  <img src="https://cdn-uploads.huggingface.co/production/uploads/670336f25ce58dd0c33e32e3/EC3s3yVe3mLlq6t6QxQGL.jpeg"
-       style="width:100%; max-width:1000px; height:auto;">
-</p>
-
-# BallNet
-
-## Model Description
-
-BallNet is an MLP model designed for the Metaball. It can predict both 6D force and 3D shape (mesh nodes) from the 6D motion of the ball.
-
-Try it out on the [Spaces demo](https://huggingface.co/spaces/han-xudong/ballnet-demo).
-
-- Developer: Xudong Han, Tianyu Wu, Fang Wan, and Chaoyang Song.
-- Model type: MLP
-- License: BSD-3-Clause
-- Github: [https://github.com/han-xudong/metaball](https://github.com/han-xudong/metaball)
-
-## Intended Use
-
-This model is intended for researchers and developers working in robotics and tactile sensing. It can be used to enhance the capabilities of robotic systems by providing accurate predictions of force and shape based on tactile data.
-
-To load the model:
-
-```python
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained("han-xudong/ballnet", trust_remote_code=True)
-x = torch.zeros((1, 6))  # Example input: batch size of 1, 6D motion
-output = model(x)
-```
-
-Or to load the ONNX version:
-
-```python
-# Example code to load onnx
-import onnxruntime as ort
-import numpy as np
-from huggingface_hub import hf_hub_download
-
-onnx_model_path = hf_hub_download("han-xudong/ballnet", filename="model.onnx")
-ort_session = ort.InferenceSession(onnx_model_path)
-
-# Example input
-x = np.zeros((1, 6), dtype=np.float32)  # Batch size of 1, 6D motion
-output = ort_session.run(None, {"motion": x})
-```
-
-## Training Data
-
-The model was trained on the [BallNet-100K](https://huggingface.co/datasets/han-xudong/ballnet-100k) dataset, which includes a variety of motion, force, and shape data collected by finite element simulations.
-
-## Citation
-
-If you use this model in your research, please cite the following papers:
-
-```bibtex
-@article{liu2024proprioceptive,
-  title={Proprioceptive learning with soft polyhedral networks},
-  author={Liu, Xiaobo and Han, Xudong and Hong, Wei and Wan, Fang and Song, Chaoyang},
-  journal={The International Journal of Robotics Research},
-  volume = {43},
-  number = {12},
-  pages = {1916-1935},
-  year = {2024},
-  publisher={SAGE Publications Sage UK: London, England},
-  doi = {10.1177/02783649241238765}
-}
-```
-
-[](https://arxiv.org/abs/2308.08538)
+---
+license: bsd-3-clause
+pipeline_tag: robotics
+tags:
+  - ballnet
+  - metaball
+  - multimodal
+  - onnx
+  - pytorch
+library_name: transformers
+datasets:
+  - han-xudong/ballnet-100k
+---
+
+# Model Card for BallNet
+
+## Table of Contents
+
+- [Model Card for BallNet](#model-card-for-ballnet)
+  - [Table of Contents](#table-of-contents)
+  - [Model Description](#model-description)
+  - [Intended Use](#intended-use)
+  - [Training Data](#training-data)
+  - [Citation](#citation)
+
+## Model Description
+
+BallNet is an MLP model designed for the Metaball. It can predict both 6D force and 3D shape (mesh nodes) from the 6D motion of the ball.
+
+Try it out on the [Spaces demo](https://huggingface.co/spaces/han-xudong/ballnet-demo).
+
+- Developer: Xudong Han, Tianyu Wu, Fang Wan, and Chaoyang Song.
+- Model type: MLP
+- License: BSD-3-Clause
+
+## Intended Use
+
+This model is intended for researchers and developers working in robotics and tactile sensing. It can be used to enhance the capabilities of robotic systems by providing accurate predictions of force and shape based on tactile data.
+
+To load the model:
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("han-xudong/ballnet", trust_remote_code=True)
+x = torch.zeros((1, 6))  # Example input: batch size of 1, 6D motion
+output = model(x)
+```
+
+Or to load the ONNX version:
+
+```python
+# Example code to load onnx
+import onnxruntime as ort
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+onnx_model_path = hf_hub_download("han-xudong/ballnet", filename="model.onnx")
+ort_session = ort.InferenceSession(onnx_model_path)
+
+# Example input
+x = np.zeros((1, 6), dtype=np.float32)  # Batch size of 1, 6D motion
+output = ort_session.run(None, {"motion": x})
+```
+
+## Training Data
+
+The model was trained on the [BallNet-100K](https://huggingface.co/datasets/han-xudong/ballnet-100k) dataset, which includes a variety of motion, force, and shape data collected by finite element simulations.
+
+## Citation
+
+If you use this model in your research, please cite the following papers:
+
+```bibtex
+@article{liu2024proprioceptive,
+  title={Proprioceptive learning with soft polyhedral networks},
+  author={Liu, Xiaobo and Han, Xudong and Hong, Wei and Wan, Fang and Song, Chaoyang},
+  journal={The International Journal of Robotics Research},
+  volume = {43},
+  number = {12},
+  pages = {1916-1935},
+  year = {2024},
+  publisher={SAGE Publications Sage UK: London, England},
+  doi = {10.1177/02783649241238765}
+}
+```
+
+[](https://arxiv.org/abs/2308.08538)

model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1138bf9e90c36e6e7e68e397dc8d109c444d96e95f3e8c45d65e5c3b756528c1
-size 15815554
+oid sha256:b543f9d911a78fdc5648eca294aa7e04e7edfeb441e43f4dfc3855a02e943f36
+size 17347939

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:960574e99b9df54c0588506eb481c7e15ea0c4ff91813276edcac8a2864bec30
+size 17346660

modeling.py

@@ -7,48 +7,119 @@ class BallNetConfig(PretrainedConfig):
 
     def __init__(
         self,
-        x_dim=[6],
-        y_dim=[6, 1800],
-        h1_dim=[100, 1000],
-        h2_dim=[100, 1000],
+        x_dim: List[int],
+        y_dim: List[int],
+        hidden_dim: List[List[int]],
+        mean: List[float],
+        std: List[float],
         **kwargs,
     ):
         super().__init__(**kwargs)
+
         self.x_dim = x_dim
         self.y_dim = y_dim
-        self.h1_dim = h1_dim
-        self.h2_dim = h2_dim
+        self.hidden_dim = hidden_dim
+        self.mean = mean
+        self.std = std
+
+
+class Normalizer(nn.Module):
+    def __init__(self, mean: Tensor, std: Tensor, eps: float = 1e-8):
+        super().__init__()
+        self.register_buffer("mean", mean)
+        self.register_buffer("std", std)
+        self.eps = eps
+
+    def normalize(self, x: Tensor) -> Tensor:
+        return (x - self.mean) / (self.std + self.eps)
+
+    def denormalize(self, x: Tensor) -> Tensor:
+        return x * (self.std + self.eps) + self.mean
 
 
-class BallNet(PreTrainedModel):
+class BallNetModel(PreTrainedModel):
     config_class = BallNetConfig
+    base_model_prefix = "ballnet"
+    supports_gradient_checkpointing = False
 
-    def __init__(self, config):
+    def __init__(self, config: BallNetConfig):
         super().__init__(config)
+
         self.x_dim = config.x_dim
         self.y_dim = config.y_dim
-        self.h1_dim = config.h1_dim
-        self.h2_dim = config.h2_dim
+        self.hidden_dim = config.hidden_dim
+
+        # ---------- split mean / std ----------
+        x_mean, x_std = [], []
+        y_mean, y_std = [], []
+
+        data_dim = self.x_dim + self.y_dim
+        data_start = 0
+
+        for i, dim in enumerate(data_dim):
+            data_end = data_start + dim
+            if i < len(self.x_dim):
+                x_mean.append(config.mean[data_start:data_end])
+                x_std.append(config.std[data_start:data_end])
+            else:
+                y_mean.append(config.mean[data_start:data_end])
+                y_std.append(config.std[data_start:data_end])
+            data_start = data_end
+
+        # ---------- normalizers ----------
+        self.x_normalizers = nn.ModuleList(
+            [
+                Normalizer(
+                    mean=torch.tensor(x_mean[i], dtype=torch.float32),
+                    std=torch.clamp(
+                        torch.tensor(x_std[i], dtype=torch.float32), min=1e-8
+                    ),
+                )
+                for i in range(len(self.x_dim))
+            ]
+        )
+
+        self.y_normalizers = nn.ModuleList(
+            [
+                Normalizer(
+                    mean=torch.tensor(y_mean[i], dtype=torch.float32),
+                    std=torch.clamp(
+                        torch.tensor(y_std[i], dtype=torch.float32), min=1e-8
+                    ),
+                )
+                for i in range(len(self.y_dim))
+            ]
+        )
+
+        # ---------- estimators ----------
+        self.estimators = nn.ModuleList()
 
-        self.model = nn.ModuleDict()
-        # Define the model architecture
         for i in range(len(self.y_dim)):
-            self.model[f"estimator_{i}"] = nn.Sequential(
-                nn.Linear(self.x_dim[0], self.h1_dim[i]),
-                nn.ReLU(),
-                nn.Linear(self.h1_dim[i], self.h2_dim[i]),
-                nn.ReLU(),
-                nn.Linear(self.h2_dim[i], self.y_dim[i]),
-            )
-
-        # initialize weights
+            layers = []
+            in_dim = self.x_dim[0]
+
+            for out_dim in self.hidden_dim[i]:
+                layers.append(nn.Linear(in_dim, out_dim))
+                layers.append(nn.ReLU())
+                in_dim = out_dim
+
+            layers.append(nn.Linear(in_dim, self.y_dim[i]))
+            self.estimators.append(nn.Sequential(*layers))
+
         self.post_init()
 
-    def forward(self, x):
+    def forward(self, x: Tensor, **kwargs):
+        """
+        x: (B, 6)
+        """
+        x = self.x_normalizers[0].normalize(x)
+
         outputs = []
         for i in range(len(self.y_dim)):
-            # Get the estimator for the i-th output
-            estimator = self.model[f"estimator_{i}"]
-            y = estimator(x)
+            y = self.estimators[i](x)
+            y = self.y_normalizers[i].denormalize(y)
             outputs.append(y)
-        return outputs
+
+        return {
+            "outputs": outputs
+        }