graph_max_pooling.py 文件源码

def forward_gpu(self, inputs):
        x = inputs[0]
        xp = cuda.get_array_module(x)
        n_batch, c, N = x.shape
        N_coarse = len(self.pooling_inds)

        with cuda.get_device(x.data):
            x = x.transpose((2, 1, 0))
            p_dim = self.pooling_inds.shape[1]
            y = xp.empty((N_coarse, c, n_batch), dtype=x.dtype)
            self.max_inds = xp.empty((N_coarse, c, n_batch), dtype=np.int32)
            pooling_inds = cuda.to_gpu(self.pooling_inds)
            gpu_graphpool_fwd(N_coarse, p_dim, pooling_inds,
                              x, y, self.max_inds)
            y = y.transpose((2, 1, 0))

        return y,