Caffe2 - C++ API
A deep learning, cross platform ML framework
tt_linear_op.cc
1 #include "caffe2/operators/tt_linear_op.h"
2 
3 namespace caffe2 {
4 namespace {
5 
6 REGISTER_CPU_OPERATOR(TT, TTLinearOp<float, CPUContext>);
7 REGISTER_CPU_OPERATOR(TTLinearGradient, TTLinearGradientOp<float, CPUContext>);
8 
9 // The TT-layer serves as a low-rank decomposition of a fully connected layer.
10 // The inputs are the same as to an FC layer, but the number of the parameters
11 // are greatly reduced.
12 OPERATOR_SCHEMA(TT)
13  .NumInputs(3)
14  .NumOutputs(1)
15  .SetDoc(R"DOC(
16 The TT-layer serves as a low-rank decomposition of a fully connected layer. The
17 inputs are the same as to a fully connected layer, but the number of parameters
18 are greatly reduced and forward computation time can be drastically reduced
19 especially for layers with large weight matrices. The multiplication is computed
20 as a product of the input vector with each of the cores that make up the TT
21 layer. Given the input sizes (inp_sizes), output sizes(out_sizes), and the ranks
22 of each of the cores (tt_ranks), the ith core will have size:
23 
24  inp_sizes[i] * tt_ranks[i] * tt_ranks[i + 1] * out_sizes[i].
25 
26 The complexity of the computation is dictated by the sizes of inp_sizes,
27 out_sizes, and tt_ranks, where there is the trade off between accuracy of the
28 low-rank decomposition and the speed of the computation.
29 )DOC")
30  .Arg(
31  "inp_sizes",
32  "(int[]) Input sizes of cores. Indicates the input size of "
33  "the individual cores; the size of the input vector X must match the "
34  "product of the inp_sizes array.")
35  .Arg(
36  "out_sizes",
37  "(int[]) Output sizes of cores. Indicates the output size "
38  "of the individual cores; the size of the output vector Y must match "
39  "the product of the out_sizes array.")
40  .Arg(
41  "tt_ranks",
42  "(int[]) Ranks of cores. Indicates the ranks of the "
43  "individual cores; lower rank means larger compression, faster "
44  "computation but reduce accuracy.")
45  .Input(
46  0,
47  "X",
48  "Input tensor from previous layer with size (M x K), where "
49  "M is the batch size and K is the input size.")
50  .Input(1, "b", "1D blob containing the bias vector")
51  .Input(
52  2,
53  "cores",
54  "1D blob containing each individual cores with sizes "
55  "specified above.")
56  .Output(
57  0,
58  "Y",
59  "Output tensor from previous layer with size (M x N), "
60  "where M is the batch size and N is the output size.");
61 
62 OPERATOR_SCHEMA(TTLinearGradient);
63 
64 GRADIENT_NOT_IMPLEMENTED_YET(TT);
65 } // namespace
66 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13