(self, module_name)
| 80 | return loga |
| 81 | |
| 82 | def initialize_one_module(self, module_name): |
| 83 | default_mean = 10 |
| 84 | if module_name == "intermediate": |
| 85 | self.intermediate_loga = self.initialize_parameters( |
| 86 | self.intermediate_size, self.num_hidden_layers, mean=default_mean) |
| 87 | self.add_one_module( |
| 88 | self.intermediate_loga, type_name="intermediate", |
| 89 | parameter_per_dim=self.params_per_intermediate_dim, size=self.intermediate_size, |
| 90 | shape=[self.num_hidden_layers, 1, 1, self.intermediate_size] |
| 91 | ) |
| 92 | self.prunable_model_size += self.params_per_mlp_layer * self.num_hidden_layers |
| 93 | elif module_name == "heads": |
| 94 | self.heads_loga = self.initialize_parameters( |
| 95 | self.num_attention_heads, self.num_hidden_layers, mean=default_mean) |
| 96 | self.add_one_module( |
| 97 | self.heads_loga, type_name="heads", |
| 98 | parameter_per_dim=self.params_per_head, size=self.num_attention_heads, |
| 99 | shape=[self.num_hidden_layers, 1, |
| 100 | self.num_attention_heads, 1, 1] |
| 101 | ) |
| 102 | self.prunable_model_size += self.params_per_head * \ |
| 103 | self.num_hidden_layers * self.num_attention_heads |
| 104 | elif module_name == "hidden": |
| 105 | self.hidden_loga = self.initialize_parameters( |
| 106 | self.hidden_size, mean=default_mean) |
| 107 | self.add_one_module( |
| 108 | self.hidden_loga, type_name="hidden", |
| 109 | parameter_per_dim=self.hidden_size * 4 + self.hidden_size * 4 * 2, |
| 110 | size=self.hidden_size, shape=[self.hidden_size] |
| 111 | ) |
| 112 | elif module_name == "layer": |
| 113 | self.ffn_loga = self.initialize_parameters( |
| 114 | self.num_hidden_layers, mean=default_mean) |
| 115 | self.add_one_module( |
| 116 | self.ffn_loga, type_name="ffn", |
| 117 | parameter_per_dim=self.params_per_mlp_layer, size=1, |
| 118 | shape=[self.num_hidden_layers] |
| 119 | ) |
| 120 | self.mha_loga = self.initialize_parameters( |
| 121 | self.num_hidden_layers, mean=default_mean) |
| 122 | self.add_one_module( |
| 123 | self.mha_loga, type_name="mha", |
| 124 | parameter_per_dim=self.params_per_head * self.num_attention_heads, size=1, |
| 125 | shape=[self.num_hidden_layers] |
| 126 | ) |
| 127 | |
| 128 | # ! init the z_logas |
| 129 | def add_one_module(self, z_loga, type_name, parameter_per_dim, size, shape): |
no test coverage detected