File size: 1,098 Bytes
cdfcfcf | 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 | """Tiny CPU smoke test; run with pytest after installing requirements."""
import torch
from torch import nn
from higher_order_pruner import PruneConfig, discover_linear_groups, prune
class TinyNet(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList([nn.Linear(5, 5), nn.Linear(5, 5)])
def forward(self, x):
return self.layers[1](self.layers[0](x).tanh())
def test_exact_budget_and_general_loss():
torch.manual_seed(17)
model = TinyNet()
batches = [(torch.randn(4, 5), torch.randn(4, 5)) for _ in range(3)]
def mse(m, batch):
x, y = batch
return (m(x) - y).square().mean()
assert len(discover_linear_groups(model)) == 2
report = prune(model, batches[:2], mse, PruneConfig(sparsity=0.3, probes=1),
validation_batches=batches[2:])
assert report.prunable_weights == 50
assert report.exact_zero_weights == 15
assert sum(int((p.weight == 0).sum()) for p in model.layers) == 15
assert all(row["loss_after"] <= row["loss_before"] + 1e-5 for row in report.layers)
|