Skip to content

Commit

Permalink
Update sample Colab notebooks to use tf-nightly. (#14901)
Browse files Browse the repository at this point in the history 
This fixes most of our Colab notebooks that have been failing. I also
adapted several notebooks to API changes (tools output MLIR bytecode by
default, which can't be loaded as text and printed without going through
a converter).

I still want some sample notebooks for JAX and PyTorch... one day...

skip-ci: tested at
https://github.com/openxla/iree/actions/runs/6043263838
  • Loading branch information
@ScottTodd
ScottTodd authored Sep 1, 2023
1 parent 171e2d1 commit 326aca2
Show file tree
Hide file tree
Showing 7 changed files with 504 additions and 351 deletions.
115 changes: 65 additions & 50 deletions samples/colab/edge_detection.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,38 +56,42 @@
"## Setup"
]
},
{
"cell_type": "code",
"source": [
"%%capture\n",
"!python -m pip install --upgrade tf-nightly # Needed for stablehlo export in TF>=2.14"
],
"metadata": {
"id": "Nt3QaBP0s65F"
},
"execution_count": 2,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "k2SdyC52_QYa",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "f9054d95-2a2d-479f-f787-178fd0291026"
"outputId": "53aee7ec-e390-4293-f51f-2638b75148d3"
},
"source": [
"!python -m pip install iree-compiler iree-runtime iree-tools-tf -f https://openxla.github.io/iree/pip-release-links.html"
],
"execution_count": 2,
"execution_count": 3,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
"Looking in links: https://openxla.github.io/iree/pip-release-links.html\n",
"Collecting iree-compiler\n",
" Downloading https://github.com/openxla/iree/releases/download/candidate-20230426.501/iree_compiler-20230426.501-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (54.2 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.2/54.2 MB\u001b[0m \u001b[31m9.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hCollecting iree-runtime\n",
" Downloading https://github.com/openxla/iree/releases/download/candidate-20230426.501/iree_runtime-20230426.501-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.6 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.6/2.6 MB\u001b[0m \u001b[31m30.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hCollecting iree-tools-tf\n",
" Downloading https://github.com/openxla/iree/releases/download/candidate-20230426.501/iree_tools_tf-20230426.501-py3-none-linux_x86_64.whl (32 kB)\n",
"Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from iree-compiler) (1.22.4)\n",
"Requirement already satisfied: PyYAML in /usr/local/lib/python3.9/dist-packages (from iree-compiler) (6.0)\n",
"Installing collected packages: iree-tools-tf, iree-runtime, iree-compiler\n",
"Successfully installed iree-compiler-20230426.501 iree-runtime-20230426.501 iree-tools-tf-20230426.501\n"
"Requirement already satisfied: iree-compiler in /usr/local/lib/python3.10/dist-packages (20230831.630)\n",
"Requirement already satisfied: iree-runtime in /usr/local/lib/python3.10/dist-packages (20230831.630)\n",
"Requirement already satisfied: iree-tools-tf in /usr/local/lib/python3.10/dist-packages (20230831.630)\n",
"Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from iree-compiler) (1.23.5)\n",
"Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from iree-compiler) (6.0.1)\n"
]
}
]
Expand All @@ -99,7 +103,7 @@
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "d0d16ea9-7a0c-45cd-eb6f-36eb6baf170e"
"outputId": "8253ff7c-a1b8-44c6-f17c-b9d93dfe30cd"
},
"source": [
"#@title Imports\n",
Expand All @@ -119,13 +123,13 @@
"# Print version information for future notebook users to reference.\n",
"print(\"TensorFlow version: \", tf.__version__)"
],
"execution_count": 15,
"execution_count": 4,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"TensorFlow version: 2.12.0\n"
"TensorFlow version: 2.15.0-dev20230831\n"
]
}
]
Expand All @@ -142,7 +146,7 @@
"ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), \"iree\", \"colab_artifacts\")\n",
"os.makedirs(ARTIFACTS_DIR, exist_ok=True)"
],
"execution_count": 4,
"execution_count": 5,
"outputs": []
},
{
Expand All @@ -160,7 +164,7 @@
" sobel_x = tf.constant([[-1.0, 0.0, 1.0],\n",
" [-2.0, 0.0, 2.0],\n",
" [-1.0, 0.0, 1.0]],\n",
" dtype=tf.float32, shape=[3, 3, 1, 1]) \n",
" dtype=tf.float32, shape=[3, 3, 1, 1])\n",
" sobel_y = tf.constant([[ 1.0, 2.0, 1.0],\n",
" [ 0.0, 0.0, 0.0],\n",
" [-1.0, -2.0, -1.0]],\n",
Expand All @@ -171,7 +175,7 @@
"\n",
"tf_module = EdgeDetectionModule()"
],
"execution_count": 5,
"execution_count": 6,
"outputs": []
},
{
Expand All @@ -180,10 +184,10 @@
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 391
"height": 356
},
"id": "Q2no8DO_s125",
"outputId": "4eb80633-b2dc-4262-83cf-7794ba9a9c09"
"outputId": "2a427588-3a06-4f6b-b15d-76242b9643ae"
},
"source": [
"#@title Load a test image of a [labrador](https://commons.wikimedia.org/wiki/File:YellowLabradorLooking_new.jpg) and run the module with TF\n",
Expand Down Expand Up @@ -219,16 +223,8 @@
"# Plot the input and output images\n",
"show_images(image, tf_edges)"
],
"execution_count": 6,
"execution_count": 7,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Downloading data from https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg\n",
"83281/83281 [==============================] - 0s 0us/step\n"
]
},
{
"output_type": "display_data",
"data": {
Expand Down Expand Up @@ -262,7 +258,7 @@
"backend_choice = backend_choice.split(\" \")[0]\n",
"backend = module_utils.BackendInfo(backend_choice)"
],
"execution_count": 7,
"execution_count": 8,
"outputs": []
},
{
Expand All @@ -273,7 +269,7 @@
"height": 356
},
"id": "2OPBr0eGroE0",
"outputId": "85c453c7-81ba-4181-f872-c2e3cb8d3b9e"
"outputId": "128e76b9-9dd7-49b7-ee4f-0612c8e25bfd"
},
"source": [
"#@title Compile and Run the EdgeDetectionModule with IREE.\n",
Expand All @@ -283,7 +279,7 @@
"iree_edges = module.edge_detect_sobel_operator(image)\n",
"show_images(image, iree_edges)"
],
"execution_count": 8,
"execution_count": 9,
"outputs": [
{
"output_type": "display_data",
Expand Down Expand Up @@ -320,56 +316,66 @@
"base_uri": "https://localhost:8080/"
},
"id": "6YGqN2uqP_7P",
"outputId": "bb1c4be0-ca97-4fd1-8344-0f6d740a415d"
"outputId": "fe3f580c-bbac-48ec-bbbb-1029fa8edce3"
},
"source": [
"#@title Construct a module containing the edge detection function\n",
"\n",
"# Do *not* further compile to a bytecode module for a particular backend.\n",
"# \n",
"#\n",
"# By stopping at stablehlo, we can more easily take advantage of future\n",
"# compiler improvements within IREE. For a production application, we would\n",
"# probably want to freeze the version of IREE used and compile as completely as\n",
"# possible ahead of time, then use some other scheme to load the module into\n",
"# the application at runtime.\n",
"compiler_module = tfc.compile_module(\n",
" EdgeDetectionModule(), import_only=True)\n",
"print(\"Edge Detection MLIR: \", compiler_module)\n",
"\n",
"edge_detection_mlir_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.mlir\")\n",
"with open(edge_detection_mlir_path, \"wt\") as output_file:\n",
"# Save the imported MLIR to disk.\n",
"imported_mlirbc_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.mlirbc\")\n",
"with open(imported_mlirbc_path, \"wb\") as output_file:\n",
" output_file.write(compiler_module)\n",
"print(f\"Wrote MLIR to path '{edge_detection_mlir_path}'\")"
"print(f\"Wrote MLIR to path '{imported_mlirbc_path}'\")\n",
"\n",
"# Copy MLIR bytecode to MLIR text and see how the compiler views this program.\n",
"imported_mlir_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.mlir\")\n",
"!iree-ir-tool copy {imported_mlirbc_path} -o {imported_mlir_path}\n",
"print(\"Edge Detection MLIR:\")\n",
"!cat {imported_mlir_path}"
],
"execution_count": 11,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Edge Detection MLIR: module {\n",
"Wrote MLIR to path '/tmp/iree/colab_artifacts/edge_detection.mlirbc'\n",
"Edge Detection MLIR:\n",
"module {\n",
" func.func @edge_detect_sobel_operator(%arg0: tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> {\n",
" %0 = stablehlo.constant dense<[[[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]], [[[-2.000000e+00]], [[0.000000e+00]], [[2.000000e+00]]], [[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]]]> : tensor<3x3x1x1xf32>\n",
" %1 = stablehlo.constant dense<[[[[1.000000e+00]], [[2.000000e+00]], [[1.000000e+00]]], [[[0.000000e+00]], [[0.000000e+00]], [[0.000000e+00]]], [[[-1.000000e+00]], [[-2.000000e+00]], [[-1.000000e+00]]]]> : tensor<3x3x1x1xf32>\n",
" %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n",
" %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n",
" %3 = stablehlo.multiply %2, %2 : tensor<1x128x128x1xf32>\n",
" %4 = stablehlo.convolution(%arg0, %1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n",
" %4 = stablehlo.convolution(%arg0, %1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n",
" %5 = stablehlo.multiply %4, %4 : tensor<1x128x128x1xf32>\n",
" %6 = stablehlo.add %3, %5 : tensor<1x128x128x1xf32>\n",
" %7 = stablehlo.sqrt %6 : tensor<1x128x128x1xf32>\n",
" return %7 : tensor<1x128x128x1xf32>\n",
" }\n",
"}\n",
"\n",
"Wrote MLIR to path '/tmp/iree/colab_artifacts/edge_detection.mlir'\n"
"}"
]
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "Ytvb5Gx_EFJl"
"id": "Ytvb5Gx_EFJl",
"outputId": "b9f8452f-1eeb-4355-8b2c-09d71dc38ec5",
"colab": {
"base_uri": "https://localhost:8080/"
}
},
"source": [
"#@title Compile and prepare to test the edge detection module\n",
Expand All @@ -383,7 +389,16 @@
"ctx.add_vm_module(vm_module)"
],
"execution_count": 12,
"outputs": []
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
"<ipython-input-12-4d337c62f971>:8: UserWarning: Making copy of unaligned VmModule buffer. It is recommended to make this deterministic by calling `copy_buffer` to always make a copy or `mmap` to efficiently load from a file. This warning can be silenced by adding `warn_if_copy=False` to `from_buffer`\n",
" vm_module = ireert.VmModule.from_flatbuffer(ctx.instance, flatbuffer_blob)\n"
]
}
]
},
{
"cell_type": "code",
Expand All @@ -393,7 +408,7 @@
"height": 356
},
"id": "OUUXxol7wl-f",
"outputId": "b6b5a2b9-89b6-4574-e7aa-49803c28a2a4"
"outputId": "32b177e5-f3c1-4389-bf2e-69b4beaa5865"
},
"source": [
"edge_detect_sobel_operator_f = ctx.modules.module[\"edge_detect_sobel_operator\"]\n",
Expand Down
Loading

0 comments on commit 326aca2

Please sign in to comment.