numpy code op test

This commit is contained in:
Gword 2020-07-06 16:51:21 +08:00
parent ff6c59f385
commit dfa1a4d999
6 changed files with 112 additions and 89 deletions

View File

@ -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()

View File

@ -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);
}
}
}

View File

@ -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
};

View File

@ -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();
}

View File

@ -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;

View File

@ -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;
}