Skip to content

Commit

Permalink
add demo of converting kornia functions from pytorch to tensorflow
Browse files Browse the repository at this point in the history
  • Loading branch information
@Sam-Armstrong
Sam-Armstrong committed Aug 23, 2024
1 parent 3a95603 commit 449c40b
Show file tree
Hide file tree
Showing 2 changed files with 332 additions and 0 deletions.
6 changes: 6 additions & 0 deletions learn_the_basics.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,11 @@ Learn the basics
.. grid:: 1 1 3 3
:gutter: 4

.. grid-item-card:: Transpiling Functions from PyTorch to TensorFlow
:link: learn_the_basics/torch_to_tf_functions.ipynb

Transpiling Kornia functions to TensorFlow.

.. grid-item-card:: Trace Code
:link: learn_the_basics/03_trace_code.ipynb

Expand All @@ -23,6 +28,7 @@ Learn the basics
:hidden:
:maxdepth: -1

learn_the_basics/torch_to_tf_functions.ipynb
learn_the_basics/03_trace_code.ipynb
learn_the_basics/05_lazy_vs_eager.ipynb
learn_the_basics/06_how_to_use_decorators.ipynb
326 changes: 326 additions & 0 deletions learn_the_basics/torch_to_tf_functions.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,326 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Transpiling Functions from PyTorch to TensorFlow"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can install the dependencies required for this notebook by running the cell below ⬇️, or check out the [Get Started](https://ivy.dev/docs/overview/get_started.html) section of the docs to find out more about installing ivy."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install ivy\n",
"!pip install torch\n",
"!pip install tensorflow\n",
"!pip install kornia"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here we'll go through an example of how any function using torch can be converted, and used in, tensorflow via `ivy.transpile`. We'll use kornia as our example, which is a state-of-the-art computer vision library built on top of torch."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"First, some boiler plate imports:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import ivy\n",
"import kornia\n",
"import numpy as np\n",
"import tensorflow as tf\n",
"import torch"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, lets transpile a kornia function that we want to use in tensorflow. The `ivy.transpile` call returns a new tensorflow function which is mathematically equivalent to the torch function we passed. This can take up to a minute to run."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"tf_rgb_to_grayscale = ivy.transpile(kornia.color.rgb_to_grayscale, source=\"torch\", target=\"tensorflow\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can now use this function exactly as the original kornia function would be, just passing tensorflow tensors rather than torch tensors:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<tf.Tensor: shape=(1, 1, 28, 28), dtype=float32, numpy=\n",
"array([[[[0.49296033, 0.85580593, 0.73077095, 0.6316749 , 0.36565533,\n",
" 0.6911214 , 0.5613365 , 0.43097103, 0.5831474 , 0.4034585 ,\n",
" 0.55593616, 0.65984905, 0.9567939 , 0.32861516, 0.47598985,\n",
" 0.69988805, 0.5131729 , 0.48460233, 0.7778542 , 0.38005152,\n",
" 0.44124812, 0.54439026, 0.13614878, 0.86290044, 0.30565348,\n",
" 0.15573153, 0.54602885, 0.32699087],\n",
" [0.7613822 , 0.8590437 , 0.6796573 , 0.5998577 , 0.37982064,\n",
" 0.26729056, 0.5365731 , 0.43443426, 0.6852775 , 0.33589965,\n",
" 0.55056024, 0.47998315, 0.38066688, 0.5961556 , 0.4367127 ,\n",
" 0.24993691, 0.43087602, 0.26046935, 0.56437266, 0.213066 ,\n",
" 0.57331777, 0.31458968, 0.47058818, 0.32165387, 0.5530722 ,\n",
" 0.4742369 , 0.46425182, 0.24489878],\n",
" [0.4989296 , 0.6239719 , 0.29515338, 0.39768082, 0.52449703,\n",
" 0.44410226, 0.70339715, 0.7027339 , 0.9493737 , 0.3313049 ,\n",
" 0.29825547, 0.25047007, 0.25240797, 0.47741622, 0.4410479 ,\n",
" 0.4535742 , 0.1399214 , 0.14160769, 0.521669 , 0.7550629 ,\n",
" 0.38794404, 0.87214035, 0.65314126, 0.7622461 , 0.39658976,\n",
" 0.5003823 , 0.3620267 , 0.5209271 ],\n",
" [0.43946576, 0.5861 , 0.6435085 , 0.54303294, 0.7460731 ,\n",
" 0.4735044 , 0.43913406, 0.72739 , 0.07200881, 0.55436087,\n",
" 0.8670199 , 0.71886593, 0.9110407 , 0.555056 , 0.3284469 ,\n",
" 0.72270083, 0.84436476, 0.8195878 , 0.5123501 , 0.6761943 ,\n",
" 0.39222774, 0.2453723 , 0.26540205, 0.532107 , 0.6674917 ,\n",
" 0.12903629, 0.78137964, 0.36030492],\n",
" [0.2759029 , 0.598225 , 0.8169377 , 0.35551432, 0.6354972 ,\n",
" 0.30139905, 0.7702066 , 0.7159862 , 0.5794634 , 0.36993578,\n",
" 0.22942673, 0.4686075 , 0.5255757 , 0.46116158, 0.25437534,\n",
" 0.50529456, 0.6788846 , 0.29976118, 0.46853146, 0.8404826 ,\n",
" 0.3029947 , 0.32095715, 0.5657457 , 0.6648414 , 0.7377707 ,\n",
" 0.45782346, 0.63786435, 0.6464489 ],\n",
" [0.48081 , 0.3674953 , 0.83497345, 0.6276602 , 0.5472318 ,\n",
" 0.23523967, 0.8159783 , 0.52615684, 0.30982086, 0.38690192,\n",
" 0.44237942, 0.43110123, 0.15921232, 0.84909004, 0.14961666,\n",
" 0.23127733, 0.6665542 , 0.47715086, 0.18885304, 0.4050523 ,\n",
" 0.2034079 , 0.40680784, 0.30552164, 0.187272 , 0.47815746,\n",
" 0.35432848, 0.33978945, 0.2588532 ],\n",
" [0.76345336, 0.72176397, 0.19500253, 0.7186791 , 0.48291487,\n",
" 0.1773673 , 0.27616158, 0.58042 , 0.65294915, 0.5219107 ,\n",
" 0.39325932, 0.30467382, 0.5634152 , 0.49274728, 0.43373075,\n",
" 0.5558106 , 0.6892457 , 0.6364119 , 0.2070568 , 0.62193185,\n",
" 0.5570237 , 0.3904049 , 0.6584081 , 0.46104455, 0.5386716 ,\n",
" 0.5024127 , 0.77666676, 0.5116049 ],\n",
" [0.33227512, 0.9302091 , 0.76431215, 0.46174574, 0.8608088 ,\n",
" 0.45231408, 0.5376172 , 0.66770905, 0.3384223 , 0.67850643,\n",
" 0.64262515, 0.5830049 , 0.16010492, 0.7425901 , 0.69794947,\n",
" 0.43080875, 0.31582838, 0.7201378 , 0.70379734, 0.7707326 ,\n",
" 0.43489042, 0.1674592 , 0.30420962, 0.64085615, 0.5562887 ,\n",
" 0.48882553, 0.72651064, 0.6645107 ],\n",
" [0.09576342, 0.5973712 , 0.51216173, 0.66430944, 0.49837804,\n",
" 0.8860815 , 0.7196047 , 0.6500621 , 0.48323685, 0.5198936 ,\n",
" 0.41961774, 0.4793267 , 0.5327019 , 0.50169057, 0.39731294,\n",
" 0.46453282, 0.6374841 , 0.20438369, 0.31915343, 0.43276027,\n",
" 0.40353265, 0.19193788, 0.49602816, 0.38916475, 0.3553547 ,\n",
" 0.4695943 , 0.31296828, 0.40816274],\n",
" [0.50160116, 0.73214316, 0.16340554, 0.3155364 , 0.3054855 ,\n",
" 0.6878571 , 0.66533625, 0.6297885 , 0.42204636, 0.62815017,\n",
" 0.28406966, 0.46879068, 0.71870494, 0.30653143, 0.3309901 ,\n",
" 0.5741293 , 0.26462665, 0.25825614, 0.47126928, 0.66076976,\n",
" 0.50600374, 0.28697857, 0.38552547, 0.6514376 , 0.74130857,\n",
" 0.7184323 , 0.6693372 , 0.770887 ],\n",
" [0.4992302 , 0.7718851 , 0.3916285 , 0.3516099 , 0.33137748,\n",
" 0.45611387, 0.44159424, 0.5586176 , 0.4695295 , 0.76106656,\n",
" 0.52622896, 0.3806168 , 0.83740973, 0.30983514, 0.5512657 ,\n",
" 0.5674867 , 0.44561777, 0.24020231, 0.7653066 , 0.7046461 ,\n",
" 0.31364638, 0.59066314, 0.57550037, 0.6942578 , 0.20950222,\n",
" 0.81518644, 0.39220342, 0.7984339 ],\n",
" [0.20109056, 0.55548745, 0.34951282, 0.46904862, 0.23029406,\n",
" 0.4211571 , 0.32920533, 0.6031544 , 0.860837 , 0.53341055,\n",
" 0.23700455, 0.6015269 , 0.49808854, 0.6447939 , 0.6604945 ,\n",
" 0.42902747, 0.42047936, 0.20613474, 0.42102173, 0.82564986,\n",
" 0.30528295, 0.4714049 , 0.5492011 , 0.28521618, 0.88891816,\n",
" 0.11472917, 0.26778728, 0.48746774],\n",
" [0.19233656, 0.3556826 , 0.5215195 , 0.6820254 , 0.16108279,\n",
" 0.6925031 , 0.552746 , 0.66557044, 0.75153124, 0.1934661 ,\n",
" 0.7330578 , 0.6626441 , 0.76607925, 0.38522887, 0.17734797,\n",
" 0.44826382, 0.6780246 , 0.27378416, 0.7493787 , 0.5452368 ,\n",
" 0.42360693, 0.5951244 , 0.5365505 , 0.2681636 , 0.6623356 ,\n",
" 0.33533204, 0.43890804, 0.6343641 ],\n",
" [0.12702148, 0.6168841 , 0.85134834, 0.64549434, 0.38682976,\n",
" 0.49264285, 0.5104517 , 0.18002039, 0.88599735, 0.34265348,\n",
" 0.5849319 , 0.6292272 , 0.4080159 , 0.6458118 , 0.45335168,\n",
" 0.46115267, 0.72040766, 0.4368204 , 0.41363248, 0.4124769 ,\n",
" 0.5162779 , 0.5182378 , 0.76471037, 0.53513986, 0.73397994,\n",
" 0.6491305 , 0.45874202, 0.46488976],\n",
" [0.60448766, 0.63621926, 0.6425189 , 0.28838447, 0.63368446,\n",
" 0.5947221 , 0.5782955 , 0.36294326, 0.48188877, 0.52637 ,\n",
" 0.18949348, 0.5689725 , 0.41363242, 0.71333516, 0.7084434 ,\n",
" 0.5964619 , 0.5581088 , 0.42732418, 0.29512846, 0.51097643,\n",
" 0.57651883, 0.81257486, 0.17240909, 0.7350769 , 0.6440831 ,\n",
" 0.7716154 , 0.7696945 , 0.42528978],\n",
" [0.5479724 , 0.81761336, 0.25962776, 0.7824863 , 0.6940273 ,\n",
" 0.11094152, 0.47357208, 0.45543537, 0.27674046, 0.6603243 ,\n",
" 0.26991653, 0.20988956, 0.36109203, 0.22456712, 0.7642828 ,\n",
" 0.74795896, 0.8624945 , 0.8977016 , 0.52269655, 0.5808172 ,\n",
" 0.30616492, 0.658983 , 0.58873963, 0.3143775 , 0.41260308,\n",
" 0.20255119, 0.32344916, 0.7811551 ],\n",
" [0.7413846 , 0.70024586, 0.45004362, 0.26675087, 0.6367631 ,\n",
" 0.36855108, 0.33443588, 0.36071956, 0.7146904 , 0.55997235,\n",
" 0.537416 , 0.8246956 , 0.3575454 , 0.48347342, 0.4314622 ,\n",
" 0.62127906, 0.37857515, 0.5016419 , 0.61000717, 0.68972874,\n",
" 0.41715318, 0.4661981 , 0.6399835 , 0.3393832 , 0.27402696,\n",
" 0.732831 , 0.63776344, 0.36250553],\n",
" [0.20105354, 0.34666434, 0.47403473, 0.5376881 , 0.23853485,\n",
" 0.23298672, 0.6009976 , 0.731824 , 0.22896147, 0.7476525 ,\n",
" 0.69204986, 0.48640496, 0.8111195 , 0.74533314, 0.35637805,\n",
" 0.2277823 , 0.41439864, 0.3882094 , 0.34607103, 0.5740501 ,\n",
" 0.45712376, 0.36859727, 0.38998622, 0.18217316, 0.30438387,\n",
" 0.7317879 , 0.4312286 , 0.69588417],\n",
" [0.3016262 , 0.73031735, 0.7620005 , 0.20053042, 0.3871917 ,\n",
" 0.38995925, 0.2710407 , 0.6272299 , 0.5048422 , 0.7078004 ,\n",
" 0.60675156, 0.3457363 , 0.2646299 , 0.5620334 , 0.47565135,\n",
" 0.8230689 , 0.44504172, 0.54050905, 0.31101573, 0.24286434,\n",
" 0.4821655 , 0.6120316 , 0.55031794, 0.29379642, 0.3952177 ,\n",
" 0.51321733, 0.6849346 , 0.36655322],\n",
" [0.56821465, 0.28197667, 0.27508613, 0.26472732, 0.3951432 ,\n",
" 0.25014058, 0.61152416, 0.5316441 , 0.2902867 , 0.6716291 ,\n",
" 0.15485194, 0.65195155, 0.6394027 , 0.48777273, 0.45825216,\n",
" 0.43302876, 0.37537804, 0.7100045 , 0.42390877, 0.419715 ,\n",
" 0.26340058, 0.5737571 , 0.3456761 , 0.3378264 , 0.22834007,\n",
" 0.71616644, 0.62237275, 0.41121426],\n",
" [0.5875386 , 0.6119703 , 0.24063084, 0.52173215, 0.7950672 ,\n",
" 0.45715544, 0.3420424 , 0.8776624 , 0.31395352, 0.35489222,\n",
" 0.7077469 , 0.34783778, 0.49901104, 0.368652 , 0.41553447,\n",
" 0.7516321 , 0.29636133, 0.63707435, 0.45460576, 0.4504552 ,\n",
" 0.38551256, 0.3776133 , 0.37716988, 0.5808645 , 0.8186358 ,\n",
" 0.7452102 , 0.13270728, 0.32439175],\n",
" [0.21365675, 0.34269702, 0.36575893, 0.20755123, 0.73255384,\n",
" 0.6269738 , 0.5145893 , 0.40890205, 0.29832682, 0.16856807,\n",
" 0.521924 , 0.635281 , 0.5942463 , 0.66077226, 0.8809012 ,\n",
" 0.73497194, 0.7199179 , 0.3170066 , 0.28845012, 0.46699166,\n",
" 0.44126272, 0.6312956 , 0.35014617, 0.9267093 , 0.4469661 ,\n",
" 0.5532285 , 0.33181012, 0.36373758],\n",
" [0.708061 , 0.39687523, 0.22090068, 0.23750389, 0.56849474,\n",
" 0.25968918, 0.46802616, 0.2708095 , 0.23563243, 0.4951679 ,\n",
" 0.49091882, 0.69873273, 0.46350205, 0.7973521 , 0.5637764 ,\n",
" 0.48444656, 0.56311876, 0.4330717 , 0.34413686, 0.57832545,\n",
" 0.3900614 , 0.6826291 , 0.70586395, 0.55363905, 0.72460383,\n",
" 0.3997658 , 0.42758685, 0.13921028],\n",
" [0.42671797, 0.60539824, 0.61464536, 0.48930198, 0.7063285 ,\n",
" 0.6610134 , 0.1488111 , 0.52273536, 0.33709767, 0.27501753,\n",
" 0.4012898 , 0.28568238, 0.76208776, 0.6255617 , 0.54422826,\n",
" 0.6596593 , 0.21829887, 0.42483532, 0.22633922, 0.6994735 ,\n",
" 0.39678702, 0.6585098 , 0.5623585 , 0.26695958, 0.27654496,\n",
" 0.6461166 , 0.37591115, 0.39935783],\n",
" [0.22347261, 0.68466955, 0.69710433, 0.51150954, 0.65092474,\n",
" 0.7845727 , 0.54139376, 0.6863614 , 0.4282503 , 0.7025439 ,\n",
" 0.70491654, 0.4540074 , 0.18350555, 0.55251014, 0.19137311,\n",
" 0.43920127, 0.301945 , 0.43204868, 0.7097841 , 0.3531679 ,\n",
" 0.28080267, 0.5700798 , 0.56655395, 0.4369979 , 0.59634125,\n",
" 0.27452317, 0.5344645 , 0.79178244],\n",
" [0.4724254 , 0.2711047 , 0.94357526, 0.34359348, 0.54977125,\n",
" 0.58226395, 0.76491785, 0.7780358 , 0.762608 , 0.591783 ,\n",
" 0.27497044, 0.38095248, 0.22212288, 0.61650527, 0.4088468 ,\n",
" 0.35023472, 0.28026333, 0.20044121, 0.542197 , 0.47522944,\n",
" 0.547221 , 0.6610136 , 0.48738474, 0.22961178, 0.66614205,\n",
" 0.5408737 , 0.28099653, 0.64303696],\n",
" [0.66587275, 0.1830825 , 0.44034603, 0.7324861 , 0.44818747,\n",
" 0.7674686 , 0.34333137, 0.6313791 , 0.3968163 , 0.51550907,\n",
" 0.2743706 , 0.4589207 , 0.64085543, 0.47129798, 0.5682373 ,\n",
" 0.7182297 , 0.60502446, 0.6077174 , 0.37393433, 0.291496 ,\n",
" 0.41727996, 0.2796298 , 0.77829576, 0.460023 , 0.7404333 ,\n",
" 0.6846879 , 0.6370108 , 0.3651577 ],\n",
" [0.69297 , 0.61812246, 0.81404895, 0.271702 , 0.66017336,\n",
" 0.66111064, 0.31377548, 0.37802222, 0.39282626, 0.12249988,\n",
" 0.74444157, 0.38005698, 0.57431 , 0.31196997, 0.6235913 ,\n",
" 0.49954063, 0.76766944, 0.2101823 , 0.77645534, 0.9046659 ,\n",
" 0.14168741, 0.4651413 , 0.769572 , 0.58584106, 0.40810975,\n",
" 0.757797 , 0.5476642 , 0.3949979 ]]]], dtype=float32)>"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# using the original torch function\n",
"kornia.color.rgb_to_grayscale(torch.rand(1, 3, 28, 28))\n",
"\n",
"# using the transpiled tensorflow function\n",
"tf_rgb_to_grayscale(tf.random.uniform((1, 3, 28, 28)))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, lets check that the outputs of both the original torch function and transpiled tensorflow are identical when given the same inputs."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"True"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"torch_x = torch.rand(1, 3, 28, 28)\n",
"tf_x = tf.convert_to_tensor(torch_x.numpy())\n",
"\n",
"torch_out = kornia.color.rgb_to_grayscale(torch_x)\n",
"tf_out = tf_rgb_to_grayscale(tf_x)\n",
"\n",
"np.allclose(torch_out.numpy(), tf_out.numpy(), atol=1e-6)"
]
}
],
"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.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

0 comments on commit 449c40b

Please sign in to comment.