Balancing function call overhead and testability in a code that is a part of the deep learning model training loop

Question

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.

Answer 1

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.

Answer 2

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