Skip to content

Commit 86eca9f

Browse files
min-jean-chomin-jean-cho
committed
added jit frozen grph lin-bn folding
1 parent d8b971e commit 86eca9f

File tree

9 files changed

+250
-0
lines changed

9 files changed

+250
-0
lines changed

‎build_variables.bzl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,10 +273,12 @@ core_sources_full_mobile_no_backend_interface = [
273273
"torch/csrc/jit/passes/remove_mutation.cpp",
274274
"torch/csrc/jit/passes/prepack_folding.cpp",
275275
"torch/csrc/jit/passes/fold_conv_bn.cpp",
276+
"torch/csrc/jit/passes/fold_linear_bn.cpp",
276277
"torch/csrc/jit/passes/dbr_quantization/remove_redundant_aliases.cpp",
277278
"torch/csrc/jit/passes/frozen_concat_linear.cpp",
278279
"torch/csrc/jit/passes/frozen_conv_add_relu_fusion.cpp",
279280
"torch/csrc/jit/passes/frozen_conv_folding.cpp",
281+
"torch/csrc/jit/passes/frozen_linear_folding.cpp",
280282
"torch/csrc/jit/passes/frozen_linear_transpose.cpp",
281283
"torch/csrc/jit/passes/frozen_ops_to_mkldnn.cpp",
282284
"torch/csrc/jit/passes/frozen_graph_optimizations.cpp",

‎test/jit/test_freezing.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2223,6 +2223,59 @@ def forward(self, x):
22232223
FileCheck().check("conv").check_not("aten::batch_norm").run(traced_model.graph)
22242224
FileCheck().check("conv").check_not("aten::add").run(traced_model.graph)
22252225

2226+
def test_linear_bn_folding(self):
2227+
module_pairs = [(nn.Linear, nn.BatchNorm1d), (nn.Linear, nn.BatchNorm2d), (nn.Linear, nn.BatchNorm3d)]
2228+
use_tracing = [True, False]
2229+
bn_running_stats = [True, False]
2230+
2231+
for modules, tracing, track_stats in product(module_pairs, use_tracing, bn_running_stats):
2232+
class LinearBN(torch.nn.Module):
2233+
def __init__(self, in_features, out_features):
2234+
super(LinearBN, self).__init__()
2235+
self.linear = modules[0](in_features, out_features)
2236+
self.bn = modules[1](out_features, eps=0.001, track_running_stats=track_stats)
2237+
2238+
def forward(self, x):
2239+
x = self.linear(x)
2240+
return self.bn(x)
2241+
2242+
mod_eager = LinearBN(32, 32).eval()
2243+
2244+
inps = [3, 32]
2245+
if modules[1] == nn.BatchNorm2d:
2246+
inps.append(inps[-1])
2247+
inps.append(inps[-1])
2248+
if modules[1] == nn.BatchNorm3d:
2249+
inps.append(inps[-1])
2250+
inps.append(inps[-1])
2251+
inps.append(inps[-1])
2252+
2253+
inp = torch.rand(inps)
2254+
2255+
if tracing:
2256+
scripted_mod = torch.jit.trace(mod_eager, (inp))
2257+
else:
2258+
scripted_mod = torch.jit.script(mod_eager)
2259+
2260+
self.run_pass("inline", scripted_mod.graph)
2261+
self.run_pass("peephole", scripted_mod.graph)
2262+
self.run_pass("constant_propagation", scripted_mod.graph)
2263+
2264+
FileCheck().check("linear").check("batch").run(scripted_mod.graph)
2265+
# successfully no-ops with non-const inputs
2266+
self.run_pass("fold_frozen_linear_bn", scripted_mod.graph)
2267+
FileCheck().check("linear").check("aten::batch_norm").run(scripted_mod.graph)
2268+
2269+
scripted_mod = torch.jit.freeze(scripted_mod)
2270+
self.run_pass("fold_frozen_linear_bn", scripted_mod.graph)
2271+
if track_stats:
2272+
FileCheck().check("linear").check_not("aten::batch_norm").run(scripted_mod.graph)
2273+
else:
2274+
FileCheck().check("linear").check("aten::batch_norm").run(scripted_mod.graph)
2275+
2276+
self.assertEqual(mod_eager(inp), scripted_mod(inp))
2277+
self.assertEqual(mod_eager(inp), scripted_mod(inp))
2278+
22262279
@unittest.skipIf(not TEST_CUDA, "Optimization currently only run for GPU")
22272280
def test_linear_concat(self):
22282281
out_dimms = [[5, 10], [1, 5]]

‎torch/csrc/jit/passes/fold_linear_bn.cpp‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#include <torch/csrc/jit/passes/fold_linear_bn.h>
2+
3+
#include <ATen/TensorOperators.h>
4+
5+
#ifndef AT_PER_OPERATOR_HEADERS
6+
#include <ATen/Functions.h>
7+
#else
8+
#include <ATen/ops/rsqrt.h>
9+
#endif
10+
11+
namespace torch {
12+
namespace jit {
13+
14+
std::tuple<at::Tensor, at::Tensor> computeUpdatedLinearWeightAndBias(
15+
const LinearBNParameters& p) {
16+
at::Tensor bn_scale = p.bn_w * at::rsqrt(p.bn_rv + p.bn_eps);
17+
at::Tensor fused_w = p.linear_w * bn_scale.unsqueeze(-1);
18+
at::Tensor fused_b = (p.linear_b - p.bn_rm) * bn_scale + p.bn_b;
19+
20+
auto linear_w_dtype = p.linear_w.dtype();
21+
auto linear_b_dtype = p.linear_b.dtype();
22+
23+
return std::make_tuple(
24+
fused_w.to(linear_w_dtype), fused_b.to(linear_b_dtype));
25+
}
26+
27+
} // namespace jit
28+
} // namespace torch

‎torch/csrc/jit/passes/fold_linear_bn.h‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
#pragma once
2+
3+
#include <torch/csrc/jit/api/module.h>
4+
5+
namespace torch {
6+
namespace jit {
7+
8+
struct TORCH_API LinearBNParameters {
9+
at::Tensor linear_w;
10+
at::Tensor linear_b;
11+
at::Tensor bn_rm;
12+
at::Tensor bn_rv;
13+
double bn_eps = 0.0;
14+
at::Tensor bn_w;
15+
at::Tensor bn_b;
16+
};
17+
18+
/**
19+
* Given the current weight and bias tensors of a Linear module and parameters
20+
* of the BatchNorm module we're folding with, compute the updated values
21+
* for the weight and bias.
22+
*
23+
* The function is basically copied from torch/nn/utils/fusion.py
24+
*/
25+
TORCH_API std::tuple<at::Tensor, at::Tensor> computeUpdatedLinearWeightAndBias(
26+
const LinearBNParameters& p);
27+
28+
} // namespace jit
29+
} // namespace torch

‎torch/csrc/jit/passes/frozen_graph_optimizations.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include <torch/csrc/jit/passes/frozen_concat_linear.h>
55
#include <torch/csrc/jit/passes/frozen_conv_folding.h>
66
#include <torch/csrc/jit/passes/frozen_graph_optimizations.h>
7+
#include <torch/csrc/jit/passes/frozen_linear_folding.h>
78
#include <torch/csrc/jit/passes/remove_dropout.h>
89
#include <torch/csrc/jit/runtime/graph_executor.h>
910
#include <torch/csrc/utils/memory.h>
@@ -24,6 +25,7 @@ void OptimizeFrozenGraph(
2425
changed |= FoldFrozenConvBatchnorm(graph);
2526
changed |= FoldFrozenConvAddOrSub(graph);
2627
changed |= FoldFrozenConvMulOrDiv(graph);
28+
changed |= FoldFrozenLinearBatchnorm(graph);
2729
} while (changed);
2830
}
2931
}

‎torch/csrc/jit/passes/frozen_graph_optimizations.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
* - FoldFrozenConvBatchnorm
99
* - FoldFrozenConvAddOrSub
1010
* - FoldFrozenConvMulOrDiv
11+
* - FoldFrozenLinearBatchnorm
1112
*/
1213

1314
namespace torch {

‎torch/csrc/jit/passes/frozen_linear_folding.cpp‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
#include <torch/csrc/jit/ir/constants.h>
2+
#include <torch/csrc/jit/ir/ir.h>
3+
#include <torch/csrc/jit/passes/dead_code_elimination.h>
4+
#include <torch/csrc/jit/passes/fold_linear_bn.h>
5+
#include <torch/csrc/jit/passes/frozen_linear_folding.h>
6+
#include <torch/csrc/jit/passes/utils/optimization_utils.h>
7+
8+
#ifndef AT_PER_OPERATOR_HEADERS
9+
#include <ATen/Functions.h>
10+
#else
11+
#include <ATen/ops/ones_like.h>
12+
#include <ATen/ops/zeros_like.h>
13+
#endif
14+
15+
namespace torch {
16+
namespace jit {
17+
18+
namespace {
19+
20+
using Tensor = at::Tensor;
21+
22+
bool supportedLinearNode(Node* n) {
23+
if (n->kind() == aten::linear) {
24+
return true;
25+
} else {
26+
return false;
27+
}
28+
}
29+
30+
bool FoldFrozenLinearBatchnorm(Block* b) {
31+
bool graph_modified = false;
32+
for (Node* n : b->nodes()) {
33+
for (Block* block : n->blocks()) {
34+
graph_modified |= FoldFrozenLinearBatchnorm(block);
35+
}
36+
37+
if (n->kind() == aten::batch_norm &&
38+
supportedLinearNode(n->inputs().at(0)->node())) {
39+
auto linear = n->inputs().at(0)->node();
40+
auto bn = n;
41+
42+
if (nonConstantParameters(linear) || nonConstantParameters(bn)) {
43+
continue;
44+
}
45+
46+
auto bn_rm_ivalue = bn->namedInput("running_mean");
47+
auto bn_rv_ivalue = bn->namedInput("running_var");
48+
49+
// check running_mean and running_var has value, if they are
50+
// None(track_running_stats=False), skiping the folding path.
51+
if (bn_rm_ivalue->type() == NoneType::get() &&
52+
bn_rv_ivalue->type() == NoneType::get()) {
53+
continue;
54+
}
55+
56+
auto bn_rm = constant_as<Tensor>(bn->namedInput("running_mean")).value();
57+
auto bn_rv = constant_as<Tensor>(bn->namedInput("running_var")).value();
58+
auto bn_eps = constant_as<double>(bn->namedInput("eps")).value();
59+
auto linear_w = constant_as<Tensor>(linear->namedInput("weight")).value();
60+
61+
// implementation taken from torch/nn/utils/fusion.py
62+
Tensor linear_b;
63+
if (linear->namedInput("bias")->type() == NoneType::get()) {
64+
linear_b = at::zeros_like(bn_rm);
65+
} else {
66+
linear_b = constant_as<Tensor>(linear->namedInput("bias")).value();
67+
}
68+
Tensor bn_w;
69+
if (bn->namedInput("weight")->type() == NoneType::get()) {
70+
bn_w = at::ones_like(bn_rm);
71+
} else {
72+
bn_w = constant_as<Tensor>(bn->namedInput("weight")).value();
73+
}
74+
Tensor bn_b;
75+
if (n->namedInput("bias")->type() == NoneType::get()) {
76+
bn_b = at::zeros_like(bn_rm);
77+
} else {
78+
bn_b = constant_as<Tensor>(bn->namedInput("bias")).value();
79+
}
80+
81+
LinearBNParameters params;
82+
params.linear_w = linear_w;
83+
params.linear_b = linear_b;
84+
params.bn_rm = bn_rm;
85+
params.bn_rv = bn_rv;
86+
params.bn_eps = bn_eps;
87+
params.bn_w = bn_w;
88+
params.bn_b = bn_b;
89+
std::tuple<Tensor, Tensor> out =
90+
computeUpdatedLinearWeightAndBias(params);
91+
WithInsertPoint guard(linear);
92+
auto fused_linear_w = b->owningGraph()->insertConstant(std::get<0>(out));
93+
auto fused_linear_b = b->owningGraph()->insertConstant(std::get<1>(out));
94+
auto linear_w_value = linear->namedInput("weight");
95+
auto linear_b_value = linear->namedInput("bias");
96+
97+
fused_linear_w->setDebugName(linear_w_value->debugName() + "_fused_bn");
98+
fused_linear_b->setDebugName(linear_b_value->debugName() + "_fused_bn");
99+
100+
linear->replaceInputWith(linear_w_value, fused_linear_w);
101+
linear->replaceInputWith(linear_b_value, fused_linear_b);
102+
103+
bn->output()->replaceAllUsesWith(linear->output());
104+
graph_modified = true;
105+
}
106+
}
107+
return graph_modified;
108+
}
109+
110+
} // namespace
111+
112+
bool FoldFrozenLinearBatchnorm(std::shared_ptr<Graph>& graph) {
113+
bool graph_modified = FoldFrozenLinearBatchnorm(graph->block());
114+
EliminateDeadCode(graph);
115+
return graph_modified;
116+
}
117+
118+
} // namespace jit
119+
} // namespace torch

‎torch/csrc/jit/passes/frozen_linear_folding.h‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
#pragma once
2+
3+
#include <torch/csrc/jit/ir/ir.h>
4+
5+
namespace torch {
6+
namespace jit {
7+
8+
// Fuses Linear -> BatchNormNd into a single Linear by
9+
// folding batchnorm weights into linear weights.
10+
// This pass only works on Frozen Graphs; otherwise it is a No-Op.
11+
TORCH_API bool FoldFrozenLinearBatchnorm(std::shared_ptr<Graph>& graph);
12+
13+
} // namespace jit
14+
} // namespace torch

‎torch/csrc/jit/python/init.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
#include <torch/csrc/jit/passes/frozen_conv_add_relu_fusion.h>
4141
#include <torch/csrc/jit/passes/frozen_conv_folding.h>
4242
#include <torch/csrc/jit/passes/frozen_graph_optimizations.h>
43+
#include <torch/csrc/jit/passes/frozen_linear_folding.h>
4344
#include <torch/csrc/jit/passes/frozen_linear_transpose.h>
4445
#include <torch/csrc/jit/passes/frozen_ops_to_mkldnn.h>
4546
#include <torch/csrc/jit/passes/fuse_linear.h>
@@ -617,6 +618,7 @@ void initJITBindings(PyObject* module) {
617618
.def("_jit_pass_fold_frozen_conv_bn", &FoldFrozenConvBatchnorm)
618619
.def("_jit_pass_fold_frozen_conv_add_or_sub", &FoldFrozenConvAddOrSub)
619620
.def("_jit_pass_fold_frozen_conv_mul_or_div", &FoldFrozenConvMulOrDiv)
621+
.def("_jit_pass_fold_frozen_linear_bn", &FoldFrozenLinearBatchnorm)
620622
.def("_jit_pass_convert_frozen_ops_to_mkldnn", &ConvertFrozenOpsToMKLDNN)
621623
.def("_jit_pass_fuse_frozen_conv_add_relu", &FuseFrozenConvAddRelu)
622624
.def("_jit_pass_transpose_frozen_linear", &FrozenLinearTranspose)

0 commit comments

Comments
 (0)