| 558 | printlog('add buffer {} as decoder_specific'.format(name)) |
| 559 | |
| 560 | def add_aio_backbone_specific(m, backbone_specific, task_sp_list=(), neck_sp_list=()): |
| 561 | for name, param in m.named_parameters(): |
| 562 | _task_sp_flag = any(name.startswith(sp_name) or name.endswith(sp_name) for sp_name in task_sp_list) |
| 563 | _neck_sp_flag = any(name.startswith(sp_name) or name.endswith(sp_name) for sp_name in neck_sp_list) |
| 564 | |
| 565 | param.task_specific = _task_sp_flag |
| 566 | param.backbone_specific = False if _task_sp_flag or _neck_sp_flag else backbone_specific |
| 567 | param.neck_specific = _neck_sp_flag |
| 568 | param.decoder_specific = False |
| 569 | if _task_sp_flag: |
| 570 | printlog('add param {} as task_specific'.format(name)) |
| 571 | elif _neck_sp_flag: |
| 572 | printlog('add param {} as neck_specific'.format(name)) |
| 573 | elif backbone_specific: |
| 574 | printlog('add param {} as backbone_specific'.format(name)) |
| 575 | |
| 576 | if not hasattr(torch.nn.Module, 'named_buffers'): |
| 577 | printlog('registering named_buffers for nn.Module at add_backbone_specific') |
| 578 | torch.nn.Module.named_buffers = named_buffers |
| 579 | |
| 580 | #m.cuda() # neccesary for broadcast in DistModule, since buffers are tensors which will be changed after .cuda() |
| 581 | for name, buffer in m.named_buffers(): |
| 582 | _task_sp_flag = any(name.startswith(sp_name) or name.endswith(sp_name) for sp_name in task_sp_list) |
| 583 | _neck_sp_flag = any(name.startswith(sp_name) or name.endswith(sp_name) for sp_name in neck_sp_list) |
| 584 | |
| 585 | buffer.task_specific = _task_sp_flag |
| 586 | buffer.backbone_specific = False if _task_sp_flag or _neck_sp_flag else backbone_specific |
| 587 | buffer.neck_specific = _neck_sp_flag |
| 588 | buffer.decoder_specific = False |
| 589 | if _task_sp_flag: |
| 590 | printlog('add buffer {} as task_specific'.format(name)) |
| 591 | elif _neck_sp_flag: |
| 592 | printlog('add buffer {} as neck_specific'.format(name)) |
| 593 | elif backbone_specific: |
| 594 | printlog('add buffer {} as backbone_specific'.format(name)) |
| 595 | |
| 596 | def add_aio_neck_specific(m, neck_specific, task_sp_list=()): |
| 597 | for name, param in m.named_parameters(): |