Prune a Conv1D layer (a model parameters) to keep only entries in index. A Conv1D work as a Linear layer (see e.g. BERT) but the weights are transposed. Return the pruned layer as a new layer with requires_grad=True. Used to remove heads.
(layer, index, dim=1)
| 1169 | |
| 1170 | |
| 1171 | def prune_conv1d_layer(layer, index, dim=1): |
| 1172 | """ Prune a Conv1D layer (a model parameters) to keep only entries in index. |
| 1173 | A Conv1D work as a Linear layer (see e.g. BERT) but the weights are transposed. |
| 1174 | Return the pruned layer as a new layer with requires_grad=True. |
| 1175 | Used to remove heads. |
| 1176 | """ |
| 1177 | index = index.to(layer.weight.device) |
| 1178 | W = layer.weight.index_select(dim, index).clone().detach() |
| 1179 | if dim == 0: |
| 1180 | b = layer.bias.clone().detach() |
| 1181 | else: |
| 1182 | b = layer.bias[index].clone().detach() |
| 1183 | new_size = list(layer.weight.size()) |
| 1184 | new_size[dim] = len(index) |
| 1185 | new_layer = Conv1D(new_size[1], new_size[0]).to(layer.weight.device) |
| 1186 | new_layer.weight.requires_grad = False |
| 1187 | new_layer.weight.copy_(W.contiguous()) |
| 1188 | new_layer.weight.requires_grad = True |
| 1189 | new_layer.bias.requires_grad = False |
| 1190 | new_layer.bias.copy_(b.contiguous()) |
| 1191 | new_layer.bias.requires_grad = True |
| 1192 | return new_layer |
| 1193 | |
| 1194 | |
| 1195 | def prune_layer(layer, index, dim=None): |
no test coverage detected