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