René's URL Explorer Experiment

{"@context":"https://schema.org","@type":"DiscussionForumPosting","headline":"Disparate Performance between Python and Java","articleBody":"\u003cem\u003ePlease make sure that this is a bug. As per our [GitHub Policy](https://github.com/tensorflow/tensorflow/blob/master/ISSUES.md), we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template\u003c/em\u003e\n\n**System information**\n- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes\n- OS Platform and Distribution (e.g., Linux Ubuntu 16.04 x86\\_64): Ubuntu 24.04 x86\\_64\n- TensorFlow installed from (source or binary): binary\n- TensorFlow version (use command below): 1.0.0 \n- Java version (i.e., the output of `java -version`): openjdk version \"21.0.6\" 2025-01-21\n- Java command line flags (e.g., GC parameters): \n- Python version (if transferring a model trained in Python): 3.12.8\n- Bazel version (if compiling from source):\n- GCC/Compiler version (if compiling from source):\n- CUDA/cuDNN version: 12.8.61/8905\n- GPU model and memory: V100 (32GB)\n\n**Describe the current behavior**\n\nExecuting the exported model using Tensorflow in Python takes significantly less time than when calling the same function from using Tensorflow Java. I suspect that I am just not using the Java API correctly, because a small change to the python can lead to comparably poor performance in the python.\n\n**Describe the expected behavior**\n\nThe function calls should take a comparable amount of time.\n\n**Code to reproduce the issue**\n\nI have the following python function:\n\n```python\n@tf.function(\n    input_signature=[\n            tf.TensorSpec(shape=[41, 2048, 2048], dtype=tf.float32, name=\"data\"),  # [k, n, n]\n            tf.TensorSpec(shape=[1, 2048, 2048], dtype=tf.float32, name=\"image\"),  # [1, n, n]\n            tf.TensorSpec(shape=[41, 2048, 2048], dtype=tf.float32, name=\"psf\"),  # [k, n, n]\n    ],\n    jit_compile=True\n)\ndef rl_step(\n    data: tf.Tensor,  # [k, n, n]\n    image: tf.Tensor, # [1, n, n]\n    psf: tf.Tensor,   # [k, n, n]\n) -\u003e tf.Tensor: # [k, n, n]\n    psf_fft = tf.signal.rfft2d(psf)\n    psft_fft = tf.signal.rfft2d(tf.reverse(psf, axis=(-2, -1)))\n    denom = tf.reduce_sum(\n        tf.signal.irfft2d(psf_fft * tf.signal.rfft2d(data)),\n        axis=0,\n        keepdims=True\n    )\n    img_err = image / denom\n    return data * tf.signal.irfft2d(tf.signal.rfft2d(img_err) * psft_fft)\n```\n\nIn python, this function is applied iteratively over the same tensor as below:\n\n```python\n    image_tensor = tf.constant(image) # [k, n, n]\n    measured_psf_tensor = tf.constant(measured_psf) # [1, n, n]\n    data_tensor = tf.constant(data) # [k, n, n]\n\n    for i in range(10):\n        start = time()\n        data = rl_step(data_tensor, image_tensor, measured_psf_tensor)\n        print(f\"Iter {i}:\", time() - start, \"seconds.\")\n```\nHere `image`, `measured_psf`, and `data` are all 3D arrays with dtype=float32 and `n=2048` and `k=41`\n\nThis prints timings around the following:\n\n```bash\nIter 0: 0.2061774730682373 seconds.\nIter 1: 0.004193544387817383 seconds.\nIter 2: 0.0007469654083251953 seconds.\nIter 3: 0.000415802001953125 seconds.\nIter 4: 0.0004220008850097656 seconds.\nIter 5: 0.0004246234893798828 seconds.\nIter 6: 0.0004112720489501953 seconds.\nIter 7: 0.00042128562927246094 seconds.\nIter 8: 0.0004055500030517578 seconds.\nIter 9: 0.00040721893310546875 seconds.\n```\n\nI tried exporting the model by adding the following after the timing code:\n\n```python\n    mod = tf.Module()\n    mod.f = rl_step\n    tf.saved_model.save(mod, \"pure_tf_export\")\n```\n\nNow I tried to use this exported mode from the Java API,\n\n\n```java\n        String modelLocation = \"./pure_tf_export\";\n        try(Graph g = new Graph(); Session s = new Session(g)){\n            SavedModelBundle model = SavedModelBundle.loader(modelLocation).load();\n\n            try (Tensor imageTensor = TFloat32.tensorOf(image);\n                Tensor psfTensor = TFloat32.tensorOf(psf);\n                Tensor dataTensor = TFloat32.tensorOf(data)\n            ){\n                Map\u003cString, Tensor\u003e inputs = new HashMap\u003cString, Tensor\u003e();\n                inputs.put(\"data\", dataTensor);\n                inputs.put(\"image\", imageTensor);\n                inputs.put(\"psf\", psfTensor);\n\n                for (int i = 0; i \u003c 10; i++){\n\n                    Instant start = Instant.now();\n\n                    Result result = model.function(\"serving_default\").call(inputs);\n                    inputs.replace(\"data\", result.get(\"output_0\").get());\n\n                    System.out.println(\"Iter \" + i + \" \" + (Duration.between(start, Instant.now()).toMillis()/1000f) + \" seconds\");\n                }\n            }\n        }\n```\nAnd I get timings as follows:\n\n```bash\nIter 0 0.701 seconds\nIter 1 0.528 seconds\nIter 2 0.874 seconds\nIter 3 0.224 seconds\nIter 4 0.254 seconds\nIter 5 1.622 seconds\nIter 6 0.241 seconds\nIter 7 0.224 seconds\nIter 8 0.231 seconds\nIter 9 0.228 seconds\n```\n\nI am pretty sure I am making a simple mistake somewhere. I suspect it is in how I am instantiating the Tensors. I know in python if you don't use `tf.constant` the timings go up a lot.\n\nAny help would be very much appreciated. I tried looking through the documentation and the tensorflow java-examples repository, but couldn't spot what I am doing wrong.\n\nThanks again!","author":{"url":"https://github.com/ryanhausen","@type":"Person","name":"ryanhausen"},"datePublished":"2025-02-11T22:12:09.000Z","interactionStatistic":{"@type":"InteractionCounter","interactionType":"https://schema.org/CommentAction","userInteractionCount":7},"url":"https://github.com/602/java/issues/602"}