Spaces:
Running
Running
File size: 1,632 Bytes
2eec02e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | from __future__ import annotations
import torch
from torch import nn
from torch.nn import functional as F
def binary_weight(weight: torch.Tensor) -> torch.Tensor:
scale = weight.abs().mean(dim=1, keepdim=True).clamp_min(1e-6)
quantized = scale * torch.where(weight >= 0, 1.0, -1.0)
return weight + (quantized - weight).detach()
def ternary_weight(weight: torch.Tensor) -> torch.Tensor:
scale = weight.abs().mean(dim=1, keepdim=True).clamp_min(1e-6)
normalized = weight / scale
quantized = scale * normalized.round().clamp(-1, 1)
return weight + (quantized - weight).detach()
class QuantizedLinear(nn.Linear):
def __init__(self, *args, bits: str = "fp32", **kwargs) -> None:
super().__init__(*args, **kwargs)
self.bits = bits
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
if self.bits == "binary":
weight = binary_weight(self.weight)
elif self.bits == "ternary":
weight = ternary_weight(self.weight)
else:
weight = self.weight
return F.linear(inputs, weight, self.bias)
class BitMLP(nn.Module):
def __init__(self, bits: str = "fp32") -> None:
super().__init__()
self.bits = bits
self.hidden = QuantizedLinear(64, 64, bits=bits)
self.output = QuantizedLinear(64, 10, bits=bits)
def forward(self, images: torch.Tensor) -> torch.Tensor:
flattened = images.flatten(1)
return self.output(F.gelu(self.hidden(flattened)))
def parameter_count(module: nn.Module) -> int:
return sum(parameter.numel() for parameter in module.parameters())
|