graph_max_pooling.py 文件源码

def backward_gpu(self, inputs, grad_outputs):
        x = inputs[0]
        xp = cuda.get_array_module(x)
        n_batch, c_in, N = x.shape

        with cuda.get_device(x.data):
            x = x.transpose((2, 1, 0))
            gy = grad_outputs[0]
            N_coarse = gy.shape[2]
            gy = gy.transpose((2, 1, 0))
            gx = xp.zeros((N, c_in, n_batch), dtype=x.dtype)
            gpu_graphpool_bwd(N, N_coarse, self.max_inds, gy, gx)
            gx = gx.transpose((2, 1, 0))
        return gx,