dmlc--dgl
xiang song(charlie.song)
0a56d65223
* upd
* fig edgebatch edges
* add test
* trigger
* Update README.md for pytorch PinSage example.
Add noting that the PinSage model example under
example/pytorch/recommendation only work with Python 3.6+
as its dataset loader depends on stanfordnlp package
which work only with Python 3.6+.
* Provid a frame agnostic API to test nn modules on both CPU and CUDA side.
1. make dgl.nn.xxx frame agnostic
2. make test.backend include dgl.nn modules
3. modify test_edge_softmax of test/mxnet/test_nn.py and
test/pytorch/test_nn.py work on both CPU and GPU
* Fix style
* Delete unused code
* Make agnostic test only related to tests/backend
1. clear all agnostic related code in dgl.nn
2. make test_graph_conv agnostic to cpu/gpu
* Fix code style
* fix
* doc
* Make all test code under tests.mxnet/pytorch.test_nn.py
work on both CPU and GPU.
* Fix syntex
* Remove rand
* Start implementing masked-mm kernel.
Add base control flow code.
* Add masked dot declare
* Update func/variable name
* Skeleton compile OK
* Update Implement. Unify BinaryDot with BinaryReduce
* New Impl of x_dot_x, reuse binary reduce template
* Compile OK.
TODO:
1. make sure x_add_x, x_sub_x, x_mul_x, x_div_x work
2. let x_dot_x work
3. make sure backward of x_add_x, x_sub_x, x_mul_x, x_div_x work
4. let x_dot_x backward work
* Fix code style
* Now we can pass the tests/compute/test_kernel.py for add/sub/mul/div forward and backward
* Fix mxnet test code
* Add u_dot_v, u_dot_e, v_dot_e unitest.
* Update doc
* Now also support v_dot_u, e_dot_u, e_dot_v
* Add unroll for some loop
* Add some Opt for cuda backward of dot builtin.
Backward is still slow for dot
* Apply UnravelRavel opt for broadcast backward
* update docstring
104 行
1.8 KiB
Python
104 行
1.8 KiB
Python
from __future__ import absolute_import
|
|
|
|
import numpy as np
|
|
import mxnet as mx
|
|
import mxnet.ndarray as nd
|
|
import mxnet.autograd as autograd
|
|
|
|
def cuda():
|
|
return mx.gpu()
|
|
|
|
def is_cuda_available():
|
|
# TODO: Does MXNet have a convenient function to test GPU availability/compilation?
|
|
try:
|
|
a = nd.array([1, 2, 3], ctx=mx.gpu())
|
|
return True
|
|
except mx.MXNetError:
|
|
return False
|
|
|
|
def array_equal(a, b):
|
|
return nd.equal(a, b).asnumpy().all()
|
|
|
|
def allclose(a, b, rtol=1e-4, atol=1e-4):
|
|
return np.allclose(a.asnumpy(), b.asnumpy(), rtol=rtol, atol=atol)
|
|
|
|
def randn(shape):
|
|
return nd.random.randn(*shape)
|
|
|
|
def attach_grad(x):
|
|
x.attach_grad()
|
|
return x
|
|
|
|
def backward(x, head_gradient=None):
|
|
x.backward(head_gradient)
|
|
|
|
def grad(x):
|
|
return x.grad
|
|
|
|
def is_no_grad(x):
|
|
return (x != 0).sum() == 0
|
|
|
|
def full(shape, fill_value, dtype, ctx):
|
|
return nd.full(shape, fill_value, dtype=dtype, ctx=ctx)
|
|
|
|
def narrow_row_set(x, start, stop, new):
|
|
x[start:stop] = new
|
|
|
|
def sparse_to_numpy(x):
|
|
return x.asscipy().todense().A
|
|
|
|
def clone(x):
|
|
return x.copy()
|
|
|
|
def reduce_sum(x):
|
|
return x.sum()
|
|
|
|
def softmax(x, dim):
|
|
return nd.softmax(x, axis=dim)
|
|
|
|
def spmm(x, y):
|
|
return nd.dot(x, y)
|
|
|
|
def add(a, b):
|
|
return a + b
|
|
|
|
def sub(a, b):
|
|
return a - b
|
|
|
|
def mul(a, b):
|
|
return a * b
|
|
|
|
def div(a, b):
|
|
return a / b
|
|
|
|
def sum(x, dim):
|
|
return x.sum(dim)
|
|
|
|
def max(x, dim):
|
|
return x.max(dim)
|
|
|
|
def min(x, dim):
|
|
return x.min(dim)
|
|
|
|
def prod(x, dim):
|
|
return x.prod(dim)
|
|
|
|
def matmul(a, b):
|
|
return nd.dot(a, b)
|
|
|
|
def dot(a, b):
|
|
return nd.sum(mul(a, b), axis=-1)
|
|
|
|
record_grad = autograd.record
|
|
|
|
|
|
class no_grad(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_traceback):
|
|
pass
|