python import torch import torch.nn as nn from sklearn.cluster import KMeans

1. 训练得到一个原始模型（此处假设已训练好）

original_weights = torch.randn(256, 512) # 示例权重

2. 展平权重并聚类

weights_np = original_weights.detach().numpy().ravel().reshape(-1, 1) k = 256 kmeans = KMeans(n_clusters=k, random_state=0).fit(weights_np) centroids = torch.tensor(kmeans.cluster_centers_.squeeze(), dtype=torch.float32) assignments = torch.tensor(kmeans.labels_, dtype=torch.long).view(original_weights.shape)

3. 构建压缩层

class ClusteredLinear(nn.Module): def init(self, centroids, assignments, bias=None): super().init() self.centroids = nn.Parameter(centroids) # 可学习质心 self.register_buffer('assignments', assignments) # 固定分配 if bias is not None: self.bias = nn.Parameter(bias) else: self.bias = None

def forward(self, x):
    # 根据 assignments 索引 centroids，重建权重矩阵
    weight = self.centroids[self.assignments]
    return nn.functional.linear(x, weight, self.bias)

4. 初始化压缩层并微调

clustered_layer = ClusteredLinear(centroids, assignments) optimizer = torch.optim.SGD(clustered_layer.parameters(), lr=0.01)

模拟几步训练

for _ in range(100): input_tensor = torch.randn(64, 256) output = clustered_layer(input_tensor) loss = output.sum() # 占位损失 optimizer.zero_grad() loss.backward() optimizer.step()