| 43 | print("=> no checkpoint found at '{}'".format(path)) |
| 44 | |
| 45 | def step(self, data, num_agent, batch_size, loss=True): |
| 46 | bev = data["bev_seq"] |
| 47 | labels = data["labels"] |
| 48 | self.optimizer.zero_grad() |
| 49 | bev = bev.permute(0, 3, 1, 2).contiguous() |
| 50 | |
| 51 | if not self.com: |
| 52 | filtered_bev = [] |
| 53 | filtered_label = [] |
| 54 | for i in range(bev.size(0)): |
| 55 | if torch.sum(bev[i]) > 1e-4: |
| 56 | filtered_bev.append(bev[i]) |
| 57 | filtered_label.append(labels[i]) |
| 58 | bev = torch.stack(filtered_bev, 0) |
| 59 | labels = torch.stack(filtered_label, 0) |
| 60 | |
| 61 | if self.kd_flag: |
| 62 | data["bev_seq_teacher"] = ( |
| 63 | data["bev_seq_teacher"].permute(0, 3, 1, 2).contiguous() |
| 64 | ) |
| 65 | |
| 66 | if self.com: |
| 67 | print("communication happening...") |
| 68 | if self.kd_flag: |
| 69 | pred, x9, x8, x7, x6, x5, fused_layer = self.model( |
| 70 | bev, data["trans_matrices"], data["num_sensor"] |
| 71 | ) |
| 72 | elif self.config.flag.startswith("when2com") or self.config.flag.startswith( |
| 73 | "who2com" |
| 74 | ): |
| 75 | if self.config.split == "train": |
| 76 | pred = self.model( |
| 77 | bev, data["trans_matrices"], data["num_sensor"], training=True |
| 78 | ) |
| 79 | else: |
| 80 | pred = self.model( |
| 81 | bev, |
| 82 | data["trans_matrices"], |
| 83 | data["num_sensor"], |
| 84 | inference=self.config.inference, |
| 85 | training=False, |
| 86 | ) |
| 87 | else: |
| 88 | pred = self.model(bev, data["trans_matrices"], data["num_sensor"]) |
| 89 | else: |
| 90 | pred = self.model(bev) |
| 91 | |
| 92 | if self.com: |
| 93 | filtered_pred = [] |
| 94 | filtered_label = [] |
| 95 | for i in range(bev.size(0)): |
| 96 | if torch.sum(bev[i]) > 1e-4: |
| 97 | filtered_pred.append(pred[i]) |
| 98 | filtered_label.append(labels[i]) |
| 99 | pred = torch.stack(filtered_pred, 0) |
| 100 | labels = torch.stack(filtered_label, 0) |
| 101 | if not loss: |
| 102 | return pred, labels |