| 469 | torch.Tensor.cuda = new_cuda |
| 470 | |
| 471 | def add_task_specific(m, task_specific): |
| 472 | for name, param in m.named_parameters(): |
| 473 | param.task_specific = task_specific |
| 474 | param.backbone_specific = False |
| 475 | param.neck_specific = False |
| 476 | param.decoder_specific = False |
| 477 | if task_specific: |
| 478 | printlog('add param {} as task_specific'.format(name)) |
| 479 | |
| 480 | if not hasattr(torch.nn.Module, 'named_buffers'): |
| 481 | printlog('registering named_buffers for nn.Module at add_task_specific') |
| 482 | torch.nn.Module.named_buffers = named_buffers |
| 483 | |
| 484 | #m.cuda() # neccesary for broadcast in DistModule, |
| 485 | # since buffers are tensors which will be changed after .cuda() |
| 486 | for name, buffer in m.named_buffers(): |
| 487 | buffer.task_specific = task_specific |
| 488 | buffer.backbone_specific = False |
| 489 | buffer.neck_specific = False |
| 490 | buffer.decoder_specific = False |
| 491 | if task_specific: |
| 492 | printlog('add buffer {} as task_specific'.format(name)) |
| 493 | |
| 494 | def add_backbone_specific(m, backbone_specific): |
| 495 | for name, param in m.named_parameters(): |