Caffe2 - Python API
A deep learning, cross platform ML framework
init_baseline.py
1 """Script to generate baseline values from PyTorch initialization algorithms"""
2 
3 import sys
4 import torch
5 
6 HEADER = """
7 #include <torch/types.h>
8 
9 #include <vector>
10 
11 namespace expected_parameters {
12 """
13 
14 FOOTER = "} // namespace expected_parameters"
15 
16 PARAMETERS = "inline std::vector<std::vector<torch::Tensor>> {}() {{"
17 
18 INITIALIZERS = {
19  "Xavier_Uniform": lambda w: torch.nn.init.xavier_uniform(w),
20  "Xavier_Normal": lambda w: torch.nn.init.xavier_normal(w),
21  "Kaiming_Normal": lambda w: torch.nn.init.kaiming_normal(w),
22  "Kaiming_Uniform": lambda w: torch.nn.init.kaiming_uniform(w)
23 }
24 
25 
26 def emit(initializer_parameter_map):
27  # Don't write generated with an @ in front, else this file is recognized as generated.
28  print("// @{} from {}".format('generated', __file__))
29  print(HEADER)
30  for initializer_name, weights in initializer_parameter_map.items():
31  print(PARAMETERS.format(initializer_name))
32  print(" return {")
33  for sample in weights:
34  print(" {")
35  for parameter in sample:
36  parameter_values = "{{{}}}".format(", ".join(map(str, parameter)))
37  print(" torch::tensor({}),".format(parameter_values))
38  print(" },")
39  print(" };")
40  print("}\n")
41  print(FOOTER)
42 
43 
44 def run(initializer):
45  torch.manual_seed(0)
46 
47  layer1 = torch.nn.Linear(7, 15)
48  INITIALIZERS[initializer](layer1.weight)
49 
50  layer2 = torch.nn.Linear(15, 15)
51  INITIALIZERS[initializer](layer2.weight)
52 
53  layer3 = torch.nn.Linear(15, 2)
54  INITIALIZERS[initializer](layer3.weight)
55 
56  weight1 = layer1.weight.data.numpy()
57  weight2 = layer2.weight.data.numpy()
58  weight3 = layer3.weight.data.numpy()
59 
60  return [weight1, weight2, weight3]
61 
62 
63 def main():
64  initializer_parameter_map = {}
65  for initializer in INITIALIZERS.keys():
66  sys.stderr.write('Evaluating {} ...\n'.format(initializer))
67  initializer_parameter_map[initializer] = run(initializer)
68 
69  emit(initializer_parameter_map)
70 
71 if __name__ == "__main__":
72  main()
xavier_uniform
Definition: init.py:429
kaiming_normal
Definition: init.py:432
kaiming_uniform
Definition: init.py:431