broadcast multi-task model parameters
(model, task_grp, share_backbone_group, share_neck_group, share_decoder_group, ignore)
| 85 | |
| 86 | |
| 87 | def broadcast_params_multitask(model, task_grp, share_backbone_group, share_neck_group, share_decoder_group, ignore): |
| 88 | """ broadcast multi-task model parameters """ |
| 89 | if task_grp is not None or share_backbone_group is not None \ |
| 90 | or share_neck_group is not None or share_decoder_group is None: |
| 91 | |
| 92 | for name,p in model.named_parameters(): |
| 93 | if ignore and name in ignore: |
| 94 | printlog('param {} ignored in broadcast'.format(name)) |
| 95 | continue |
| 96 | assert p.task_specific + p.backbone_specific + p.neck_specific + p.decoder_specific <= 1.5, \ |
| 97 | "param could not be task_specific, backbone_specific, neck_specific, decoder_specific at same time" |
| 98 | try: |
| 99 | # import pdb; pdb.set_trace() |
| 100 | if p.task_specific: |
| 101 | printlog(f'broadcasting task-specific param {name}\tgroup_idx={task_grp}') |
| 102 | broadcast(p, _get_global_rank(task_grp, 0), group=task_grp) |
| 103 | elif p.backbone_specific: |
| 104 | printlog(f'broadcasting backbone-specific param {name}\tgroup_idx={share_backbone_group}') |
| 105 | broadcast(p, _get_global_rank(task_grp, 0), group=share_backbone_group) |
| 106 | elif p.neck_specific: |
| 107 | printlog(f'broadcasting neck-specific param {name}\tgroup_idx={share_neck_group}') |
| 108 | broadcast(p, _get_global_rank(task_grp, 0), group=share_neck_group) |
| 109 | elif p.decoder_specific: |
| 110 | printlog(f'broadcasting decoder-specific param {name}\tgroup_idx={share_decoder_group}') |
| 111 | broadcast(p, _get_global_rank(task_grp, 0), group=share_decoder_group) |
| 112 | else: |
| 113 | printlog(f'broadcasting non-specific param {name}') |
| 114 | broadcast(p, 0) |
| 115 | except: |
| 116 | raise RuntimeError('param {} does not have task_specific or backbone_specific or neck_specific or decoder_specific'.format(name)) |
| 117 | |
| 118 | for name,b in model.named_buffers(): |
| 119 | if ignore and name in ignore: |
| 120 | printlog('buffer {} ignored in broadcast'.format(name)) |
| 121 | continue |
| 122 | assert b.task_specific + b.backbone_specific + b.neck_specific + b.decoder_specific <= 1, \ |
| 123 | "buffer could not be task_specific, backbone_specific, neck_specific, decoder_specific at same time" |
| 124 | try: |
| 125 | if b.task_specific: |
| 126 | printlog('broadcasting task-specific buffer {}'.format(name)) |
| 127 | broadcast(b, _get_global_rank(task_grp, 0), group=task_grp) |
| 128 | elif b.backbone_specific: |
| 129 | printlog('broadcasting backbone-specific buffer {}'.format(name)) |
| 130 | broadcast(b, _get_global_rank(task_grp, 0), group=share_backbone_group) |
| 131 | elif b.neck_specific: |
| 132 | printlog('broadcasting neck-specific buffer {}'.format(name)) |
| 133 | broadcast(b, _get_global_rank(task_grp, 0), group=share_neck_group) |
| 134 | elif b.decoder_specific: |
| 135 | printlog('broadcasting decoder-specific buffer {}'.format(name)) |
| 136 | broadcast(b, _get_global_rank(task_grp, 0), group=share_decoder_group) |
| 137 | else: |
| 138 | broadcast(b, 0) |
| 139 | except: |
| 140 | raise RuntimeError('buffer {} does not have task_specific'.format(name, id(b))) |
| 141 | else: |
| 142 | for name,p in model.named_parameters(): |
| 143 | if ignore and name in ignore: |
| 144 | printlog('param {} ignored in broadcast'.format(name)) |