lzhengning
/
JittorMirror
mirror of https://github.com/Jittor/Jittor
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

test numpy code op

This commit is contained in:
Gword 2020-07-03 11:44:19 +08:00
parent cc20479fd1
commit fae73d9ca1
4 changed files with 134 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
python/jittor/test/test_numpy_code_op.py Normal file
View File

@ -0,0 +1,47 @@
# ***************************************************************
# Copyright (c) 2020 Jittor. Authors: 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.
# ***************************************************************
import unittest
import jittor as jt
import numpy as np
class TestCodeOp(unittest.TestCase):
def forward_code(np, data):
a,b = data["inputs"]
c,d = data["outputs"]
np.add(a,b,out=c)
np.substract(a,b,out=d)
def backward_code1(np, data):
dout = data["dout"]
da, db = data["outputs"]
np.copyto(dout, da)
np.copyto(dout, db)
def backward_code2(np, data):
dout = data["dout"]
da, db = data["outputs"]
np.copyto(dout, da)
np.negtive(dout, db)
def test(self):
a = jt.random((5,1))
b = jt.random((5,1))
c, d = jt.numpy_code(
[a.shape, a.shape],
[a.dtype, a.dtype],
[a, b],
self.forward_code,
[self.backward_code1,self.backward_code2],
)
print("a:",a)
print("b:",b)
print("a+b:",c)
print("a-b:",d)
if __name__ == "__main__":
unittest.main()

75
src/ops/numpy_code_op.cc
View File

@ -13,7 +13,7 @@
namespace jittor {
static auto make_numpy_code = get_op_info("numpy_code")
.get_constructor<VarPtr, NanoVector, NanoString, vector<Var*>&&, NumpyFunc&&, NumpyResults&&>();
.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";
@ -22,8 +22,8 @@ static inline void check_vary_shape(NanoVector v) {
<< "Vary shape should only occur in the first dimension:" << v;
}
NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs={}, NumpyFunc&& forward, vector<NumpyFunc>&& backward)
: _inputs(inputs), forward(forward),backward(backward)
NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs, NumpyFunc&& forward, vector<NumpyFunc>&& sbackward)
: _inputs(inputs), forward(move(forward))
{
_outputs.push_back(create_output(shape, dtype));
CHECKop(_inputs.size(),<=,10);
@ -32,10 +32,13 @@ NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inpu
flags.set(NodeFlags::_vary_shape);
check_vary_shape(_outputs[0]->shape);
}
for (int i=0; i<sbackward.size(); i++) {
backward.push_back(move(sbackward[i]));
}
}
NumpyCodeOp::NumpyCodeOp(vector<NanoVector>&& shapes, vector<NanoString>&& dtypes, vector<Var*>&& inputs={}, NumpyFunc&& forward, vector<NumpyFunc>&& backward)
: _inputs(inputs), forward(forward),backward(backward)
NumpyCodeOp::NumpyCodeOp(vector<NanoVector>&& shapes, vector<NanoString>&& dtypes, vector<Var*>&& inputs, NumpyFunc&& forward, vector<NumpyFunc>&& sbackward)
: _inputs(inputs), forward(move(forward))
{
CHECKop(shapes.size(),==,dtypes.size()) << "Number of outputs' shapes and dtypes should be the same";
_outputs.resize(shapes.size());
@ -49,10 +52,13 @@ NumpyCodeOp::NumpyCodeOp(vector<NanoVector>&& shapes, vector<NanoString>&& dtype
check_vary_shape(_outputs[i]->shape);
}
}
for (int i=0; i<sbackward.size(); i++) {
backward.push_back(move(sbackward[i]));
}
}
NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs={}, NumpyFunc&& forward, NumpyResults&& results)
: _inputs(inputs), forward(forward), _results(results)
NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs, NumpyFunc&& forward, NumpyResult&& results)
: _inputs(inputs), forward(move(forward)), _results(move(results))
{
_outputs.push_back(create_output(shape, dtype));
CHECKop(_inputs.size(),<=,10);
@ -63,17 +69,23 @@ NumpyCodeOp::NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inpu
}
}
void NumpyCodeOp::grad(Var* out, Var* dout, Var* v, int v_index) {
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);
result.arrays["dout"] = ArrayArgs{
dout->ptr,
dout->shape,
dout->dtype(),
};
// result.ints["dout_index"] = _outputs.find(out);
for (int i=0; i<_outputs.size(); i++) {
if (_outputs[i] == out) {
result.ints["dout_index"] = i;
break;
}
}
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(
@ -87,14 +99,13 @@ void NumpyCodeOp::grad(Var* out, Var* dout, Var* v, int v_index) {
return make_numpy_code(
_inputs[v_index]->shape,
_inputs[v_index]->dtype(),
inputs,
backward[v_index],
result,
)
move(inputs),
move(backward[v_index]),
move(result));
}
void NumpyCodeOp::run() {
NumpyResult result=_results;
NumpyResult result=move(_results);
vector<ArrayArgs> inputs(_inputs.size());
vector<ArrayArgs> outputs(_outputs.size());
/*
@ -102,21 +113,19 @@ void NumpyCodeOp::run() {
NanoVector shape;
NanoString dtype;
*/
for (int i=0; i<inputs.size(); i++)
inputs[i] = ArrayArgs{
_inputs[i]->ptr,
_inputs[i]->shape,
_inputs[i]->dtype(),
};
for (int i=0; i<outputs.size(); i++)
outputs[i] = ArrayArgs{
_outputs[i]->ptr,
_outputs[i]->shape,
_outputs[i]->dtype(),
};
for (int i=0; i<inputs.size(); i++) {
inputs[i].ptr=_inputs[i]->ptr<ArrayArgs>();
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].shape=_outputs[i]->shape;
outputs[i].dtype=_outputs[i]->dtype();
}
result.varrays["inputs"] = move(inputs);
result.varrays["outputs"] = move(outputs);
forward.callback(&results);
forward.callback(&result);
}
} // jittor

7
src/ops/numpy_code_op.h
View File

@ -5,6 +5,7 @@
// ***************************************************************
#pragma once
#include "op.h"
#include "numpy_func.h"
namespace jittor {
@ -15,13 +16,13 @@ struct NumpyCodeOp : Op {
vector<NumpyFunc> backward;
NumpyResult _results;
NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs={}, NumpyFunc&& forward, vector<NumpyFunc>&& backward);
NumpyCodeOp(NanoVector shape, NanoString dtype, vector<Var*>&& inputs, NumpyFunc&& forward, vector<NumpyFunc>&& backward);
// @attrs(multiple_outputs)
NumpyCodeOp(vector<NanoVector>&& shapes, vector<NanoString>&& dtypes, vector<Var*>&& inputs={}, NumpyFunc&& forward, vector<NumpyFunc>&& backward);
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, NumpyResults&& 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;

77
src/pyjt/py_converter.h
View File

@ -67,7 +67,7 @@ DEF_IS(int, bool) is_type(PyObject* obj) {
return PyLong_CheckExact(obj);
}
DEF_IS(int, PyObject*) to_py_object(const T& a) {
DEF_IS(int, PyObject*) to_py_object(const int& a) {
return PyLong_FromLong(a);
}
@ -374,6 +374,14 @@ DEF_IS(DataView, PyObject*) to_py_object(T a) {
return oh.release();
}
struct NumpyFunc;
DEF_IS(NumpyFunc, bool) is_type(PyObject* obj) {
return PyCallable_Check(obj);
}
DEF_IS(NumpyFunc, T) from_py_object(PyObject* obj);
#define CHECK_IS_1(check_type) \
template<typename T> struct is_##check_type : public std::false_type {}; \
template<typename T> \
@ -455,41 +463,6 @@ DEF_IS(FetchFunc, T) from_py_object(PyObject* obj) {
return func;
}
struct NumpyFunc;
DEF_IS(NumpyFunc, bool) is_type(PyObject* obj) {
return PyCallable_Check(obj);
}
DEF_IS(NumpyFunc, T) from_py_object(PyObject* obj) {
// PyObject_Call
Py_INCREF(obj);
T func(
// callback
[obj](typename T::R* result) {
// import numpy
PyObjHolder np(PyImport_ImportModule("numpy"));
// data = {}
PyObjHolder data(to_py_object(results->varrays));
PyObjHolder data2(to_py_object(results->ints));
PyObjHolder data3(to_py_object(results->arrays));
// data.update(data2)
PyDict_Update(data.obj, data2.obj);
// data.update(data3)
PyDict_Update(data.obj, data3.obj);
// args = []
PyObjHolder args(PyList_new());
auto ok = PyList_Append(args.obj, np.obj);
ASSERT(ok);
auto ok = PyList_Append(args.obj, data.obj);
ASSERT(ok);
PyObjHolder ret(PyObject_Call(obj, args.obj/* PyObject* */, nullptr));
},
// deleter
[obj]() { Py_DECREF(obj); }
);
return func;
}
#define CHECK_IS_2(check_type) \
template<typename T> struct is_##check_type : public std::false_type {}; \
@ -583,4 +556,36 @@ DEF_IS_1(fast_shared_ptr, T) from_py_object(PyObject* obj) {
}
DEF_IS(NumpyFunc, T) from_py_object(PyObject* obj) {
// PyObject_Call
Py_INCREF(obj);
T func(
// callback
[obj](typename T::R* result) {
// 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 ret(PyObject_Call(obj, args.obj, nullptr));
},
// deleter
[obj]() { Py_DECREF(obj); }
);
return func;
}
} // jittor