jupyter notebook to explore performance data

Peter · Peter · commit 61b74ea79dfc · 2023-10-30T11:52:01.000-04:00
diff --git a/src/pytorch_kinematics/chain.py b/src/pytorch_kinematics/chain.py
@@ -317,59 +317,6 @@ def forward_kinematics(self, th, frame_indices: Optional = None):
 
         return frame_names_and_transform3ds
 
-    def forward_kinematics_py(self, th, frame_indices: Optional = None):
-        if frame_indices is None:
-            frame_indices = self.get_all_frame_indices()
-
-        th = self.ensure_tensor(th)
-        th = torch.atleast_2d(th)
-
-        b = th.shape[0]
-
-        axes_expanded = self.axes.unsqueeze(0).repeat(b, 1, 1)
-
-        frame_transforms = {}
-
-        # compute all joint transforms at once first
-        # in order to handle multiple joint types without branching, we create all possible transforms
-        # for all joint types and then select the appropriate one for each joint.
-        rev_jnt_transform = tensor_axis_and_angle_to_matrix(axes_expanded, th)
-        pris_jnt_transform = tensor_axis_and_d_to_pris_matrix(axes_expanded, th)
-
-        for frame_idx in frame_indices:
-            frame_transform = torch.eye(4).to(th).unsqueeze(0).repeat(b, 1, 1)
-
-            # iterate down the list and compose the transform
-            for chain_idx in self.parents_indices[frame_idx]:
-                if chain_idx.item() in frame_transforms:
-                    frame_transform = frame_transforms[chain_idx.item()]
-                else:
-                    link_offset_i = self.link_offsets[chain_idx]
-                    if link_offset_i is not None:
-                        frame_transform = frame_transform @ link_offset_i
-
-                    joint_offset_i = self.joint_offsets[chain_idx]
-                    if joint_offset_i is not None:
-                        frame_transform = frame_transform @ joint_offset_i
-
-                    jnt_idx = self.joint_indices[chain_idx]
-                    jnt_type = self.joint_type_indices[chain_idx]
-                    if jnt_type == 0:
-                        pass
-                    elif jnt_type == 1:
-                        jnt_transform_i = rev_jnt_transform[:, jnt_idx]
-                        frame_transform = frame_transform @ jnt_transform_i
-                    elif jnt_type == 2:
-                        jnt_transform_i = pris_jnt_transform[:, jnt_idx]
-                        frame_transform = frame_transform @ jnt_transform_i
-
-            frame_transforms[frame_idx.item()] = frame_transform
-
-        frame_names_and_transform3ds = {self.idx_to_frame[frame_idx]: tf.Transform3d(matrix=transform) for
-                                        frame_idx, transform in frame_transforms.items()}
-
-        return frame_names_and_transform3ds
-
     def ensure_tensor(self, th):
         """
         Converts a number of possible types into a tensor. The order of the tensor is determined by the order
@@ -502,18 +449,6 @@ def forward_kinematics(self, th, end_only: bool = True):
         else:
             return mat
 
-    def forward_kinematics_py(self, th, end_only: bool = True):
-        """ Like the base class, except `th` only needs to contain the joints in the SerialChain, not all joints. """
-        frame_indices, th = self.convert_serial_inputs_to_chain_inputs(end_only, th)
-
-        mat = super().forward_kinematics_py(th, frame_indices)
-
-        if end_only:
-            return mat[self._serial_frames[-1].name]
-        else:
-            return mat
-
-
     def convert_serial_inputs_to_chain_inputs(self, end_only, th):
         if end_only:
             frame_indices = self.get_frame_indices(self._serial_frames[-1].name)
@@ -534,29 +469,3 @@ def convert_serial_inputs_to_chain_inputs(self, end_only, th):
                 if frame.joint.joint_type != 'fixed':
                     th[jnt_idx] = partial_th_i
         return frame_indices, th
-
-    def forward_kinematics_slow(self, th, world=None, end_only=True):
-        if world is None:
-            world = tf.Transform3d()
-        if world.dtype != self.dtype or world.device != self.device:
-            world = world.to(dtype=self.dtype, device=self.device, copy=True)
-        th, N = ensure_2d_tensor(th, self.dtype, self.device)
-        zeros = torch.zeros([N, 1], dtype=world.dtype, device=world.device)
-
-        theta_idx = 0
-        link_transforms = {}
-        trans = tf.Transform3d(matrix=world.get_matrix().repeat(N, 1, 1))
-        for f in self._serial_frames:
-            if f.link.offset is not None:
-                trans = trans.compose(f.link.offset)
-
-            if f.joint.joint_type == "fixed":  # If fixed
-                trans = trans.compose(f.get_transform(zeros))
-            else:
-                joint_transform = f.get_transform(th[:, theta_idx].view(N, 1))
-                trans = trans.compose(joint_transform)
-                theta_idx += 1
-
-            link_transforms[f.link.name] = trans
-
-        return link_transforms[self._serial_frames[-1].link.name] if end_only else link_transforms
diff --git a/tests/.ipynb_checkpoints/viz_fk_perf-checkpoint.ipynb b/tests/.ipynb_checkpoints/viz_fk_perf-checkpoint.ipynb
@@ -0,0 +1,89 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-10-30T15:37:54.975496787Z",
+     "start_time": "2023-10-30T15:37:54.933122254Z"
+    }
+   },
+   "outputs": [
+    {
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'torch'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
+      "\u001b[0;32m<ipython-input-1-5e79930a6286>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Load the data\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mpickle\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'fk_perf.pkl'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'torch'"
+     ]
+    }
+   ],
+   "source": [
+    "# Load the data\n",
+    "import torch\n",
+    "import pickle\n",
+    "\n",
+    "with open('fk_perf.pkl', 'rb') as f:\n",
+    "    headers, data = pickle.load(f)\n",
+    "\n",
+    "import pandas as pd\n",
+    "df = pd.DataFrame(data, columns=headers)\n",
+    "\n",
+    "df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'/usr/bin/python3'"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import sys\n",
+    "sys.executable\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/tests/test_fk_perf.py b/tests/test_fk_perf.py
@@ -5,93 +5,6 @@
 import pytorch_kinematics as pk
 import numpy as np
 
-N = 10000
-number = 100
-
-
-def test_val_fk_correctness():
-    val = pk.build_chain_from_mjcf(open('val.xml').read())
-    val = val.to(dtype=torch.float32, device='cuda')
-
-    th = torch.zeros(N, 20, dtype=torch.float32, device='cuda')
-
-    frame_indices = val.get_frame_indices('left_tool', 'right_tool')
-    t_py = val.forward_kinematics_py(th, frame_indices)
-    t_cpp = val.forward_kinematics(th, frame_indices)
-    l_py = t_py['left_tool'].get_matrix()
-    l_cpp = t_cpp['left_tool'].get_matrix()
-    r_py = t_py['right_tool'].get_matrix()
-    r_cpp = t_cpp['right_tool'].get_matrix()
-
-    assert torch.allclose(l_py, l_cpp)
-    assert torch.allclose(r_py, r_cpp)
-
-
-def test_val_fk_perf():
-    val = pk.build_serial_chain_from_mjcf(open('val.xml').read(), end_link_name='left_tool')
-    val = val.to(dtype=torch.float32, device='cuda')
-
-    th = torch.zeros(N, 20, dtype=torch.float32, device='cuda')
-
-    def _val_old_fk():
-        tg = val.forward_kinematics_slow(th, end_only=True)
-        m = tg.get_matrix()
-        return m
-
-    def _val_new_py_fk():
-        tg = val.forward_kinematics_py(th, end_only=True)
-        m = tg.get_matrix()
-        return m
-
-    def _val_new_cpp_fk():
-        tg = val.forward_kinematics(th, end_only=True)
-        m = tg.get_matrix()
-        return m
-
-    val_old_dt = timeit.timeit(_val_old_fk, number=number)
-    print(f'Val FK OLD dt: {val_old_dt / number:.4f}')
-
-    val_new_py_dt = timeit.timeit(_val_new_py_fk, number=number)
-    print(f'Val FK NEW dt: {val_new_py_dt / number:.4f}')
-
-    val_new_cpp_dt = timeit.timeit(_val_new_cpp_fk, number=number)
-    print(f'Val FK NEW C++ dt: {val_new_cpp_dt / number:.4f}')
-
-    assert val_old_dt > val_new_cpp_dt
-
-
-def test_kuka_fk_perf():
-    kuka = pk.build_serial_chain_from_urdf(open('kuka_iiwa.urdf').read(), end_link_name='lbr_iiwa_link_7')
-    kuka = kuka.to(dtype=torch.float32, device='cuda')
-
-    th = torch.zeros(N, 7, dtype=torch.float32, device='cuda')
-
-    def _kuka_old_fk():
-        tg = kuka.forward_kinematics_slow(th, end_only=True)
-        m = tg.get_matrix()
-        return m
-
-    def _kuka_new_py_fk():
-        tg = kuka.forward_kinematics_py(th, end_only=True)
-        m = tg.get_matrix()
-        return m
-
-    def _kuka_new_cpp_fk():
-        tg = kuka.forward_kinematics(th, end_only=True)
-        m = tg.get_matrix()
-        return m
-
-    kuka_old_dt = timeit.timeit(_kuka_old_fk, number=number)
-    print(f'Kuka FK OLD dt: {kuka_old_dt / number:.4f}')
-
-    kuka_new_py_dt = timeit.timeit(_kuka_new_py_fk, number=number)
-    print(f'Kuka FK NEW dt: {kuka_new_py_dt / number:.4f}')
-
-    kuka_new_cpp_dt = timeit.timeit(_kuka_new_cpp_fk, number=number)
-    print(f'Kuka FK NEW C++ dt: {kuka_new_cpp_dt / number:.4f}')
-
-    assert kuka_old_dt > kuka_new_cpp_dt
-
 
 def main():
     # do an in-depth analysis of multiple models, devices, data types, batch sizes, etc.
@@ -108,6 +21,7 @@ def main():
     devices = ['cpu', 'cuda']
     dtypes = [torch.float32, torch.float64]
     batch_sizes = [1, 10, 100, 1_000, 10_000, 100_000]
+    number = 100
 
     # iterate over all combinations and store in a pandas dataframe
     headers = ['chain', 'device', 'dtype', 'batch_size', 'time']
diff --git a/tests/viz_fk_perf.ipynb b/tests/viz_fk_perf.ipynb
