Compressed Linear Layer.
| 30 | |
| 31 | |
| 32 | class CLinear(nn.Module): |
| 33 | """Compressed Linear Layer.""" |
| 34 | |
| 35 | def __init__(self, weight=None, bias=None, device=None): |
| 36 | super().__init__() |
| 37 | if weight is None: |
| 38 | self.weight = None |
| 39 | elif isinstance(weight, Tensor): |
| 40 | self.weight = compress(weight.data.to(device), default_compression_config) |
| 41 | else: |
| 42 | self.weight = weight |
| 43 | self.bias = bias |
| 44 | |
| 45 | def forward(self, input: Tensor) -> Tensor: |
| 46 | weight = decompress(self.weight, default_compression_config) |
| 47 | if self.bias is None: |
| 48 | return F.linear(input.to(weight.dtype), weight) |
| 49 | return F.linear(input.to(weight.dtype), weight, self.bias.to(weight.dtype)) |
| 50 | |
| 51 | |
| 52 | def compress_module(module, target_device): |
no outgoing calls
no test coverage detected