(self, model, inputs, return_outputs=False)
| 124 | return mse_loss(student_logits, teacher_logits) |
| 125 | |
| 126 | def compute_loss(self, model, inputs, return_outputs=False): |
| 127 | with torch.no_grad(): |
| 128 | teacher_outputs = self.teacher_model( |
| 129 | **inputs |
| 130 | # **inputs, output_hidden_states=True, output_attentions=True |
| 131 | ) |
| 132 | teacher_logits = teacher_outputs.get("logits") |
| 133 | del teacher_outputs |
| 134 | |
| 135 | # forward pass |
| 136 | student_outputs = model(**inputs) |
| 137 | # get attributes |
| 138 | student_logits = student_outputs.get("logits") |
| 139 | |
| 140 | if not return_outputs: |
| 141 | del student_outputs |
| 142 | |
| 143 | # torch.save(student_logits, "/root/model/acr_duda/code/rej_analysis/sft_student_logits.pt") |
| 144 | # torch.save(teacher_logits, "/root/model/acr_duda/code/rej_analysis/sft_teacher_logits.pt") |
| 145 | # torch.save(inputs, "/root/model/acr_duda/code/rej_analysis/sft_inputs.pt") |
| 146 | # raise 1 |
| 147 | |
| 148 | kd_loss = 0.0 |
| 149 | size_average = True |
| 150 | if model.kd_loss_scale > 0.0: |
| 151 | if self.loss_type == "reverse": |
| 152 | kd_loss = self.re_loss(inputs['labels'], student_logits, teacher_logits) |
| 153 | elif self.loss_type == "forward": |
| 154 | kd_loss = self.ce_loss(inputs['labels'], student_logits, teacher_logits) |
| 155 | elif self.loss_type == "tlsd": |
| 156 | kd_loss = self.TLSD_loss(inputs['labels'], student_logits, teacher_logits) |
| 157 | elif self.loss_type == "cakld": |
| 158 | kd_loss = self.cakld_loss(inputs['labels'], student_logits, teacher_logits, self.mean_prob) |
| 159 | elif self.loss_type == "jsd": |
| 160 | kd_loss = self.jsd_loss(inputs['labels'], student_logits, teacher_logits, 0.5) |
| 161 | |
| 162 | del teacher_logits |
| 163 | del student_logits |
| 164 | |
| 165 | tok_loss = model.kd_loss_scale * kd_loss |
| 166 | return (tok_loss, student_outputs) if return_outputs else tok_loss |
| 167 | |
| 168 |
nothing calls this directly
no test coverage detected