[resolved] How to implement k-max pooling for CNN text classification?

#1

A possible implmentation is

def kmax(self, x, k):
    return x.sort(dim = 3)[0][:, :, :, -k:]

However, this cannot keep the information of relative position.

#2

If you remove the [0] part you’ll get a tuple of (sorted_values, sorted_indices), so it should be possible to get the position. Also, you might want to use topk instead of sort.

#3

Could you provide some demonstration code? I cannot find any differential API for such kind of index selection. index_select seems not suitable for this. @apaszke

#4

Then I think I don’t understand the problem. What’s the exact formula for k-pooling and what’s the problem with your implementation?

#5

I found the exact solution. The key API is torch.gather:

import torch
def kmax_pooling(x, dim, k):
    index = x.topk(k, dim = dim)[1].sort(dim = dim)[0]
    return x.gather(dim, index)
x = torch.rand(4, 5, 6, 10)
y = kmax_pooling(x, 3, 5)
print(x[0, 0])
print(y[0, 0])

Output:

 0.2762  0.3788  0.5708  0.3251  0.0568  0.2483  0.3930  0.1249  0.1874  0.1113
 0.9230  0.7428  0.0957  0.2301  0.6187  0.8898  0.3007  0.2653  0.5313  0.1032
 0.6376  0.9639  0.6584  0.1502  0.0250  0.5792  0.9283  0.1783  0.9545  0.1681
 0.8456  0.6135  0.2860  0.9366  0.5178  0.0113  0.4864  0.9308  0.3005  0.5403
 0.3280  0.8755  0.2290  0.0899  0.9093  0.6971  0.1557  0.2412  0.7991  0.9169
 0.5389  0.4603  0.7291  0.4070  0.0113  0.3571  0.3860  0.3354  0.4081  0.0209
[torch.FloatTensor of size 6x10]


 0.2762  0.3788  0.5708  0.3251  0.3930
 0.9230  0.7428  0.6187  0.8898  0.5313
 0.6376  0.9639  0.6584  0.9283  0.9545
 0.8456  0.6135  0.9366  0.9308  0.5403
 0.8755  0.9093  0.6971  0.7991  0.9169
 0.5389  0.4603  0.7291  0.4070  0.4081
[torch.FloatTensor of size 6x5]

Can also be used on autograd.Variable.

8 Likes