FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Add a "Snapshot" kernel that always makes a copy of its input. This s… · codeplaysoftware/tensorflow@54c2584 · GitHub

  1. codeplaysoftware
  2. tensorflow

Repository navigation

Commit 54c2584

Browse files
tensorflower-gardener
committed
Add a "Snapshot" kernel that always makes a copy of its input. This should be used in place of unreadable hacks like "x = y + 0" in places that need to force a copy of a tensor to be made.
PiperOrigin-RevId: 178439351
1 parent 28807c5 commit 54c2584

9 files changed

Lines changed: 173 additions & 0 deletions

File tree

‎tensorflow/core/api_def/base_api/api_def_Snapshot.pbtxt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
op {
2+
graph_op_name: "Snapshot"
3+
summary: "Returns a copy of the input tensor."
4+
}

‎tensorflow/core/api_def/python_api/api_def_Snapshot.pbtxt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
op {
2+
graph_op_name: "Snapshot"
3+
visibility: HIDDEN
4+
}

‎tensorflow/core/kernels/BUILD‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -605,6 +605,7 @@ cc_library(
605605
":reverse_sequence_op",
606606
":shape_ops",
607607
":slice_op",
608+
":snapshot_op",
608609
":split_op",
609610
":split_v_op",
610611
":strided_slice_op",
@@ -801,6 +802,12 @@ tf_kernel_library(
801802
deps = ARRAY_DEPS + [":strided_slice_op"],
802803
)
803804

805+
tf_kernel_library(
806+
name = "snapshot_op",
807+
prefix = "snapshot_op",
808+
deps = ARRAY_DEPS,
809+
)
810+
804811
tf_kernel_library(
805812
name = "split_op",
806813
gpu_srcs = ["cuda_device_array.h"],

‎tensorflow/core/kernels/snapshot_op.cc‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
// See docs in ../ops/array_ops.cc.
17+
#include "tensorflow/core/kernels/snapshot_op.h"
18+
19+
#include "tensorflow/core/framework/op_kernel.h"
20+
#include "tensorflow/core/framework/register_types.h"
21+
#include "tensorflow/core/framework/types.h"
22+
23+
namespace tensorflow {
24+
typedef Eigen::ThreadPoolDevice CPUDevice;
25+
26+
#define REGISTER_KERNEL(TYPE) \
27+
REGISTER_KERNEL_BUILDER( \
28+
Name("Snapshot").Device(DEVICE_CPU).TypeConstraint<TYPE>("T"), \
29+
SnapshotOp<CPUDevice, TYPE>);
30+
31+
TF_CALL_POD_TYPES(REGISTER_KERNEL);
32+
#undef REGISTER_KERNEL
33+
34+
#if TENSORFLOW_USE_SYCL
35+
typedef Eigen::SyclDevice SyclDevice;
36+
#define REGISTER_SYCL_KERNEL(TYPE) \
37+
REGISTER_KERNEL_BUILDER( \
38+
Name("Snapshot").Device(DEVICE_SYCL).TypeConstraint<TYPE>("T"), \
39+
SnapshotOp<SyclDevice, TYPE>);
40+
41+
TF_CALL_POD_TYPES(REGISTER_SYCL_KERNEL);
42+
43+
#undef REGISTER_SYCL_KERNEL
44+
#endif // TENSORFLOW_USE_SYCL
45+
46+
} // namespace tensorflow

‎tensorflow/core/kernels/snapshot_op.h‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#ifndef TENSORFLOW_KERNELS_SNAPSHOT_OP_H_
17+
#define TENSORFLOW_KERNELS_SNAPSHOT_OP_H_
18+
19+
#if GOOGLE_CUDA
20+
#define EIGEN_USE_GPU
21+
#endif
22+
23+
#define EIGEN_USE_THREADS
24+
25+
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
26+
#include "tensorflow/core/framework/op_kernel.h"
27+
28+
namespace tensorflow {
29+
30+
template <typename Device, typename Scalar>
31+
class SnapshotOp : public OpKernel {
32+
public:
33+
explicit SnapshotOp(OpKernelConstruction* context) : OpKernel(context) {}
34+
35+
void Compute(OpKernelContext* context) override {
36+
const Tensor& input = context->input(0);
37+
Tensor* output = nullptr;
38+
OP_REQUIRES_OK(context,
39+
context->allocate_output(0, input.shape(), &output));
40+
const Device& device = context->eigen_device<Device>();
41+
device.memcpy(output->template flat<Scalar>().data(),
42+
input.template flat<Scalar>().data(),
43+
input.NumElements() * sizeof(Scalar));
44+
}
45+
};
46+
47+
} // namespace tensorflow
48+
49+
#endif // TENSORFLOW_KERNELS_SNAPSHOT_OP_H_

‎tensorflow/core/kernels/snapshot_op_gpu.cu.cc‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
#if GOOGLE_CUDA
16+
17+
// See docs in ../ops/array_ops.cc.
18+
#include "tensorflow/core/kernels/snapshot_op.h"
19+
20+
#include "tensorflow/core/framework/op_kernel.h"
21+
#include "tensorflow/core/framework/register_types.h"
22+
#include "tensorflow/core/framework/types.h"
23+
24+
namespace tensorflow {
25+
typedef Eigen::GpuDevice GPUDevice;
26+
27+
#define REGISTER_KERNEL(TYPE) \
28+
REGISTER_KERNEL_BUILDER( \
29+
Name("Snapshot").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \
30+
SnapshotOp<GPUDevice, TYPE>);
31+
32+
TF_CALL_POD_TYPES(REGISTER_KERNEL);
33+
#undef REGISTER_KERNEL
34+
35+
} // namespace tensorflow
36+
37+
#endif // GOOGLE_CUDA

‎tensorflow/core/ops/array_ops.cc‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1722,6 +1722,20 @@ REGISTER_OP("Identity")
17221722
Return a tensor with the same shape and contents as the input tensor or value.
17231723
)Doc");
17241724

1725+
REGISTER_OP("Snapshot")
1726+
.Input("input: T")
1727+
.Output("output: T")
1728+
.Attr("T: type")
1729+
.SetShapeFn([](shape_inference::InferenceContext* c) {
1730+
c->set_output(0, c->input(0));
1731+
auto* handle_data = c->input_handle_shapes_and_types(0);
1732+
if (handle_data != nullptr) {
1733+
c->set_output_handle_shapes_and_types(0, *handle_data);
1734+
}
1735+
return Status::OK();
1736+
})
1737+
.Doc(R"Doc(Returns a copy of the input tensor.)Doc");
1738+
17251739
#ifdef INTEL_MKL
17261740
REGISTER_OP("_MklIdentity")
17271741
.Input("input: T")

‎tensorflow/python/kernel_tests/array_ops_test.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
from tensorflow.python.framework import test_ops
3434
from tensorflow.python.framework import test_util
3535
from tensorflow.python.ops import array_ops
36+
from tensorflow.python.ops import gen_array_ops
3637
from tensorflow.python.ops import gradients_impl
3738
from tensorflow.python.ops import init_ops
3839
from tensorflow.python.ops import math_ops
@@ -1124,5 +1125,15 @@ def testResourceRejection(self):
11241125
guarantee_a.eval()
11251126

11261127

1128+
class SnapshotOpTest(test_util.TensorFlowTestCase):
1129+
1130+
def testInvertPermutation(self):
1131+
for dtype in [dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64]:
1132+
with self.test_session(use_gpu=True):
1133+
x = constant_op.constant([0, 1, 2, 3], dtype=dtype)
1134+
y = gen_array_ops._snapshot(x)
1135+
self.assertAllEqual(y.eval(), [0, 1, 2, 3])
1136+
1137+
11271138
if __name__ == "__main__":
11281139
test_lib.main()

‎tensorflow/python/ops/hidden_ops.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ ParallelConcat
2121
Placeholder
2222
RefIdentity
2323
Reverse
24+
Snapshot
2425
SpaceToBatch
2526
Split
2627
SplitV

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL