| 121 | return pred, loss_data |
| 122 | |
| 123 | def get_kd_loss(self, batch_size, data, fused_layer, num_agent, x5, x6, x7): |
| 124 | if not self.kd_flag: |
| 125 | return 0 |
| 126 | |
| 127 | bev_seq_teacher = data["bev_seq_teacher"].type(torch.cuda.FloatTensor) |
| 128 | kd_weight = data["kd_weight"] |
| 129 | ( |
| 130 | logit_teacher, |
| 131 | x9_teacher, |
| 132 | x8_teacher, |
| 133 | x7_teacher, |
| 134 | x6_teacher, |
| 135 | x5_teacher, |
| 136 | x4_teacher, |
| 137 | ) = self.teacher(bev_seq_teacher) |
| 138 | kl_loss_mean = nn.KLDivLoss(size_average=True, reduce=True) |
| 139 | |
| 140 | target_x5 = x5_teacher.permute(0, 2, 3, 1).reshape( |
| 141 | num_agent * batch_size * 16 * 16, -1 |
| 142 | ) |
| 143 | student_x5 = x5.permute(0, 2, 3, 1).reshape( |
| 144 | num_agent * batch_size * 16 * 16, -1 |
| 145 | ) |
| 146 | kd_loss_x5 = kl_loss_mean( |
| 147 | F.log_softmax(student_x5, dim=1), F.softmax(target_x5, dim=1) |
| 148 | ) |
| 149 | |
| 150 | target_x6 = x6_teacher.permute(0, 2, 3, 1).reshape( |
| 151 | num_agent * batch_size * 32 * 32, -1 |
| 152 | ) |
| 153 | student_x6 = x6.permute(0, 2, 3, 1).reshape( |
| 154 | num_agent * batch_size * 32 * 32, -1 |
| 155 | ) |
| 156 | kd_loss_x6 = kl_loss_mean( |
| 157 | F.log_softmax(student_x6, dim=1), F.softmax(target_x6, dim=1) |
| 158 | ) |
| 159 | |
| 160 | target_x7 = x7_teacher.permute(0, 2, 3, 1).reshape( |
| 161 | num_agent * batch_size * 64 * 64, -1 |
| 162 | ) |
| 163 | student_x7 = x7.permute(0, 2, 3, 1).reshape( |
| 164 | num_agent * batch_size * 64 * 64, -1 |
| 165 | ) |
| 166 | kd_loss_x7 = kl_loss_mean( |
| 167 | F.log_softmax(student_x7, dim=1), F.softmax(target_x7, dim=1) |
| 168 | ) |
| 169 | |
| 170 | target_x4 = x4_teacher.permute(0, 2, 3, 1).reshape( |
| 171 | num_agent * batch_size * 32 * 32, -1 |
| 172 | ) |
| 173 | student_x4 = fused_layer.permute(0, 2, 3, 1).reshape( |
| 174 | num_agent * batch_size * 32 * 32, -1 |
| 175 | ) |
| 176 | kd_loss_fused_layer = kl_loss_mean( |
| 177 | F.log_softmax(student_x4, dim=1), F.softmax(target_x4, dim=1) |
| 178 | ) |
| 179 | |
| 180 | return kd_weight * (kd_loss_x5 + kd_loss_x6 + kd_loss_x7 + kd_loss_fused_layer) |