[wasm] Fix AvgPool and MaxPool for 1x1 kernels #6969
Conversation
Create identity_pool_test with a function to test similar ops against 1x1 kernels and use it to test MaxPool. Fix MaxPool in the wasm backend.
tfjs-backend-wasm/src/cc/util.h
Outdated
| @@ -152,6 +152,15 @@ const std::vector<size_t> assert_and_get_broadcast_shape( | |||
| const std::vector<size_t> get_broadcast_dims( | |||
| const std::vector<size_t> in_shape, const std::vector<size_t> out_shape); | |||
|
|
|||
| // Generates the output for AvgPool, MaxPool, etc where xnnpack does not support | |||
| // a 1x1 filter. Applies batching, channels, and strides. | |||
| // TODO(mattsoulanille): Support padding. | |||
There was a problem hiding this comment.
Could we either add an error in identity_pool or add a condition in the references for cases padding != 0? Otherwise, this will silently produce a wrong result.
There was a problem hiding this comment.
I did some more reading on AvgPool padding, and it seems like it's meaningless for the 1x1 kernel case.
Unlike for convolution ops, padding for AvgPool and MaxPool does not take into account the padded values when computing the result at a given index. It only uses the values that are part of the original tensor (and padding just allows it to compute at a given index instead of returning a tensor of smaller dimension than the input). Pooling with a 1x1 kernel will always output the same dimension as the input, so padding should not affect the result.
I also think I discovered a bug in the CPU backend; Padding a 1x1 kernel can result in NaNs in the output tensor.
As far as I can tell right now, the wasm version is actually correct to completely ignore the padding (for 1x1), so I'll remove the TODO.
XNNPack does not support 1x1 kernels for AvgPool or MaxPool. Implement these cases manually, including support for strides.
Support for padding is still TODO.
#6867 could be fixed by this issue, but I've yet to confirm that.
To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.
This change is