Fixup

Author: qubvel

examples/convert_to_onnx.ipynb

@@ -52,7 +52,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model = smp.Unet('resnet34', encoder_weights='imagenet', classes=1)\n",
+    "model = smp.Unet(\"resnet34\", encoder_weights=\"imagenet\", classes=1)\n",
     "model = model.eval()"
    ]
   },
@@ -70,23 +70,23 @@
    "outputs": [],
    "source": [
     "# dynamic_axes is used to specify the variable length axes. it can be just batch size\n",
-    "dynamic_axes = {0: 'batch_size', 2: \"height\", 3: \"width\"}\n",
+    "dynamic_axes = {0: \"batch_size\", 2: \"height\", 3: \"width\"}\n",
     "\n",
-    "onnx_model_name = 'unet_resnet34.onnx'\n",
+    "onnx_model_name = \"unet_resnet34.onnx\"\n",
     "\n",
     "onnx_model = torch.onnx.export(\n",
-    "    model,                          # model being run\n",
-    "    torch.randn(1, 3, 224, 224),    # model input\n",
-    "    onnx_model_name,                # where to save the model (can be a file or file-like object) \n",
-    "    export_params=True,             # store the trained parameter weights inside the model file\n",
-    "    opset_version=17,               # the ONNX version to export\n",
-    "    do_constant_folding=True,       # whether to execute constant folding for optimization\n",
-    "    input_names=['input'],          # the model's input names\n",
-    "    output_names=['output'],        # the model's output names\n",
-    "    dynamic_axes={                  # variable length axes\n",
-    "        'input': dynamic_axes,\n",
-    "        'output': dynamic_axes\n",
-    "    }\n",
+    "    model,  # model being run\n",
+    "    torch.randn(1, 3, 224, 224),  # model input\n",
+    "    onnx_model_name,  # where to save the model (can be a file or file-like object)\n",
+    "    export_params=True,  # store the trained parameter weights inside the model file\n",
+    "    opset_version=17,  # the ONNX version to export\n",
+    "    do_constant_folding=True,  # whether to execute constant folding for optimization\n",
+    "    input_names=[\"input\"],  # the model's input names\n",
+    "    output_names=[\"output\"],  # the model's output names\n",
+    "    dynamic_axes={  # variable length axes\n",
+    "        \"input\": dynamic_axes,\n",
+    "        \"output\": dynamic_axes,\n",
+    "    },\n",
     ")"
    ]
   },
@@ -153,11 +153,13 @@
     }
    ],
    "source": [
-    "# create sample with different batch size, height and width \n",
+    "# create sample with different batch size, height and width\n",
     "# from what we used in export above\n",
-    "sample = torch.randn(2, 3, 512, 512)  \n",
+    "sample = torch.randn(2, 3, 512, 512)\n",
     "\n",
-    "ort_session = onnxruntime.InferenceSession(onnx_model_name, providers=[\"CPUExecutionProvider\"])\n",
+    "ort_session = onnxruntime.InferenceSession(\n",
+    "    onnx_model_name, providers=[\"CPUExecutionProvider\"]\n",
+    ")\n",
     "\n",
     "# compute ONNX Runtime output prediction\n",
     "ort_inputs = {\"input\": sample.numpy()}\n",