| 27 | |
| 28 | |
| 29 | class Classification(MegatronModule): |
| 30 | |
| 31 | def __init__(self, num_classes, num_tokentypes=2): |
| 32 | super(Classification, self).__init__() |
| 33 | args = get_args() |
| 34 | |
| 35 | self.num_classes = num_classes |
| 36 | init_method = init_method_normal(args.init_method_std) |
| 37 | |
| 38 | self.language_model, self._language_model_key = get_language_model( |
| 39 | attention_mask_func=bert_attention_mask_func, |
| 40 | num_tokentypes=num_tokentypes, |
| 41 | add_pooler=True, |
| 42 | init_method=init_method, |
| 43 | scaled_init_method=scaled_init_method_normal(args.init_method_std, |
| 44 | args.num_layers)) |
| 45 | |
| 46 | # Multi-choice head. |
| 47 | self.classification_dropout = torch.nn.Dropout(args.hidden_dropout) |
| 48 | self.classification_head = get_linear_layer(args.hidden_size, |
| 49 | self.num_classes, |
| 50 | init_method) |
| 51 | self._classification_head_key = 'classification_head' |
| 52 | |
| 53 | def forward(self, input_ids, attention_mask, tokentype_ids): |
| 54 | |
| 55 | extended_attention_mask = bert_extended_attention_mask( |
| 56 | attention_mask, next(self.language_model.parameters()).dtype) |
| 57 | position_ids = bert_position_ids(input_ids) |
| 58 | |
| 59 | _, pooled_output = self.language_model(input_ids, |
| 60 | position_ids, |
| 61 | extended_attention_mask, |
| 62 | tokentype_ids=tokentype_ids) |
| 63 | |
| 64 | # Output. |
| 65 | classification_output = self.classification_dropout(pooled_output) |
| 66 | classification_logits = self.classification_head(classification_output) |
| 67 | |
| 68 | # Reshape back to separate choices. |
| 69 | classification_logits = classification_logits.view(-1, self.num_classes) |
| 70 | |
| 71 | return classification_logits |
| 72 | |
| 73 | def state_dict_for_save_checkpoint(self, destination=None, prefix='', |
| 74 | keep_vars=False): |
| 75 | """For easy load when model is combined with other heads, |
| 76 | add an extra key.""" |
| 77 | |
| 78 | state_dict_ = {} |
| 79 | state_dict_[self._language_model_key] \ |
| 80 | = self.language_model.state_dict_for_save_checkpoint( |
| 81 | destination, prefix, keep_vars) |
| 82 | state_dict_[self._classification_head_key] \ |
| 83 | = self.classification_head.state_dict( |
| 84 | destination, prefix, keep_vars) |
| 85 | return state_dict_ |
| 86 |
no outgoing calls
no test coverage detected