RLPack
 
Loading...
Searching...
No Matches
C_GradAccumulator.h
Go to the documentation of this file.
1//
2// Created by Kartik Rajeshwaran on 2022-12-14.
3//
4
5#ifndef RLPACK_BINARIES_GRAD_ACCUMULATOR_CGRADACCUMULATOR_H_
6#define RLPACK_BINARIES_GRAD_ACCUMULATOR_CGRADACCUMULATOR_H_
7
8#include <pybind11/pybind11.h>
9#include <torch/extension.h>
10
11#include <string>
12#include <vector>
13
24public:
25 C_GradAccumulator(std::vector<std::string> &parameterKeys, int64_t boostStrapRounds);
27
28 void accumulate(std::map<std::string, torch::Tensor> &namedParameters);
29 std::map<std::string, torch::Tensor> mean_reduce();
30 std::map<std::string, torch::Tensor> sum_reduce();
31 std::map<std::string, torch::Tensor> get_item(int64_t index);
32 void set_item(int64_t index, std::map<std::string, torch::Tensor> &namedParameters);
33 void delete_item(int64_t index);
34 size_t size();
35 void clear();
36
37private:
41 std::vector<std::string> parameterKeys_;
43 std::vector<std::map<std::string, torch::Tensor>> namedParametersGrads_;
45 std::map<std::string, torch::Tensor> reducedParams_;
46};
53#endif//RLPACK_BINARIES_GRAD_ACCUMULATOR_CGRADACCUMULATOR_H_
Definition: C_GradAccumulator.h:23
void clear()
Definition: C_GradAccumulator.cpp:159
size_t size()
Definition: C_GradAccumulator.cpp:155
int64_t bootstrapRounds_
The number of boostrap rounds over which accumulation and reduction is to take place.
Definition: C_GradAccumulator.h:39
std::map< std::string, torch::Tensor > reducedParams_
The map to store final results of reduced parameters.
Definition: C_GradAccumulator.h:45
std::map< std::string, torch::Tensor > mean_reduce()
Definition: C_GradAccumulator.cpp:43
std::map< std::string, torch::Tensor > sum_reduce()
Definition: C_GradAccumulator.cpp:69
void accumulate(std::map< std::string, torch::Tensor > &namedParameters)
Definition: C_GradAccumulator.cpp:21
std::map< std::string, torch::Tensor > get_item(int64_t index)
Definition: C_GradAccumulator.cpp:95
void set_item(int64_t index, std::map< std::string, torch::Tensor > &namedParameters)
Definition: C_GradAccumulator.cpp:113
void delete_item(int64_t index)
Definition: C_GradAccumulator.cpp:139
std::vector< std::map< std::string, torch::Tensor > > namedParametersGrads_
The vector to accumulate the gradients.
Definition: C_GradAccumulator.h:43
std::vector< std::string > parameterKeys_
The parameter keys for the model for which gradient accumulation is being done.
Definition: C_GradAccumulator.h:41