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