Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update sample Colab notebooks to use tf-nightly. #14901

Merged
merged 1 commit into from
Sep 1, 2023
+504 −351
Merged

Update sample Colab notebooks to use tf-nightly. #14901

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 65 additions & 50 deletions samples/colab/edge_detection.ipynb
View file
Open in desktop
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
Loading