(model, layer_idx)
| 24 | |
| 25 | # Extract BERT Layer |
| 26 | def extract_bert_layer(model, layer_idx): |
| 27 | desired_layer = '' |
| 28 | for _, submodule in model.named_children(): |
| 29 | for n, s in submodule.named_children(): |
| 30 | if n == 'layers' or n == 'resblocks': |
| 31 | for n2, s2 in s.named_children(): |
| 32 | if n2 == str(layer_idx): |
| 33 | desired_layer = s2 |
| 34 | return desired_layer |
| 35 | |
| 36 | # Get an estimator for the compression term |
| 37 | def get_compression_estimator(var, layer, features): |
no outgoing calls
no test coverage detected