Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
wmsofts
Paddle
提交
b294f054
P
Paddle
项目概览
wmsofts
/
Paddle
与 Fork 源项目一致
Fork自
PaddlePaddle / Paddle
通知
1
Star
0
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
P
Paddle
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
未验证
提交
b294f054
编写于
9月 15, 2022
作者:
W
wanghuancoder
提交者:
GitHub
9月 15, 2022
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
[Eager] saved_tensors_hooks (#45763)
* saved_tensors_hooks
上级
0c40d889
变更
12
隐藏空白更改
内联
并排
Showing
12 changed file
with
649 addition
and
9 deletion
+649
-9
paddle/fluid/eager/CMakeLists.txt
paddle/fluid/eager/CMakeLists.txt
+6
-1
paddle/fluid/eager/hooks.h
paddle/fluid/eager/hooks.h
+14
-0
paddle/fluid/eager/saved_tensors_hooks.cc
paddle/fluid/eager/saved_tensors_hooks.cc
+107
-0
paddle/fluid/eager/saved_tensors_hooks.h
paddle/fluid/eager/saved_tensors_hooks.h
+97
-0
paddle/fluid/eager/tensor_wrapper.h
paddle/fluid/eager/tensor_wrapper.h
+50
-2
paddle/fluid/pybind/CMakeLists.txt
paddle/fluid/pybind/CMakeLists.txt
+2
-1
paddle/fluid/pybind/eager.h
paddle/fluid/pybind/eager.h
+3
-0
paddle/fluid/pybind/eager_functions.cc
paddle/fluid/pybind/eager_functions.cc
+32
-0
paddle/fluid/pybind/eager_py_layer.cc
paddle/fluid/pybind/eager_py_layer.cc
+133
-5
python/paddle/autograd/__init__.py
python/paddle/autograd/__init__.py
+2
-0
python/paddle/autograd/saved_tensors_hooks.py
python/paddle/autograd/saved_tensors_hooks.py
+111
-0
python/paddle/fluid/tests/unittests/test_saved_tensors_hooks.py
.../paddle/fluid/tests/unittests/test_saved_tensors_hooks.py
+92
-0
未找到文件。
paddle/fluid/eager/CMakeLists.txt
浏览文件 @
b294f054
...
...
@@ -14,7 +14,8 @@ set(eager_deps
grad_node_info
grad_tensor_holder
accumulation_node
custom_operator_node
)
custom_operator_node
python
)
set
(
fluid_deps
tracer
...
...
@@ -77,6 +78,10 @@ cc_library(
autograd_meta
hook_utils
)
cc_library
(
saved_tensors_hooks
SRCS saved_tensors_hooks.cc
DEPS hook_utils
)
if
(
NOT
((
NOT WITH_PYTHON
)
AND ON_INFER
))
add_subdirectory
(
tests
)
endif
()
paddle/fluid/eager/hooks.h
浏览文件 @
b294f054
...
...
@@ -62,4 +62,18 @@ class CppVoidHook : public VoidHook {
std
::
function
<
void
()
>
fn_
;
};
class
PackHookBase
{
public:
virtual
~
PackHookBase
()
=
default
;
virtual
void
*
operator
()(
const
paddle
::
experimental
::
Tensor
&
tensor
)
=
0
;
virtual
void
*
operator
()(
void
*
py_tensor
)
=
0
;
};
class
UnPackHookBase
{
public:
virtual
~
UnPackHookBase
()
=
default
;
virtual
paddle
::
experimental
::
Tensor
operator
()(
void
*
packed_value
)
=
0
;
virtual
void
*
operator
()(
void
*
packed_value
,
void
*
other
)
=
0
;
};
}
// namespace egr
paddle/fluid/eager/saved_tensors_hooks.cc
0 → 100644
浏览文件 @
b294f054
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "paddle/fluid/eager/saved_tensors_hooks.h"
#include "paddle/fluid/eager/api/utils/global_utils.h"
#if !(defined(PADDLE_NO_PYTHON) && defined(PADDLE_ON_INFERENCE))
#include "paddle/fluid/pybind/eager.h"
#include "paddle/fluid/pybind/eager_utils.h"
#endif
namespace
egr
{
#if !(defined(PADDLE_NO_PYTHON) && defined(PADDLE_ON_INFERENCE))
PackHook
::
PackHook
(
PyObject
*
hook
)
:
hook_
(
hook
)
{
Py_INCREF
(
hook_
);
}
PackHook
::~
PackHook
()
{
::
pybind11
::
gil_scoped_acquire
gil
;
Py_DECREF
(
hook_
);
}
void
*
PackHook
::
operator
()(
const
paddle
::
experimental
::
Tensor
&
tensor
)
{
bool
grad_tmp
=
egr
::
Controller
::
Instance
().
HasGrad
();
egr
::
Controller
::
Instance
().
SetHasGrad
(
false
);
::
pybind11
::
gil_scoped_acquire
gil
;
auto
args
=
PyTuple_New
(
1
);
PyTuple_SET_ITEM
(
args
,
0
,
paddle
::
pybind
::
ToPyObject
(
tensor
));
PyObject
*
ret
=
PyObject_Call
(
hook_
,
args
,
nullptr
);
Py_XDECREF
(
args
);
egr
::
Controller
::
Instance
().
SetHasGrad
(
grad_tmp
);
return
reinterpret_cast
<
void
*>
(
ret
);
}
void
*
PackHook
::
operator
()(
void
*
py_tensor
)
{
bool
grad_tmp
=
egr
::
Controller
::
Instance
().
HasGrad
();
egr
::
Controller
::
Instance
().
SetHasGrad
(
false
);
::
pybind11
::
gil_scoped_acquire
gil
;
auto
args
=
PyTuple_New
(
1
);
Py_INCREF
(
reinterpret_cast
<
PyObject
*>
(
py_tensor
));
PyTuple_SET_ITEM
(
args
,
0
,
reinterpret_cast
<
PyObject
*>
(
py_tensor
));
PyObject
*
ret
=
PyObject_Call
(
hook_
,
args
,
nullptr
);
Py_XDECREF
(
args
);
egr
::
Controller
::
Instance
().
SetHasGrad
(
grad_tmp
);
return
reinterpret_cast
<
void
*>
(
ret
);
}
UnPackHook
::
UnPackHook
(
PyObject
*
hook
)
:
hook_
(
hook
)
{
Py_INCREF
(
hook_
);
}
UnPackHook
::~
UnPackHook
()
{
::
pybind11
::
gil_scoped_acquire
gil
;
Py_DECREF
(
hook_
);
}
paddle
::
experimental
::
Tensor
UnPackHook
::
operator
()(
void
*
packed_value
)
{
bool
grad_tmp
=
egr
::
Controller
::
Instance
().
HasGrad
();
egr
::
Controller
::
Instance
().
SetHasGrad
(
false
);
::
pybind11
::
gil_scoped_acquire
gil
;
auto
args
=
PyTuple_New
(
1
);
Py_INCREF
(
reinterpret_cast
<
PyObject
*>
(
packed_value
));
PyTuple_SET_ITEM
(
args
,
0
,
reinterpret_cast
<
PyObject
*>
(
packed_value
));
PyObject
*
ret
=
PyObject_Call
(
hook_
,
args
,
nullptr
);
Py_XDECREF
(
args
);
egr
::
Controller
::
Instance
().
SetHasGrad
(
grad_tmp
);
PADDLE_ENFORCE_EQ
(
paddle
::
pybind
::
IsEagerTensor
(
ret
),
true
,
paddle
::
platform
::
errors
::
InvalidArgument
(
"paddle.autograd.saved_tensors_hooks only one pair "
"of hooks is allowed at a time."
));
auto
tensor
=
reinterpret_cast
<
paddle
::
pybind
::
TensorObject
*>
(
ret
)
->
tensor
;
Py_XDECREF
(
ret
);
return
tensor
;
}
void
*
UnPackHook
::
operator
()(
void
*
packed_value
,
void
*
other
)
{
bool
grad_tmp
=
egr
::
Controller
::
Instance
().
HasGrad
();
egr
::
Controller
::
Instance
().
SetHasGrad
(
false
);
::
pybind11
::
gil_scoped_acquire
gil
;
auto
args
=
PyTuple_New
(
1
);
Py_INCREF
(
reinterpret_cast
<
PyObject
*>
(
packed_value
));
PyTuple_SET_ITEM
(
args
,
0
,
reinterpret_cast
<
PyObject
*>
(
packed_value
));
PyObject
*
ret
=
PyObject_Call
(
hook_
,
args
,
nullptr
);
Py_XDECREF
(
args
);
egr
::
Controller
::
Instance
().
SetHasGrad
(
grad_tmp
);
PADDLE_ENFORCE_EQ
(
paddle
::
pybind
::
IsEagerTensor
(
ret
),
true
,
paddle
::
platform
::
errors
::
InvalidArgument
(
"paddle.autograd.saved_tensors_hooks only one pair "
"of hooks is allowed at a time."
));
return
reinterpret_cast
<
void
*>
(
ret
);
}
#endif
}
// namespace egr
paddle/fluid/eager/saved_tensors_hooks.h
0 → 100644
浏览文件 @
b294f054
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <Python.h>
#include "paddle/fluid/eager/api/utils/global_utils.h"
#include "paddle/fluid/eager/hooks.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/phi/core/errors.h"
namespace
egr
{
#if !(defined(PADDLE_NO_PYTHON) && defined(PADDLE_ON_INFERENCE))
class
PackHook
:
public
PackHookBase
{
public:
explicit
PackHook
(
PyObject
*
hook
);
~
PackHook
();
void
*
operator
()(
const
paddle
::
experimental
::
Tensor
&
tensor
)
override
;
void
*
operator
()(
void
*
py_tensor
)
override
;
private:
PyObject
*
hook_
;
};
class
UnPackHook
:
public
UnPackHookBase
{
public:
explicit
UnPackHook
(
PyObject
*
hook
);
~
UnPackHook
();
paddle
::
experimental
::
Tensor
operator
()(
void
*
packed_value
)
override
;
void
*
operator
()(
void
*
packed_value
,
void
*
other
)
override
;
private:
PyObject
*
hook_
;
};
#endif
class
SavedTensorsHooks
{
public:
SavedTensorsHooks
()
=
default
;
~
SavedTensorsHooks
()
{}
void
SetHooks
(
PyObject
*
pack_hook
,
PyObject
*
unpack_hook
)
{
#if !(defined(PADDLE_NO_PYTHON) && defined(PADDLE_ON_INFERENCE))
PADDLE_ENFORCE_EQ
(
pack_hook_
==
nullptr
&&
unpack_hook_
==
nullptr
,
true
,
paddle
::
platform
::
errors
::
InvalidArgument
(
"paddle.autograd.saved_tensors_hooks only one pair "
"of hooks is allowed at a time."
));
pack_hook_
=
std
::
make_shared
<
PackHook
>
(
pack_hook
);
unpack_hook_
=
std
::
make_shared
<
UnPackHook
>
(
unpack_hook
);
is_enable_
=
true
;
#endif
}
void
ResetHooks
()
{
#if !(defined(PADDLE_NO_PYTHON) && defined(PADDLE_ON_INFERENCE))
pack_hook_
=
nullptr
;
unpack_hook_
=
nullptr
;
is_enable_
=
false
;
#endif
}
bool
IsEnable
()
{
return
is_enable_
;
}
std
::
shared_ptr
<
PackHookBase
>
GetPackHook
()
{
return
pack_hook_
;
}
std
::
shared_ptr
<
UnPackHookBase
>
GetUnPackHook
()
{
return
unpack_hook_
;
}
static
SavedTensorsHooks
&
GetInstance
()
{
static
SavedTensorsHooks
instance
;
return
instance
;
}
private:
std
::
shared_ptr
<
PackHookBase
>
pack_hook_
;
std
::
shared_ptr
<
UnPackHookBase
>
unpack_hook_
;
bool
is_enable_
{
false
};
};
}
// namespace egr
paddle/fluid/eager/tensor_wrapper.h
浏览文件 @
b294f054
...
...
@@ -27,6 +27,7 @@
#pragma once
#include "paddle/fluid/eager/autograd_meta.h"
#include "paddle/fluid/eager/grad_node_info.h"
#include "paddle/fluid/eager/saved_tensors_hooks.h"
#include "paddle/fluid/eager/utils.h"
#include "paddle/phi/api/lib/utils/allocator.h"
...
...
@@ -69,7 +70,20 @@ class TensorWrapper {
"Unrecognized tensor type for no_need_buffer feature"
));
}
}
else
{
intermidiate_tensor_
.
set_impl
(
tensor
.
impl
());
if
(
SavedTensorsHooks
::
GetInstance
().
IsEnable
()
&&
tensor
.
is_dense_tensor
())
{
phi
::
DenseTensor
*
dense_tensor
=
static_cast
<
phi
::
DenseTensor
*>
(
tensor
.
impl
().
get
());
intermidiate_tensor_
.
set_impl
(
std
::
move
(
std
::
make_shared
<
phi
::
DenseTensor
>
(
std
::
make_shared
<
phi
::
Allocation
>
(
nullptr
,
0
,
tensor
.
place
()),
dense_tensor
->
meta
())));
auto
pack_hook
=
SavedTensorsHooks
::
GetInstance
().
GetPackHook
();
unpack_hook_
=
SavedTensorsHooks
::
GetInstance
().
GetUnPackHook
();
packed_value_
=
reinterpret_cast
<
PyObject
*>
((
*
pack_hook
)(
tensor
));
}
else
{
intermidiate_tensor_
.
set_impl
(
tensor
.
impl
());
}
}
if
(
VLOG_IS_ON
(
7
))
{
...
...
@@ -86,6 +100,29 @@ class TensorWrapper {
}
}
TensorWrapper
(
const
TensorWrapper
&
other
)
{
no_need_buffer_
=
other
.
no_need_buffer_
;
intermidiate_tensor_
=
other
.
intermidiate_tensor_
;
weak_grad_node_
=
other
.
weak_grad_node_
;
inplace_version_snapshot_
=
other
.
inplace_version_snapshot_
;
packed_value_
=
other
.
packed_value_
;
unpack_hook_
=
other
.
unpack_hook_
;
Py_XINCREF
(
packed_value_
);
}
TensorWrapper
&
operator
=
(
const
TensorWrapper
&
other
)
{
no_need_buffer_
=
other
.
no_need_buffer_
;
intermidiate_tensor_
=
other
.
intermidiate_tensor_
;
weak_grad_node_
=
other
.
weak_grad_node_
;
inplace_version_snapshot_
=
other
.
inplace_version_snapshot_
;
packed_value_
=
other
.
packed_value_
;
unpack_hook_
=
other
.
unpack_hook_
;
Py_XINCREF
(
packed_value_
);
return
*
this
;
}
~
TensorWrapper
()
{
Py_XDECREF
(
packed_value_
);
}
paddle
::
experimental
::
Tensor
recover
()
{
VLOG
(
6
)
<<
"Recover tensor: "
<<
intermidiate_tensor_
.
name
()
<<
" for wrapper"
;
...
...
@@ -94,7 +131,16 @@ class TensorWrapper {
return
paddle
::
experimental
::
Tensor
();
}
check_inplace_version
();
if
(
packed_value_
&&
unpack_hook_
)
{
auto
tensor_unpacked
=
(
*
unpack_hook_
)(
reinterpret_cast
<
void
*>
(
packed_value_
));
auto
src_dense_tensor
=
static_cast
<
phi
::
DenseTensor
*>
(
tensor_unpacked
.
impl
().
get
());
static_cast
<
phi
::
DenseTensor
*>
(
intermidiate_tensor_
.
impl
().
get
())
->
ResetHolder
(
src_dense_tensor
->
MoveMemoryHolder
());
}
else
{
check_inplace_version
();
}
paddle
::
experimental
::
Tensor
recovered_tensor
=
intermidiate_tensor_
;
...
...
@@ -168,5 +214,7 @@ class TensorWrapper {
paddle
::
experimental
::
Tensor
intermidiate_tensor_
;
std
::
weak_ptr
<
egr
::
GradNodeBase
>
weak_grad_node_
;
uint32_t
inplace_version_snapshot_
=
0
;
PyObject
*
packed_value_
{
nullptr
};
std
::
shared_ptr
<
UnPackHookBase
>
unpack_hook_
;
};
}
// namespace egr
paddle/fluid/pybind/CMakeLists.txt
浏览文件 @
b294f054
...
...
@@ -41,7 +41,8 @@ set(PYBIND_DEPS
new_profiler
auto_parallel
jit_layer
jit_property
)
jit_property
saved_tensors_hooks
)
if
(
WITH_PSCORE
)
set
(
PYBIND_DEPS
${
PYBIND_DEPS
}
ps_service
)
...
...
paddle/fluid/pybind/eager.h
浏览文件 @
b294f054
...
...
@@ -12,6 +12,7 @@ limitations under the License. */
#include <Python.h>
#include "paddle/fluid/eager/hooks.h"
#include "paddle/fluid/eager/pylayer/py_layer_node.h"
#include "paddle/phi/core/dense_tensor.h"
#include "pybind11/pybind11.h"
...
...
@@ -28,6 +29,8 @@ typedef struct {
typedef
struct
{
PyObject_HEAD
PyObject
*
container
;
bool
container_be_packed
;
std
::
shared_ptr
<
egr
::
UnPackHookBase
>
unpack_hook
;
PyObject
*
non_differentiable
;
PyObject
*
not_inplace_tensors
;
bool
materialize_grads
;
...
...
paddle/fluid/pybind/eager_functions.cc
浏览文件 @
b294f054
...
...
@@ -25,6 +25,7 @@ typedef SSIZE_T ssize_t;
#include "paddle/fluid/eager/autograd_meta.h"
#include "paddle/fluid/eager/backward.h"
#include "paddle/fluid/eager/custom_operator/custom_operator_node.h"
#include "paddle/fluid/eager/saved_tensors_hooks.h"
#include "paddle/fluid/eager/utils.h"
#include "paddle/fluid/framework/convert_utils.h"
#include "paddle/fluid/framework/custom_operator.h"
...
...
@@ -591,6 +592,29 @@ static PyObject* eager_api_sparse_csr_tensor(PyObject* self,
return
ToPyObject
(
tensor
);
EAGER_CATCH_AND_THROW_RETURN_NULL
}
static
PyObject
*
eager_api_register_saved_tensors_hooks
(
PyObject
*
self
,
PyObject
*
args
,
PyObject
*
kwargs
)
{
EAGER_TRY
if
(
egr
::
Controller
::
Instance
().
HasGrad
())
{
auto
pack_hook
=
PyTuple_GET_ITEM
(
args
,
0
);
auto
unpack_hook
=
PyTuple_GET_ITEM
(
args
,
1
);
egr
::
SavedTensorsHooks
::
GetInstance
().
SetHooks
(
pack_hook
,
unpack_hook
);
}
RETURN_PY_NONE
EAGER_CATCH_AND_THROW_RETURN_NULL
}
static
PyObject
*
eager_api_reset_saved_tensors_hooks
(
PyObject
*
self
,
PyObject
*
args
,
PyObject
*
kwargs
)
{
EAGER_TRY
egr
::
SavedTensorsHooks
::
GetInstance
().
ResetHooks
();
RETURN_PY_NONE
EAGER_CATCH_AND_THROW_RETURN_NULL
}
#if defined(PADDLE_WITH_CUDA)
static
PyObject
*
eager_api_async_read
(
PyObject
*
self
,
PyObject
*
args
,
...
...
@@ -965,6 +989,14 @@ PyMethodDef variable_functions[] = {
(
PyCFunction
)(
void
(
*
)(
void
))
eager_api_sparse_csr_tensor
,
METH_VARARGS
|
METH_KEYWORDS
,
NULL
},
{
"register_saved_tensors_hooks"
,
(
PyCFunction
)(
void
(
*
)(
void
))
eager_api_register_saved_tensors_hooks
,
METH_VARARGS
|
METH_KEYWORDS
,
NULL
},
{
"reset_saved_tensors_hooks"
,
(
PyCFunction
)(
void
(
*
)(
void
))
eager_api_reset_saved_tensors_hooks
,
METH_VARARGS
|
METH_KEYWORDS
,
NULL
},
/**sparse functions**/
#if defined(PADDLE_WITH_CUDA)
{
"async_read"
,
...
...
paddle/fluid/pybind/eager_py_layer.cc
浏览文件 @
b294f054
...
...
@@ -20,6 +20,7 @@ limitations under the License. */
#include "paddle/fluid/eager/api/all.h"
#include "paddle/fluid/eager/autograd_meta.h"
#include "paddle/fluid/eager/pylayer/py_layer_node.h"
#include "paddle/fluid/eager/saved_tensors_hooks.h"
#include "paddle/fluid/eager/utils.h"
#include "paddle/fluid/framework/convert_utils.h"
#include "paddle/fluid/memory/allocation/allocator.h"
...
...
@@ -78,6 +79,7 @@ PyObject* PyLayerNew(PyTypeObject* type, PyObject* args, PyObject* kwargs) {
if
(
obj
)
{
auto
v
=
reinterpret_cast
<
PyLayerObject
*>
(
obj
);
v
->
materialize_grads
=
true
;
v
->
container_be_packed
=
false
;
new
(
&
v
->
grad_node
)
std
::
weak_ptr
<
egr
::
GradNodePyLayer
>
();
new
(
&
v
->
forward_input_tensor_is_duplicable
)
std
::
vector
<
bool
>
();
new
(
&
v
->
forward_output_tensor_is_duplicable
)
std
::
vector
<
bool
>
();
...
...
@@ -96,6 +98,7 @@ static void PyLayerDealloc(PyLayerObject* self) {
Py_DECREF
(
self
->
not_inplace_tensors
);
}
self
->
grad_node
.
~
weak_ptr
<
egr
::
GradNodePyLayer
>
();
self
->
unpack_hook
=
nullptr
;
self
->
forward_input_tensor_is_duplicable
.
~
vector
();
self
->
forward_output_tensor_is_duplicable
.
~
vector
();
Py_TYPE
(
self
)
->
tp_free
(
reinterpret_cast
<
PyObject
*>
(
self
));
...
...
@@ -455,23 +458,148 @@ PyObject* pylayer_method_apply(PyObject* cls,
EAGER_CATCH_AND_THROW_RETURN_NULL
}
PyObject
*
call_unpack_hook
(
PyLayerObject
*
self
)
{
auto
unpack_hook
=
self
->
unpack_hook
;
auto
packed_value
=
self
->
container
;
auto
packed_value_size
=
PyTuple_GET_SIZE
(
packed_value
);
auto
unpacked_value
=
PyTuple_New
(
packed_value_size
);
for
(
Py_ssize_t
i
=
0
;
i
<
packed_value_size
;
i
++
)
{
PyObject
*
obj
=
PyTuple_GET_ITEM
(
packed_value
,
i
);
if
(
PyList_Check
(
obj
))
{
Py_ssize_t
len
=
PyList_Size
(
obj
);
auto
tmp_list
=
PyList_New
(
len
);
for
(
Py_ssize_t
j
=
0
;
j
<
len
;
j
++
)
{
PyObject
*
o
=
PyList_GetItem
(
obj
,
j
);
PyTuple_SET_ITEM
(
tmp_list
,
j
,
reinterpret_cast
<
PyObject
*>
(((
*
unpack_hook
)(
reinterpret_cast
<
void
*>
(
o
),
nullptr
))));
}
PyTuple_SET_ITEM
(
unpacked_value
,
i
,
tmp_list
);
}
else
if
(
PyTuple_Check
(
obj
))
{
Py_ssize_t
len
=
PyTuple_Size
(
obj
);
auto
tmp_tuple
=
PyTuple_New
(
len
);
for
(
Py_ssize_t
j
=
0
;
j
<
len
;
j
++
)
{
PyObject
*
o
=
PyTuple_GetItem
(
obj
,
j
);
PyTuple_SET_ITEM
(
tmp_tuple
,
j
,
reinterpret_cast
<
PyObject
*>
((
*
unpack_hook
)(
reinterpret_cast
<
void
*>
(
o
),
nullptr
)));
}
PyTuple_SET_ITEM
(
unpacked_value
,
i
,
tmp_tuple
);
}
else
{
PyTuple_SET_ITEM
(
unpacked_value
,
i
,
reinterpret_cast
<
PyObject
*>
((
*
unpack_hook
)(
reinterpret_cast
<
void
*>
(
obj
),
nullptr
)));
}
}
return
unpacked_value
;
}
PyObject
*
tensor_properties_get_container
(
PyLayerObject
*
self
,
void
*
closure
)
{
EAGER_TRY
if
(
self
->
container
==
nullptr
)
{
RETURN_PY_NONE
;
}
Py_INCREF
(
self
->
container
);
return
self
->
container
;
if
(
self
->
container_be_packed
)
{
return
call_unpack_hook
(
self
);
}
else
{
Py_INCREF
(
self
->
container
);
return
self
->
container
;
}
EAGER_CATCH_AND_THROW_RETURN_NULL
}
void
call_pack_hook
(
PyLayerObject
*
self
,
PyObject
*
value
)
{
PyObject
*
saved_value
=
nullptr
;
if
(
PyTuple_Check
(
value
))
{
saved_value
=
value
;
}
else
if
(
PyList_Check
(
value
))
{
saved_value
=
PyList_AsTuple
(
value
);
}
else
{
saved_value
=
PyTuple_New
(
1
);
Py_INCREF
(
value
);
PyTuple_SET_ITEM
(
saved_value
,
0
,
value
);
}
auto
pack_hook
=
egr
::
SavedTensorsHooks
::
GetInstance
().
GetPackHook
();
self
->
unpack_hook
=
egr
::
SavedTensorsHooks
::
GetInstance
().
GetUnPackHook
();
auto
saved_value_size
=
PyTuple_GET_SIZE
(
saved_value
);
PyObject
*
packed_value
=
PyTuple_New
(
saved_value_size
);
for
(
Py_ssize_t
i
=
0
;
i
<
saved_value_size
;
i
++
)
{
PyObject
*
obj
=
PyTuple_GET_ITEM
(
saved_value
,
i
);
if
(
IsEagerTensor
(
obj
))
{
PyTuple_SET_ITEM
(
packed_value
,
i
,
reinterpret_cast
<
PyObject
*>
(
(
*
pack_hook
)(
reinterpret_cast
<
void
*>
(
obj
))));
}
else
if
(
PyList_Check
(
obj
))
{
Py_ssize_t
len
=
PyList_Size
(
obj
);
auto
tmp_list
=
PyList_New
(
len
);
for
(
Py_ssize_t
j
=
0
;
j
<
len
;
j
++
)
{
PyObject
*
o
=
PyList_GetItem
(
obj
,
j
);
if
(
IsEagerTensor
(
o
))
{
PyTuple_SET_ITEM
(
tmp_list
,
j
,
reinterpret_cast
<
PyObject
*>
(
(
*
pack_hook
)(
reinterpret_cast
<
void
*>
(
o
))));
}
else
{
PADDLE_THROW
(
platform
::
errors
::
InvalidArgument
(
"save_for_backward only support Tensor, list of Tensor, tuple of "
"Tensor."
));
}
}
PyTuple_SET_ITEM
(
packed_value
,
i
,
tmp_list
);
}
else
if
(
PyTuple_Check
(
obj
))
{
Py_ssize_t
len
=
PyTuple_Size
(
obj
);
auto
tmp_tuple
=
PyTuple_New
(
len
);
for
(
Py_ssize_t
j
=
0
;
j
<
len
;
j
++
)
{
PyObject
*
o
=
PyTuple_GetItem
(
obj
,
j
);
if
(
IsEagerTensor
(
o
))
{
PyTuple_SET_ITEM
(
tmp_tuple
,
j
,
reinterpret_cast
<
PyObject
*>
(
(
*
pack_hook
)(
reinterpret_cast
<
void
*>
(
o
))));
}
else
{
PADDLE_THROW
(
platform
::
errors
::
InvalidArgument
(
"save_for_backward only support Tensor, list of Tensor, tuple of "
"Tensor."
));
}
}
PyTuple_SET_ITEM
(
packed_value
,
i
,
tmp_tuple
);
}
else
{
PADDLE_THROW
(
platform
::
errors
::
InvalidArgument
(
"save_for_backward only support Tensor, list of Tensor, tuple of "
"Tensor."
));
}
}
if
(
PyTuple_Check
(
value
))
{
Py_XDECREF
(
saved_value
);
}
Py_XDECREF
(
self
->
container
);
self
->
container
=
packed_value
;
self
->
container_be_packed
=
true
;
}
int
tensor_properties_set_container
(
PyLayerObject
*
self
,
PyObject
*
value
,
void
*
closure
)
{
EAGER_TRY
Py_XINCREF
(
value
);
Py_XDECREF
(
self
->
container
);
self
->
container
=
value
;
if
(
egr
::
SavedTensorsHooks
::
GetInstance
().
IsEnable
())
{
call_pack_hook
(
self
,
value
);
}
else
{
Py_XINCREF
(
value
);
Py_XDECREF
(
self
->
container
);
self
->
container
=
value
;
}
return
0
;
EAGER_CATCH_AND_THROW_RETURN_NEG
}
...
...
python/paddle/autograd/__init__.py
浏览文件 @
b294f054
...
...
@@ -26,9 +26,11 @@ else:
from
.py_layer
import
LegacyPyLayerContext
as
PyLayerContext
# noqa: F401
from
..framework
import
set_grad_enabled
,
is_grad_enabled
# noqa: F401
from
..fluid.dygraph.base
import
no_grad_
as
no_grad
# noqa: F401
from
.saved_tensors_hooks
import
saved_tensors_hooks
__all__
=
[
# noqa
'backward'
,
'PyLayer'
,
'PyLayerContext'
,
'saved_tensors_hooks'
,
]
python/paddle/autograd/saved_tensors_hooks.py
0 → 100644
浏览文件 @
b294f054
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from
paddle.fluid
import
core
__all__
=
[]
class
saved_tensors_hooks
():
"""
Dynamic graph, registers a pair of pack / unpack hooks for saved tensors.
Parameters:
pack_hook (function): The pack hook will be called every time the forward
operation inputs/outputs tensors need be saved for backward. Then you
can save it to CPU or Disk. The input of `pack_hook` is a tensor need
be saved. The output of `pack_hook` is then stored information instead
of the original tensor. `pack_hook` will also be called while any
tensor need be saved by `PyLayerContext.save_for_backward`. If a tensor
saved for backward is no need buffer, `pack_hook` will not be called.
Only the thensor saved for backward is LoDTensor, `pack_hook` will be
called.
unpack_hook (function): The unpack hook will be called every time the
backward need use the saved inputs/outputs tensors. Then you can reload
the tensor and return it to paddle framework. The input of `unpack_hook`
is the information returned by `pack_hook`. The output of `unpack_hook`
is a tensor reloaded by the information, and the tensor mast has the same
content as the original tensor passed as input to the corresponding
`pack_hook`.
Returns:
None
Examples:
.. code-block:: python
# Example1
import paddle
def pack_hook(x):
print("Packing", x)
return x.numpy()
def unpack_hook(x):
print("UnPacking", x)
return paddle.to_tensor(x)
a = paddle.ones([3,3])
b = paddle.ones([3,3]) * 2
a.stop_gradient = False
b.stop_gradient = False
with paddle.autograd.saved_tensors_hooks(pack_hook, unpack_hook):
y = paddle.multiply(a, b)
y.sum().backward()
# Example2
import paddle
from paddle.autograd import PyLayer
class cus_multiply(PyLayer):
@staticmethod
def forward(ctx, a, b):
y = paddle.multiply(a, b)
ctx.save_for_backward(a, b)
return y
@staticmethod
def backward(ctx, dy):
a,b = ctx.saved_tensor()
grad_a = dy * a
grad_b = dy * b
return grad_a, grad_b
def pack_hook(x):
print("Packing", x)
return x.numpy()
def unpack_hook(x):
print("UnPacking", x)
return paddle.to_tensor(x)
a = paddle.ones([3,3])
b = paddle.ones([3,3]) * 2
a.stop_gradient = False
b.stop_gradient = False
with paddle.autograd.saved_tensors_hooks(pack_hook, unpack_hook):
y = cus_multiply.apply(a, b)
y.sum().backward()
"""
def
__init__
(
self
,
pack_hook
,
unpack_hook
):
self
.
pack_hook
=
pack_hook
self
.
unpack_hook
=
unpack_hook
def
__enter__
(
self
):
core
.
eager
.
register_saved_tensors_hooks
(
self
.
pack_hook
,
self
.
unpack_hook
)
def
__exit__
(
self
,
*
args
):
core
.
eager
.
reset_saved_tensors_hooks
()
python/paddle/fluid/tests/unittests/test_saved_tensors_hooks.py
0 → 100644
浏览文件 @
b294f054
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import
unittest
import
numpy
as
np
import
paddle
from
paddle.autograd
import
PyLayer
class
TestSavedTensorsHooks
(
unittest
.
TestCase
):
def
test_save_for_multiply
(
self
):
def
pack_hook
(
x
):
return
x
.
numpy
()
def
unpack_hook
(
x
):
return
paddle
.
to_tensor
(
x
)
a
=
paddle
.
ones
([
3
,
3
])
b
=
paddle
.
ones
([
3
,
3
])
*
2
a
.
stop_gradient
=
False
b
.
stop_gradient
=
False
with
paddle
.
autograd
.
saved_tensors_hooks
(
pack_hook
,
unpack_hook
):
y
=
paddle
.
multiply
(
a
,
b
)
y
.
sum
().
backward
()
aa
=
paddle
.
ones
([
3
,
3
])
bb
=
paddle
.
ones
([
3
,
3
])
*
2
aa
.
stop_gradient
=
False
bb
.
stop_gradient
=
False
yy
=
paddle
.
multiply
(
aa
,
bb
)
yy
.
sum
().
backward
()
self
.
assertTrue
(
paddle
.
equal_all
(
aa
.
grad
,
a
.
grad
))
self
.
assertTrue
(
paddle
.
equal_all
(
bb
.
grad
,
b
.
grad
))
def
test_save_for_pylayer
(
self
):
class
cus_multiply
(
PyLayer
):
@
staticmethod
def
forward
(
ctx
,
a
,
b
):
y
=
paddle
.
multiply
(
a
,
b
)
ctx
.
save_for_backward
(
a
,
b
)
return
y
@
staticmethod
def
backward
(
ctx
,
dy
):
a
,
b
=
ctx
.
saved_tensor
()
grad_a
=
dy
*
a
grad_b
=
dy
*
b
return
grad_a
,
grad_b
def
pack_hook
(
x
):
return
x
.
numpy
()
def
unpack_hook
(
x
):
return
paddle
.
to_tensor
(
x
)
a
=
paddle
.
ones
([
3
,
3
])
b
=
paddle
.
ones
([
3
,
3
])
*
2
a
.
stop_gradient
=
False
b
.
stop_gradient
=
False
with
paddle
.
autograd
.
saved_tensors_hooks
(
pack_hook
,
unpack_hook
):
y
=
cus_multiply
.
apply
(
a
,
b
)
y
.
sum
().
backward
()
aa
=
paddle
.
ones
([
3
,
3
])
bb
=
paddle
.
ones
([
3
,
3
])
*
2
aa
.
stop_gradient
=
False
bb
.
stop_gradient
=
False
yy
=
cus_multiply
.
apply
(
aa
,
bb
)
yy
.
sum
().
backward
()
self
.
assertTrue
(
paddle
.
equal_all
(
aa
.
grad
,
a
.
grad
))
self
.
assertTrue
(
paddle
.
equal_all
(
bb
.
grad
,
b
.
grad
))
if
__name__
==
'__main__'
:
unittest
.
main
()
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录