selected_rows.h 4.3 KB
Newer Older
1 2
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

Q
qijun 已提交
3 4 5
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
6

Q
qijun 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
8

Q
qijun 已提交
9 10 11 12 13 14 15
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
16

Y
Yancey1989 已提交
17
#include <algorithm>
Y
Yancey1989 已提交
18
#include <utility>
19 20
#include <vector>

Y
Yi Wang 已提交
21 22
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/tensor.h"
Y
Yancey1989 已提交
23
#include "paddle/fluid/memory/memcpy.h"
Q
qijun 已提交
24 25 26 27 28

namespace paddle {
namespace framework {

class SelectedRows {
Y
Yancey1989 已提交
29 30 31 32 33 34 35 36 37 38 39
  /*
   * @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.
Y
Yancey1989 已提交
40
   *  Get(keys, value*), get value by given key list and apply it to the given
Y
Yancey1989 已提交
41 42 43 44
   * value pointer
   *    with the specified offset.
   *
   */
Q
qijun 已提交
45 46 47 48 49 50
 public:
  SelectedRows(const std::vector<int64_t>& rows, const int64_t& height)
      : rows_(rows), height_(height) {
    value_.reset(new Tensor());
  }

Q
QI JUN 已提交
51 52 53 54
  SelectedRows() {
    height_ = 0;
    value_.reset(new Tensor());
  }
Q
qijun 已提交
55 56 57

  platform::Place place() const { return value_->place(); }

Q
qijun 已提交
58 59 60
  const Tensor& value() const { return *value_; }

  Tensor* mutable_value() { return value_.get(); }
Q
qijun 已提交
61 62 63 64 65

  int64_t height() const { return height_; }

  void set_height(int64_t height) { height_ = height; }

Q
qijun 已提交
66
  const Vector<int64_t>& rows() const { return rows_; }
Q
qijun 已提交
67

Q
QI JUN 已提交
68 69
  Vector<int64_t>* mutable_rows() { return &rows_; }

Q
qijun 已提交
70
  void set_rows(const Vector<int64_t>& rows) { rows_ = rows; }
Q
qijun 已提交
71

Y
Yancey1989 已提交
72 73 74 75
  /*
   * @brief wheter has the specified key in the table.
   *
   * @return true if the key is exists.
Q
qiaolongfei 已提交
76
   */
Y
Yancey1989 已提交
77 78 79
  bool HasKey(int64_t key) const;

  /*
Y
Yancey1989 已提交
80
   * @brief Get value by the key list, if the
Y
Yancey1989 已提交
81
   *
Y
Yancey1989 已提交
82 83
   * @return a list of pair which contains the non-exists key and the index in
   * the value
Y
Yancey1989 已提交
84
   */
Y
Yancey1989 已提交
85 86
  std::vector<std::pair<int64_t, int64_t>> Get(std::vector<int64_t> keys,
                                               framework::Tensor* value) const;
Y
Yancey1989 已提交
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111

  /*
   * @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 key) const {
    auto it = std::find(rows_.begin(), rows_.end(), key);
    if (it == rows_.end()) {
      return static_cast<int64_t>(-1);
    }
Q
qiaolongfei 已提交
112 113 114
    return static_cast<int64_t>(std::distance(rows_.begin(), it));
  }

Q
qijun 已提交
115 116 117 118 119 120 121
  DDim GetCompleteDims() const {
    std::vector<int64_t> dims = vectorize(value_->dims());
    dims[0] = height_;
    return make_ddim(dims);
  }

 private:
Q
qijun 已提交
122
  // Notice: rows can be duplicate. We can have {0, 4, 7, 0, 5, 7, 9} here.
123
  // SelectedRows are simply concated when adding together. Until a
Q
qijun 已提交
124
  // SelectedRows add a Tensor, will the duplicate rows be handled.
Q
qijun 已提交
125
  Vector<int64_t> rows_;
Q
qijun 已提交
126 127
  std::unique_ptr<Tensor> value_{nullptr};
  int64_t height_;
128
  std::mutex auto_grown_mutex_;
Q
qijun 已提交
129 130
};

131 132 133 134 135 136 137
/*
 * Serialize/Desiralize SelectedRows to std::ostream
 * You can pass ofstream or ostringstream to serilize to file
 * or to a in memory string. GPU tensor will be copied to CPU.
 */
void SerializeToStream(std::ostream& os, const SelectedRows& selected_rows,
                       const platform::DeviceContext& dev_ctx);
Y
Yancey 已提交
138 139
void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows,
                           const platform::DeviceContext& dev_ctx);
140

Q
qijun 已提交
141 142
}  // namespace framework
}  // namespace paddle