| 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(): |
| 598 | _task_sp_flag = any(name.startswith(sp_name) or name.endswith(sp_name) for sp_name in task_sp_list) |
| 599 | |
| 600 | param.task_specific = _task_sp_flag |
| 601 | param.backbone_specific = False |
| 602 | param.neck_specific = False if _task_sp_flag else neck_specific |
| 603 | param.decoder_specific = False |
| 604 | if _task_sp_flag: |
| 605 | printlog('add param {} as task_specific'.format(name)) |
| 606 | elif neck_specific: |
| 607 | printlog('add param {} as neck_specific'.format(name)) |
| 608 | |
| 609 | if not hasattr(torch.nn.Module, 'named_buffers'): |
| 610 | printlog('registering named_buffers for nn.Module at add_neck_specific') |
| 611 | torch.nn.Module.named_buffers = named_buffers |
| 612 | |
| 613 | #m.cuda() # neccesary for broadcast in DistModule, since buffers are tensors which will be changed after .cuda() |
| 614 | for name, buffer in m.named_buffers(): |
| 615 | _task_sp_flag = any(name.startswith(sp_name) or name.endswith(sp_name) for sp_name in task_sp_list) |
| 616 | |
| 617 | buffer.task_specific = _task_sp_flag |
| 618 | buffer.backbone_specific = False |
| 619 | buffer.neck_specific = False if _task_sp_flag else neck_specific |
| 620 | buffer.decoder_specific = False |
| 621 | if _task_sp_flag: |
| 622 | printlog('add buffer {} as task_specific'.format(name)) |
| 623 | elif neck_specific: |
| 624 | printlog('add buffer {} as neck_specific'.format(name)) |
| 625 | |
| 626 | def add_aio_decoder_specific(m, decoder_specific, task_sp_list=(), neck_sp_list=()): |
| 627 | for name, param in m.named_parameters(): |