mirror of https://github.com/Jittor/Jittor
numpy code op test
This commit is contained in:
parent
ff6c59f385
commit
dfa1a4d999
|
@ -1,5 +1,8 @@
|
|||
# ***************************************************************
|
||||
# Copyright (c) 2020 Jittor. Authors: Dun Liang <randonlang@gmail.com>. All Rights Reserved.
|
||||
# Copyright (c) 2020 Jittor. Authors:
|
||||
# Guowei Yang <471184555@qq.com>
|
||||
# Dun Liang <randonlang@gmail.com>.
|
||||
# All Rights Reserved.
|
||||
# This file is subject to the terms and conditions defined in
|
||||
# file 'LICENSE.txt', which is part of this source code package.
|
||||
# ***************************************************************
|
||||
|
@ -8,23 +11,27 @@ import jittor as jt
|
|||
import numpy as np
|
||||
|
||||
class TestCodeOp(unittest.TestCase):
|
||||
def forward_code(np, data):
|
||||
def forward_code(self, np, data):
|
||||
a,b = data["inputs"]
|
||||
c,d = data["outputs"]
|
||||
np.add(a,b,out=c)
|
||||
np.substract(a,b,out=d)
|
||||
np.subtract(a,b,out=d)
|
||||
p, r = c.__array_interface__['data']
|
||||
|
||||
def backward_code1(np, data):
|
||||
def backward_code1(self, np, data):
|
||||
dout = data["dout"]
|
||||
da, db = data["outputs"]
|
||||
np.copyto(dout, da)
|
||||
np.copyto(dout, db)
|
||||
a,b,dout = data["inputs"]
|
||||
out = data["outputs"][0]
|
||||
np.copyto(out, dout)
|
||||
|
||||
def backward_code2(np, data):
|
||||
def backward_code2(self, np, data):
|
||||
dout = data["dout"]
|
||||
da, db = data["outputs"]
|
||||
np.copyto(dout, da)
|
||||
np.negtive(dout, db)
|
||||
out_index = data["out_index"]
|
||||
out = data["outputs"][0]
|
||||
if out_index==0:
|
||||
np.copyto(out, dout)
|
||||
else:
|
||||
np.negative(dout, out)
|
||||
|
||||
def test(self):
|
||||
a = jt.random((5,1))
|
||||
|
@ -38,10 +45,16 @@ class TestCodeOp(unittest.TestCase):
|
|||
[self.backward_code1,self.backward_code2],
|
||||
)
|
||||
|
||||
print("a:",a)
|
||||
print("b:",b)
|
||||
print("a+b:",c)
|
||||
print("a-b:",d)
|
||||
assert np.allclose(c.data,(a+b).data)
|
||||
assert np.allclose(d.data,(a-b).data)
|
||||
dca, dcb = jt.grad(c,[a,b])
|
||||
dda, ddb = jt.grad(d,[a,b])
|
||||
one=np.ones(a.shape)
|
||||
mone=one*-1.0
|
||||
assert np.allclose(dca.data,one)
|
||||
assert np.allclose(dcb.data,one)
|
||||
assert np.allclose(dda.data,one)
|
||||
assert np.allclose(ddb.data,mone)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
44
src/grad.cc
44
src/grad.cc
|
@ -93,28 +93,30 @@ vector<VarPtr> grad(Var* loss, vector<Var*> targets) {
|
|||
auto index = it.index;
|
||||
if (op->tflag != nt) continue;
|
||||
// TODO: support two outputs backprop.
|
||||
Var* out = op->outputs().back();
|
||||
Var* dout = grads[out->custom_data];
|
||||
VarPtr dvar = make_grad(op, out, dout, var, index);
|
||||
registe_node_trace_grad(dvar.ptr, op, index);
|
||||
if (dvar)
|
||||
ASSERT(dvar->num==var->num && dvar->shape.size()==var->shape.size())
|
||||
<< "dvar" << dvar << "var" << var;
|
||||
if (!grad)
|
||||
grad = move(dvar);
|
||||
else if (dvar) {
|
||||
grad = make_binary(grad, dvar, ns_add);
|
||||
#ifdef PREVENT_LARGE_FUSED_OP
|
||||
gsum ++;
|
||||
if (gsum>=PREVENT_LARGE_FUSED_OP) {
|
||||
// TODO: this is a dirty fix for
|
||||
// stopping fuse lots of op together,
|
||||
// try to find a better solution
|
||||
grad->flags.set(NodeFlags::_stop_fuse);
|
||||
for (Var* out : op->outputs()) {
|
||||
if (out->tflag != nt) continue;
|
||||
Var* dout = grads[out->custom_data];
|
||||
VarPtr dvar = make_grad(op, out, dout, var, index);
|
||||
registe_node_trace_grad(dvar.ptr, op, index);
|
||||
if (dvar)
|
||||
ASSERT(dvar->num==var->num && dvar->shape.size()==var->shape.size())
|
||||
<< "dvar" << dvar << "var" << var;
|
||||
if (!grad)
|
||||
grad = move(dvar);
|
||||
else if (dvar) {
|
||||
grad = make_binary(grad, dvar, ns_add);
|
||||
#ifdef PREVENT_LARGE_FUSED_OP
|
||||
gsum ++;
|
||||
if (gsum>=PREVENT_LARGE_FUSED_OP) {
|
||||
// TODO: this is a dirty fix for
|
||||
// stopping fuse lots of op together,
|
||||
// try to find a better solution
|
||||
grad->flags.set(NodeFlags::_stop_fuse);
|
||||
}
|
||||
#endif
|
||||
assign_attrs(grad.ptr, var);
|
||||
registe_node_trace_grad(grad.ptr, var, index);
|
||||
}
|
||||
#endif
|
||||
assign_attrs(grad.ptr, var);
|
||||
registe_node_trace_grad(grad.ptr, var, index);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,5 +1,8 @@
|
|||
// ***************************************************************
|
||||
// Copyright (c) 2020 Jittor. Authors: Dun Liang <randonlang@gmail.com>. All Rights Reserved.
|
||||
// Copyright (c) 2020 Jittor. Authors:
|
||||
// Guowei Yang <471184555@qq.com>
|
||||
// Dun Liang <randonlang@gmail.com>.
|
||||
// All Rights Reserved.
|
||||
// This file is subject to the terms and conditions defined in
|
||||
// file 'LICENSE.txt', which is part of this source code package.
|
||||
// ***************************************************************
|
||||
|
@ -17,15 +20,21 @@ struct NumpyFunc {
|
|||
typedef NumpyResult R;
|
||||
std::function<void(R*)> callback;
|
||||
std::function<void()> deleter;
|
||||
std::function<void()> inc_ref;
|
||||
NumpyFunc() = default;
|
||||
NumpyFunc(NumpyFunc&& other) : callback(other.callback), deleter(other.deleter) {
|
||||
NumpyFunc(NumpyFunc&& other) : callback(other.callback), deleter(other.deleter), inc_ref(other.inc_ref) {
|
||||
other.callback = nullptr;
|
||||
other.deleter = nullptr;
|
||||
other.inc_ref = nullptr;
|
||||
};
|
||||
NumpyFunc(const NumpyFunc& other) : callback(other.callback), deleter(other.deleter), inc_ref(other.inc_ref) {
|
||||
inc_ref();
|
||||
};
|
||||
NumpyFunc(const NumpyFunc&) = delete;
|
||||
NumpyFunc(std::function<void(R*)>&& callback) : callback(move(callback)) {}
|
||||
NumpyFunc(std::function<void(R*)>&& callback, std::function<void()>&& deleter)
|
||||
: callback(move(callback)), deleter(move(deleter)) {};
|
||||
NumpyFunc(std::function<void(R*)>&& callback, std::function<void()>&& deleter, std::function<void()>&& inc_ref)
|
||||
: callback(move(callback)), deleter(move(deleter)), inc_ref(move(inc_ref)) {};
|
||||
~NumpyFunc() {
|
||||
if (deleter) {
|
||||
deleter();
|
||||
|
@ -36,9 +45,9 @@ struct NumpyFunc {
|
|||
|
||||
struct NumpyResult {
|
||||
// vector<Allocation> allocations;
|
||||
map<string, vector<ArrayArgs>> varrays;
|
||||
map<string, vector<DataView>> varrays;
|
||||
map<string, int> ints;
|
||||
map<string, ArrayArgs> arrays;
|
||||
map<string, DataView> arrays;
|
||||
// mem ptr, dtype, shape --> numpy array
|
||||
};
|
||||
|
||||
|
|
|
@ -1,5 +1,8 @@
|
|||
// ***************************************************************
|
||||
// Copyright (c) 2020 Jittor. Authors: Dun Liang <randonlang@gmail.com>. All Rights Reserved.
|
||||
// Copyright (c) 2020 Jittor. Authors:
|
||||
// Guowei Yang <471184555@qq.com>
|
||||
// Dun Liang <randonlang@gmail.com>.
|
||||
// All Rights Reserved.
|
||||
// This file is subject to the terms and conditions defined in
|
||||
// file 'LICENSE.txt', which is part of this source code package.
|
||||
// ***************************************************************
|
||||
|
@ -13,7 +16,7 @@
|
|||
namespace jittor {
|
||||
|
||||
static auto make_numpy_code = get_op_info("numpy_code")
|
||||
.get_constructor<VarPtr, NanoVector, NanoString, vector<Var*>&&, NumpyFunc&&, NumpyResult&&>();
|
||||
.get_constructor<VarPtr, NanoVector, NanoString, vector<Var*>&&, NumpyFunc, NumpyResult&&>();
|
||||
|
||||
static inline void check_vary_shape(NanoVector v) {
|
||||
ASSERT(v.size()) << "Vary shape should not be zero dimension";
|
||||
|
@ -33,7 +36,7 @@ NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inpu
|
|||
check_vary_shape(_outputs[0]->shape);
|
||||
}
|
||||
for (int i=0; i<sbackward.size(); i++) {
|
||||
backward.push_back(move(sbackward[i]));
|
||||
backward.push_back(sbackward[i]);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -53,12 +56,12 @@ NumpyCodeOp::NumpyCodeOp(vector<NanoVector>&& shapes, vector<NanoString>&& dtype
|
|||
}
|
||||
}
|
||||
for (int i=0; i<sbackward.size(); i++) {
|
||||
backward.push_back(move(sbackward[i]));
|
||||
backward.push_back(sbackward[i]);
|
||||
}
|
||||
}
|
||||
|
||||
NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs, NumpyFunc&& forward, NumpyResult&& results)
|
||||
: _inputs(inputs), forward(move(forward)), _results(move(results))
|
||||
NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs, NumpyFunc forward, NumpyResult&& results)
|
||||
: _inputs(inputs), forward(forward), _results(move(results))
|
||||
{
|
||||
_outputs.push_back(create_output(shape, dtype));
|
||||
CHECKop(_inputs.size(),<=,10);
|
||||
|
@ -71,55 +74,48 @@ NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inpu
|
|||
|
||||
VarPtr NumpyCodeOp::grad(Var* out, Var* dout, Var* v, int v_index) {
|
||||
NumpyResult result;
|
||||
// set results
|
||||
// set dout index
|
||||
// result.ints["dout_index"] = _outputs.find(out);
|
||||
|
||||
int out_index=-1;
|
||||
for (int i=0; i<_outputs.size(); i++) {
|
||||
if (_outputs[i] == out) {
|
||||
result.ints["dout_index"] = i;
|
||||
out_index = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
ASSERT(out_index!=-1);
|
||||
result.ints["out_index"] = out_index;
|
||||
result.arrays["dout"].ptr=dout;
|
||||
result.arrays["dout"].shape=dout->shape;
|
||||
result.arrays["dout"].dtype=dout->dtype();
|
||||
auto inputs = clone(_inputs);
|
||||
inputs.push_back(dout);
|
||||
|
||||
// code op:
|
||||
/*
|
||||
return make_code(
|
||||
_inputs[v_index]->shape,
|
||||
_inputs[v_index]->dtype(),
|
||||
move(inputs),
|
||||
move(cpu_src), {}, alias+cpu_header,
|
||||
move(cuda_src), {}, alias+cuda_header
|
||||
);
|
||||
*/
|
||||
return make_numpy_code(
|
||||
_inputs[v_index]->shape,
|
||||
_inputs[v_index]->dtype(),
|
||||
move(inputs),
|
||||
move(backward[v_index]),
|
||||
move(inputs),
|
||||
backward[v_index],
|
||||
move(result));
|
||||
}
|
||||
|
||||
void NumpyCodeOp::run() {
|
||||
NumpyResult result=move(_results);
|
||||
vector<ArrayArgs> inputs(_inputs.size());
|
||||
vector<ArrayArgs> outputs(_outputs.size());
|
||||
/*
|
||||
const void* ptr;
|
||||
NanoVector shape;
|
||||
NanoString dtype;
|
||||
*/
|
||||
NumpyResult result;
|
||||
result.varrays = _results.varrays;
|
||||
result.ints = _results.ints;
|
||||
result.arrays = _results.arrays;
|
||||
|
||||
if (result.arrays.count("dout") > 0){
|
||||
result.arrays["dout"].ptr=((Var*)result.arrays["dout"].ptr)->ptr<DataView>();
|
||||
}
|
||||
vector<DataView> inputs(_inputs.size());
|
||||
vector<DataView> outputs(_outputs.size());
|
||||
for (int i=0; i<inputs.size(); i++) {
|
||||
inputs[i].ptr=_inputs[i]->ptr<ArrayArgs>();
|
||||
inputs[i].ptr=_inputs[i]->ptr<DataView>();
|
||||
inputs[i].shape=_inputs[i]->shape;
|
||||
inputs[i].dtype=_inputs[i]->dtype();
|
||||
}
|
||||
for (int i=0; i<outputs.size(); i++) {
|
||||
outputs[i].ptr=_outputs[i]->ptr<ArrayArgs>();
|
||||
outputs[i].ptr=_outputs[i]->ptr<DataView>();
|
||||
outputs[i].shape=_outputs[i]->shape;
|
||||
outputs[i].dtype=_outputs[i]->dtype();
|
||||
}
|
||||
|
|
|
@ -1,5 +1,8 @@
|
|||
// ***************************************************************
|
||||
// Copyright (c) 2020 Jittor. Authors: Dun Liang <randonlang@gmail.com>. All Rights Reserved.
|
||||
// Copyright (c) 2020 Jittor. Authors:
|
||||
// Guowei Yang <471184555@qq.com>
|
||||
// Dun Liang <randonlang@gmail.com>.
|
||||
// All Rights Reserved.
|
||||
// This file is subject to the terms and conditions defined in
|
||||
// file 'LICENSE.txt', which is part of this source code package.
|
||||
// ***************************************************************
|
||||
|
@ -22,7 +25,7 @@ struct NumpyCodeOp : Op {
|
|||
NumpyCodeOp(vector<NanoVector>&& shapes, vector<NanoString>&& dtypes, vector<Var*>&& inputs, NumpyFunc&& forward, vector<NumpyFunc>&& backward);
|
||||
|
||||
// @pybind(None)
|
||||
NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs, NumpyFunc&& forward, NumpyResult&& results);
|
||||
NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs, NumpyFunc forward, NumpyResult&& results);
|
||||
|
||||
const char* name() const override { return "numpy_code"; }
|
||||
VarPtr grad(Var* out, Var* dout, Var* v, int v_index) override;
|
||||
|
|
|
@ -354,7 +354,6 @@ DEF_IS(VarHolder*, T) from_py_object(PyObject* obj, unique_ptr<VarHolder>& holde
|
|||
|
||||
struct DataView;
|
||||
DEF_IS(DataView, PyObject*) to_py_object(T a) {
|
||||
auto obj = GET_OBJ_FROM_RAW_PTR(a.vh);
|
||||
int64 dims[a.shape.size()];
|
||||
for (int i=0; i<a.shape.size(); i++)
|
||||
dims[i] = a.shape[i];
|
||||
|
@ -369,10 +368,13 @@ DEF_IS(DataView, PyObject*) to_py_object(T a) {
|
|||
NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_WRITEABLE, // flags
|
||||
NULL // obj
|
||||
));
|
||||
Py_INCREF(obj);
|
||||
PyObjHolder oh2(obj);
|
||||
ASSERT(PyArray_SetBaseObject(oh.obj, oh2.obj)==0);
|
||||
oh2.release();
|
||||
if (a.vh) {
|
||||
auto obj = GET_OBJ_FROM_RAW_PTR(a.vh);
|
||||
PyObjHolder oh2(obj);
|
||||
Py_INCREF(obj);
|
||||
ASSERT(PyArray_SetBaseObject(oh.obj, oh2.obj)==0);
|
||||
oh2.release();
|
||||
}
|
||||
return oh.release();
|
||||
}
|
||||
|
||||
|
@ -568,24 +570,22 @@ DEF_IS(NumpyFunc, T) from_py_object(PyObject* obj) {
|
|||
// import numpy
|
||||
PyObjHolder np(PyImport_ImportModule("numpy"));
|
||||
// data = {}
|
||||
//PyObjHolder data(to_py_object<map<string, vector<ArrayArgs>>>(result->varrays));
|
||||
PyObjHolder data(to_py_object(result->varrays));
|
||||
PyObjHolder data2(to_py_object(result->ints));
|
||||
PyObjHolder data3(to_py_object(result->arrays));
|
||||
// data.update(data2)
|
||||
PyDict_Update(data.obj, data2.obj);
|
||||
// data.update(data3)
|
||||
PyDict_Update(data.obj, data3.obj);
|
||||
|
||||
// args = []
|
||||
PyObjHolder args(PyList_New(0));
|
||||
int ok = PyList_Append(args.obj, np.obj);
|
||||
ASSERT(ok);
|
||||
ok = PyList_Append(args.obj, data.obj);
|
||||
ASSERT(ok);
|
||||
PyObjHolder args(PyTuple_New(2));
|
||||
PyTuple_SET_ITEM(args.obj, 0, np.obj);
|
||||
PyTuple_SET_ITEM(args.obj, 1, data.obj);
|
||||
PyObjHolder ret(PyObject_Call(obj, args.obj, nullptr));
|
||||
},
|
||||
// deleter
|
||||
[obj]() { Py_DECREF(obj); }
|
||||
[obj]() { Py_DECREF(obj); },
|
||||
// inc_ref
|
||||
[obj]() { Py_INCREF(obj); }
|
||||
);
|
||||
return func;
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue