diff --git a/paddle/fluid/framework/selected_rows.cc b/paddle/fluid/framework/selected_rows.cc index d9d6b7dd67f1c6e4bbd6a4e1a8f0843d4cb93c05..794e7f743413b068119afd5df232bfc2bb91a8c7 100644 --- a/paddle/fluid/framework/selected_rows.cc +++ b/paddle/fluid/framework/selected_rows.cc @@ -17,6 +17,52 @@ limitations under the License. */ namespace paddle { namespace framework { +struct ReAllocateVisitor { + ReAllocateVisitor(framework::Tensor* tensor, const framework::DDim& dims) + : tensor_(tensor), dims_(dims) {} + + template + void operator()() const { + framework::Tensor cpu_tensor; + platform::CPUPlace cpu; + T* ptr = cpu_tensor.mutable_data(dims_, cpu); + const T* old_ptr = + tensor_->memory_size() == 0 ? nullptr : tensor_->data(); + if (old_ptr != nullptr) { + std::copy(old_ptr, old_ptr + tensor_->numel(), ptr); + } + tensor_->ShareDataWith(cpu_tensor); + } + + framework::Tensor* tensor_; + framework::DDim dims_; +}; + +struct TensorCopyVisitor { + TensorCopyVisitor(framework::Tensor* dst, int64_t dst_offset, + const framework::Tensor src, int64_t src_offset, + int64_t size) + : dst_(dst), + dst_offset_(dst_offset), + src_(src), + src_offset_(src_offset), + size_(size) {} + + template + void operator()() const { + // TODO(Yancey1989): support other place + platform::CPUPlace cpu; + memory::Copy(cpu, dst_->mutable_data(cpu) + dst_offset_, cpu, + src_.data() + src_offset_, size_ * sizeof(T)); + } + + framework::Tensor* dst_; + int64_t dst_offset_; + framework::Tensor src_; + int64_t src_offset_; + int64_t size_; +}; + void SerializeToStream(std::ostream& os, const SelectedRows& selected_rows, const platform::DeviceContext& dev_ctx) { { // the 1st field, uint32_t version @@ -69,5 +115,66 @@ void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows, TensorFromStream(is, selected_rows->mutable_value(), dev_ctx); } +bool SelectedRows::HasKey(int64_t key) const { + return std::find(rows_.begin(), rows_.end(), key) == rows_.end() ? false + : true; +} + +std::vector SelectedRows::Get(std::vector keys, + framework::Tensor* value) const { + PADDLE_ENFORCE(value->IsInitialized(), + "The value tensor should be initialized."); + std::vector non_keys; + int64_t value_width = value_->numel() / value_->dims()[0]; + PADDLE_ENFORCE_EQ(value_width, value->numel() / value->dims()[0], + "output tensor should have the same shape with table " + "execpt the dims[0]."); + + for (size_t i = 0; i < keys.size(); ++i) { + int64_t index = Index(keys[i]); + if (index == -1) { + non_keys.push_back(keys[i]); + } else { + framework::VisitDataType( + framework::ToDataType(value_->type()), + TensorCopyVisitor(value, i * value_width, *value_.get(), + index * value_width, value_width)); + } + } + return non_keys; +} + +bool SelectedRows::Set(int64_t key, const framework::Tensor& value) { + PADDLE_ENFORCE(value.IsInitialized(), "The value should be initialized."); + if (value_->IsInitialized()) { + PADDLE_ENFORCE_EQ( + value.type(), value_->type(), + "The type of the value should be same with the original value"); + } + PADDLE_ENFORCE_EQ(value.dims()[0], static_cast(1), + "The first dim of value should be 1."); + auto index = Index(key); + bool is_new_key = false; + if (index == -1) { + rows_.push_back(key); + index = rows_.size() - 1; + is_new_key = true; + // whether need to resize the table + if (static_cast(rows_.size()) > value_->dims()[0]) { + auto dims = value_->dims(); + dims[0] = (dims[0] + 1) << 1; + framework::VisitDataType(framework::ToDataType(value.type()), + ReAllocateVisitor(value_.get(), dims)); + } + } + + framework::VisitDataType( + framework::ToDataType(value.type()), + TensorCopyVisitor(value_.get(), + index * value_->numel() / value_->dims()[0], value, + static_cast(0), value.numel())); + return is_new_key; +} + } // namespace framework } // namespace paddle diff --git a/paddle/fluid/framework/selected_rows.h b/paddle/fluid/framework/selected_rows.h index 8e2d9470d3954e0f66c74828a8d8292c2875a8f4..d6c9507b1681855e759a4b1b9d3dddf6fcb2fc13 100644 --- a/paddle/fluid/framework/selected_rows.h +++ b/paddle/fluid/framework/selected_rows.h @@ -14,15 +14,33 @@ limitations under the License. */ #pragma once +#include #include #include "paddle/fluid/framework/lod_tensor.h" #include "paddle/fluid/framework/tensor.h" +#include "paddle/fluid/memory/memcpy.h" namespace paddle { namespace framework { class SelectedRows { + /* + * @brief We can use the SelectedRows structure to reproduce a sparse table. + * A sparse table is a key-value structure that the key is an `int64_t` + * number, + * and the value is a Tensor which the first dimension is 0. + * You can use the following interface to operate the sparse table, and you + * can find + * some detail information from the comments of each interface: + * + * HasKey(key), whether the sparse table has the specified key. + * Set(key, value), set a key-value pair into the sparse table. + * Get(keys, value*), get value by given key list and apply it to the given + * value pointer + * with the specified offset. + * + */ public: SelectedRows(const std::vector& rows, const int64_t& height) : rows_(rows), height_(height) { @@ -50,12 +68,45 @@ class SelectedRows { void set_rows(const Vector& rows) { rows_ = rows; } - /** - * get the index of id in rows + /* + * @brief wheter has the specified key in the table. + * + * @return true if the key is exists. + */ + bool HasKey(int64_t key) const; + + /* + * @brief Get value by the key list, if the + * + * @return a list of keys which does not exists in table + */ + std::vector Get(std::vector keys, + framework::Tensor* tensor) const; + + /* + * @brief Set a key-value pair into the table. + * This function will double the value memory if it's not engouth. + * + * @note: + * 1. The first dim of the value should be 1 + * 2. The value should be initialized and the data type + * should be the same with the table. + * + * @return true if the key is a new one, otherwise false + * + */ + bool Set(int64_t key, const Tensor& value); + + /* + * @brief Get the index of key in rows + * + * @return -1 if the key does not exists. */ - int64_t index(int64_t id) const { - auto it = std::find(rows_.begin(), rows_.end(), id); - PADDLE_ENFORCE(it != rows_.end(), "id should be in rows"); + int64_t Index(int64_t key) const { + auto it = std::find(rows_.begin(), rows_.end(), key); + if (it == rows_.end()) { + return static_cast(-1); + } return static_cast(std::distance(rows_.begin(), it)); } diff --git a/paddle/fluid/framework/selected_rows_test.cc b/paddle/fluid/framework/selected_rows_test.cc index 960d8d64f04a819217413ff881977ce5fb5a30f2..39fe6d92940606084c28eec1a4d6486cb58844ce 100644 --- a/paddle/fluid/framework/selected_rows_test.cc +++ b/paddle/fluid/framework/selected_rows_test.cc @@ -17,7 +17,7 @@ namespace framework { class SelectedRowsTester : public ::testing::Test { public: - virtual void SetUp() override { + void SetUp() override { std::vector rows{0, 4, 7}; int64_t height = 10; int64_t row_numel = 100; @@ -59,5 +59,40 @@ TEST_F(SelectedRowsTester, SerializeAndDeseralize) { ASSERT_EQ(selected_rows_->GetCompleteDims(), dst_tensor.GetCompleteDims()); } +TEST_F(SelectedRowsTester, Table) { + platform::CPUPlace cpu; + SelectedRows table; + // initialize a sparse table + table.mutable_value()->Resize(framework::make_ddim({1, 100})); + table.mutable_value()->mutable_data(cpu); + table.mutable_rows()->push_back(1); + + int64_t key = 10000; + int64_t non_key = 999; + framework::Tensor value; + value.Resize(framework::make_ddim({1, 100})); + auto ptr = value.mutable_data(cpu); + ptr[0] = static_cast(10); + + ASSERT_EQ(table.rows().size(), static_cast(1)); + ASSERT_EQ(table.HasKey(key), false); + + table.Set(key, value); + + ASSERT_EQ(table.rows().size(), static_cast(2)); + ASSERT_EQ(table.HasKey(key), true); + // check re-allocate + ASSERT_EQ(table.value().dims()[0], static_cast(4)); + + framework::Tensor get_value; + get_value.mutable_data(framework::make_ddim({2, 100}), cpu); + std::vector keys({non_key, key}); + auto non_keys = table.Get(keys, &get_value); + + ASSERT_EQ(get_value.data()[100], static_cast(10)); + ASSERT_EQ(non_keys.size(), static_cast(1)); + ASSERT_EQ(non_keys[0], non_key); +} + } // namespace framework } // namespace paddle diff --git a/paddle/fluid/operators/lookup_table_op.h b/paddle/fluid/operators/lookup_table_op.h index cb088c267bcc028ff11583cd73de5ca1722a9b69..d482506bf0361c11a019e32efbf348a64aaf5164 100644 --- a/paddle/fluid/operators/lookup_table_op.h +++ b/paddle/fluid/operators/lookup_table_op.h @@ -103,7 +103,8 @@ class LookupTableKernel : public framework::OpKernel { memset(output + i * row_width, 0, row_width * sizeof(T)); } else { PADDLE_ENFORCE_GE(ids[i], 0); - auto id_index = table_t.index(ids[i]); + auto id_index = table_t.Index(ids[i]); + PADDLE_ENFORCE_GE(id_index, 0, "the input key should be exists."); memcpy(output + i * row_width, table + id_index * row_width, row_width * sizeof(T)); } diff --git a/paddle/fluid/operators/sgd_op.h b/paddle/fluid/operators/sgd_op.h index cfc8793e1e05a7d4fa9207ae77a664b391b9a986..f3e88b0a0b05ef792b2cc8e880bdfddb6e6124d1 100644 --- a/paddle/fluid/operators/sgd_op.h +++ b/paddle/fluid/operators/sgd_op.h @@ -107,7 +107,9 @@ class SGDOpKernel : public framework::OpKernel { for (size_t i = 0; i < grad.rows().size(); i++) { PADDLE_ENFORCE(grad.rows()[i] < grad.height(), "Input rows index should less than height"); - int64_t id_index = param.index(grad.rows()[i]); + int64_t id_index = param.Index(grad.rows()[i]); + PADDLE_ENFORCE_GE(id_index, static_cast(0), + "id should be in the table"); for (size_t j = 0; j < grad_row_width; j++) { out_data[id_index * grad_row_width + j] -= lr[0] * grad_data[i * grad_row_width + j];