CODEWITHSUNDEEP
pytorch
Ink
Ink

Reputation: 963

What's the difference between torch.mean and torch.nn.avg_pool?

Taking a tensor with shape [4,8,12] as an example, what's the difference between the two lines:

torch.mean(x, dim=2)
torch.nn.functional.avg_pool1d(x, kernel_size=12)

Upvotes: 4

Views: 7856

Answers (1)

saiden
saiden

Reputation: 342

With the very example you provided the result is the same, but only because you specified dim=2 and kernel_size equal to the dimensionality of the third (index 2) dimension. But in principle, you are applying two different functions, that sometimes just happen to collide with specific choices of hyperparameters.

torch.mean is effectively a dimensionality reduction function, meaning that when you average all values across one dimension, you effectively get rid of that dimension. On the other hand, average 1-dimensional pooling is more powerful in this regard, as it gives you a lot more flexibility in choosing kernel size, padding and stride like you would normally do when using a convolutional layer.

You can see the first function as a specific case of 1-d pooling.

Upvotes: 4

Related Questions