mxnet
kvstore.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
27 #ifndef MXNET_CPP_KVSTORE_H_
28 #define MXNET_CPP_KVSTORE_H_
29 
30 #include <string>
31 #include <vector>
32 #include "mxnet-cpp/ndarray.h"
33 
34 namespace mxnet {
35 namespace cpp {
36 
37 class KVStore {
38  public:
39  static void SetType(const std::string& type);
40  static void RunServer();
41  static void Init(int key, const NDArray& val);
42  static void Init(const std::vector<int>& keys, const std::vector<NDArray>& vals);
43  static void Push(int key, const NDArray& val, int priority = 0);
44  static void Push(const std::vector<int>& keys,
45  const std::vector<NDArray>& vals, int priority = 0);
46  static void Pull(int key, NDArray* out, int priority = 0);
47  static void Pull(const std::vector<int>& keys, std::vector<NDArray>* outs, int priority = 0);
48  // TODO(lx): put lr in optimizer or not?
49  static void SetOptimizer(std::unique_ptr<Optimizer> optimizer, bool local = false);
50  static std::string GetType();
51  static int GetRank();
52  static int GetNumWorkers();
53  static void Barrier();
54  static std::string GetRole();
55 
56  private:
57  KVStore();
58  static KVStoreHandle& get_handle();
59  static std::unique_ptr<Optimizer>& get_optimizer();
60  static KVStore*& get_kvstore();
61  static void Controller(int head, const char* body, void* controller_handle);
62  static void Updater(int key, NDArrayHandle recv, NDArrayHandle local, void* handle_);
63 };
64 
65 } // namespace cpp
66 } // namespace mxnet
67 
68 #endif // MXNET_CPP_KVSTORE_H_
static std::string GetRole()
static void Pull(int key, NDArray *out, int priority=0)
static void Init(int key, const NDArray &val)
static void Barrier()
static void SetType(const std::string &type)
namespace of mxnet
Definition: base.h:118
static void SetOptimizer(std::unique_ptr< Optimizer > optimizer, bool local=false)
NDArray interface.
Definition: ndarray.h:121
static void Push(int key, const NDArray &val, int priority=0)
void * KVStoreHandle
handle to KVStore
Definition: c_api.h:85
void * NDArrayHandle
handle to NDArray
Definition: c_api.h:67
static int GetNumWorkers()
static std::string GetType()
static int GetRank()
static void RunServer()
Definition: kvstore.h:37