Modified from: https://arxiv.org/pdf/2012.12395.pdf Args: config (object): The Config object. layer (int, optional): Collaborate on which layer. Defaults to 3. in_channels (int, optional): The input channels. Defaults to 13. kd_flag (bool, optional): Whe
| 281 | |
| 282 | |
| 283 | class CNNNet(nn.Module): |
| 284 | """ |
| 285 | Modified from: |
| 286 | https://arxiv.org/pdf/2012.12395.pdf |
| 287 | |
| 288 | Args: |
| 289 | config (object): The Config object. |
| 290 | layer (int, optional): Collaborate on which layer. Defaults to 3. |
| 291 | in_channels (int, optional): The input channels. Defaults to 13. |
| 292 | kd_flag (bool, optional): Whether to use knowledge distillation (for DiscoNet to ues). Defaults to True. |
| 293 | num_agent (int, optional): The number of agents (including RSU). Defaults to 5. |
| 294 | """ |
| 295 | |
| 296 | def __init__( |
| 297 | self, |
| 298 | config, |
| 299 | layer=3, |
| 300 | in_channels=13, |
| 301 | kd_flag=True, |
| 302 | num_agent=5, |
| 303 | compress_level=0, |
| 304 | train_completion=True, |
| 305 | ): |
| 306 | super().__init__() |
| 307 | self.config = config |
| 308 | self.kd_flag = kd_flag |
| 309 | self.in_channels = in_channels |
| 310 | self.num_agent = num_agent |
| 311 | self.train_completion = train_completion |
| 312 | self.stpn = STPN_KD(config.map_dims[2], compress_level, train_completion) |
| 313 | |
| 314 | def get_feature_maps_size(self, feature_maps: tuple): |
| 315 | size = list(feature_maps.shape) |
| 316 | # NOTE: batch size will change the shape[0]. We need to manually set it to 1. |
| 317 | size[0] = 1 |
| 318 | size = tuple(size) |
| 319 | return size |
| 320 | |
| 321 | # get feat maps for each agent [10 512 32 32] -> [2 5 512 32 32] |
| 322 | def build_feature_list(self, batch_size: int, feat_maps: dict) -> list: |
| 323 | feature_map = {} |
| 324 | # [5,256,32,32] |
| 325 | feature_list = [] |
| 326 | |
| 327 | for i in range(self.num_agent): |
| 328 | feature_map[i] = torch.unsqueeze(feat_maps[batch_size * i:batch_size * (i + 1)], 1) |
| 329 | # feature_map[i]: [B,1,256,32,32] |
| 330 | feature_list.append(feature_map[i]) |
| 331 | |
| 332 | return feature_list |
| 333 | |
| 334 | # [2 5 512 16 16] [batch, agent, channel, height, width] |
| 335 | @staticmethod |
| 336 | def build_local_communication_matrix(feature_list: list): |
| 337 | return torch.cat(tuple(feature_list), 1) |
| 338 | |
| 339 | @staticmethod |
| 340 | # FIXME: rename 'j' |