(layer: torch.nn.Linear,
index: torch.LongTensor,
dim: int = 0)
| 427 | |
| 428 | |
| 429 | def prune_linear_layer(layer: torch.nn.Linear, |
| 430 | index: torch.LongTensor, |
| 431 | dim: int = 0) -> torch.nn.Linear: |
| 432 | # Copied from transformers, the latest version of transformers deletes this function |
| 433 | index = index.to(layer.weight.device) |
| 434 | W = layer.weight.index_select(dim, index).detach().clone() |
| 435 | if layer.bias is not None: |
| 436 | if dim == 1: |
| 437 | b = layer.bias.detach().clone() |
| 438 | else: |
| 439 | b = layer.bias[index].detach().clone() |
| 440 | new_size = list(layer.weight.size()) |
| 441 | new_size[dim] = len(index) |
| 442 | new_layer = torch.nn.Linear( |
| 443 | new_size[1], new_size[0], bias=layer.bias |
| 444 | is not None).to(layer.weight.device) |
| 445 | new_layer.weight.requires_grad = False |
| 446 | new_layer.weight.copy_(W.contiguous()) |
| 447 | new_layer.weight.requires_grad = True |
| 448 | if layer.bias is not None: |
| 449 | new_layer.bias.requires_grad = False |
| 450 | new_layer.bias.copy_(b.contiguous()) |
| 451 | new_layer.bias.requires_grad = True |
| 452 | return new_layer |
no test coverage detected
searching dependent graphs…