I am currently implementing the transformer architecture for sequence to sequence problems. Key part of the model is the attention mechanism, which is basically a matrix multiplication, followed by a masking operation and a softmax function. My initial thought was to wrap this 3 steps in a function, that looks like this:

    def attention(self, matrix_1, matrix_2, mask=None, trans_1=False, trans_2=False):
        att_stage_1 = F.matmul(matrix_1, matrix_2, transa=trans_1, transb=trans_2)*self.scale_score
        att_stage_2 = F.where(mask, att_stage_1, self.np.ones(att_stage_1.shape, 'f')*(-1e9))
        return F.softmax(att_stage_2, axis=3)

I want to write unit tests for this function to test whether the output is what I expect it to be. The problem, however, is that this function, as it is, performs 3 separate operations: matmul, masking and softmax. I would prefer to determine that each of this operations does produces correct output, but as it is I could only check the final effect. This leads me to a design where I would wrap each of this 3 operations to a separate, dedicated function and test them separately. What I am concerned, however, is that the overhead of python functions calls in a training loop function that is called on each forward pass may be unnecessary.

Thus, the question is, what would be the correct approach to balance design and reliability vs performance in this scenario? Maybe I am missing some obvious approach here.

2 Answers 2


What is the correct approach to balance design and reliability vs performance?

Your question makes very little sense.

  • Performance requires a design that supports performance. Otherwise it merely appears fast, because it happens to be faster than what was necessary/desired. Like tumbling down a hill. It is fast in comparison to walking. Yet a car, and a road can probably get you down the hill much faster.
  • Performance is useless unless it is reliable. Otherwise its only being performant at doing the wrong thing. Again like tumbling down a hill, its fast but it is probably doing the wrong things (like breaking bones). While a car can do it without the downsides of tumbling, and still allow you to descend that hill again.

I'm going to go out on a limb here and rephrase your question.

  • Should I decompose the function to support Unit Testing?
  • Will the decomposition incur an intolerable runtime cost?

Decomposition and Unit Tests

Will decomposing the function, in comparison to the original function:

  • Reduce the amount of test setup needed?
  • Reduce the overall number of tests?
  • Make the tests shorter and easier to understand?
  • Illuminate any meaningful duplication of behaviour?
  • Help explain how the algorithm works?

As a general rule of thumb, avoid making work.

  • If the decomposition/s that you can think of satisfy the above, then refactor. You will spend a lot of time rereading code, make it easy on your future self.

  • If for whatever reason the decompositions make life harder then you have already found a good level of analysis for the problem.

Runtime Costs

Python is an interpreted language. Some interpreters will compile optimised byte/native codes for common execution paths. Given that this function is a dataflow algorithm with no branching, the optimiser will have a field day here. It will notice that each function is always being called, in a particular order and generally it will optimise out the function call, among other enhancements.

That being said Python is an interpreted language, so you will need to consult the documentation on the interpreter to learn how it would handle this. Regardless of what the documentation says though, nothing speaks louder than data about how it really does do it. So collect statistics - which is good technique for any performance question about any language/platform.

  1. Write the code both ways.
  2. Execute it over a large set of varied data.
  3. Measure the amount of time it takes to completely process all of the data.
  4. Repeat the measurement several times (at least 5 per set/implementation) to get a feel for variability.
  5. Compare them.
  6. Consider the trade off between Developer concerns and Operational concerns - code that is a pleasure to work with vs. Code that meets the necessary/desired performance and reliability.
  7. Make your decision.
  • CPython – the only Python implementation that is suitable for OP's scenario – performs zero optimizations. It is a dumb bytecode interpreter. The only optimization it is capable of is deactivating assertions with a CLI flag. No JIT, no inlining. That means optimizations such as inlining do have to be performed in the source code. I've done that to tickle more performance out of hot spots, but I used a profiler to find them rather than trying to preemptively avoid function calls. For ML, the bottleneck is likely to be in native libraries do that Python micro-optimizations would be useless.
    – amon
    Apr 18, 2019 at 7:00
  • Thanks for the answer and taking time to rephrase my question! I totally agree with the second point, that performance is useless unless it is reliable and it is precisely what I wanted to address here. I guess the other approach would be a bottom-up one, when I would first write tests and functions for every step and then refactor if that increases performance, still having a reliable function. It is for the sake of being able and confident to refactor it to, for example, reuse some pieces in a new architecture. Have it pleasant to work with as you mentioned.
    – dkwasny
    Apr 18, 2019 at 8:00
  • Seconding @amon on this. I would be shocked if function call overhead even showed up on a profiler next to the matrix multiplication. Apr 18, 2019 at 12:09

You are correct that Python function calls are comparatively expensive, but you are probably not right to conclude that you should avoid function calls here. Numpy itself typically consists of a couple of Python layers before array operations are performed by native code. That means you are already using Python function calls, and one more level likely won't matter.

If you really want to create a testable design without extra function calls, consider returning a tuple with the results of each stage. That might however prevent optimizations such as modifying an array in-place.

As a QA technique, opinions are split on whether such detail tests are good tests. I'd argue that such tests do help build confidence in the software, and are therefore beneficial – even if such tests can be very fragile.

For a test with the purpose of verifying your understanding of the math rather than primarily verifying the behaviour of the software, it might also be feasible to re-implement the stages of the software in the test. This would give you insight into the stages without runtime overhead, but result in more fragile tests:

def test_stages():
  # verify stage 1
  result1 = ... something with data ...
  assert result1 == expected1

  # verify stage 2
  result2 = ... something with result1 ...
  assert result2 == expected2

  # verify final stage
  result = ... something with result2 ...
  assert result == expected_final

  # compare with real implementation
  actual = attention(data)
  assert actual == expected_final
  • Thanks for the advice! Yes, one of the purpose is the confidence that math is correct, I am still inexperienced and looking for the optimal approach for me. The second thing is that sometimes I see AI research code being very messy, not reusable to the minimal extent and without any tests. I think this is a situation in which it is extremely easy to introduce some bugs through copy pasting, and also the copy-pasting approach is a lot of overhead if your model differs by some small implementation that could be easily configurable if the piece of code was written more modular.
    – dkwasny
    Apr 18, 2019 at 8:06

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.