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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
| """
SCULPT: Training Edge Vision Models for Post-Training Quantization Readiness
论文核心方法复现
两个核心组件:
1. Topology-aware Activation Regularizer: 抑制重尾分布
2. Percentile-based Clipping: 学习部署就绪的激活边界
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Dict, Optional
import numpy as np
class ActivationStatistics:
"""
激活分布统计
监控偏度(skewness)和峰度(kurtosis)，
量化不友好的分布特征
"""
@staticmethod
def compute_skewness(x: torch.Tensor) -> torch.Tensor:
"""偏度: 分布不对称性"""
mean = x.mean()
std = x.std()
if std == 0:
return torch.tensor(0.0)
return ((x - mean) ** 3).mean() / (std ** 3)
@staticmethod
def compute_kurtosis(x: torch.Tensor) -> torch.Tensor:
"""峰度: 尾部厚度 (正态分布=3)"""
mean = x.mean()
std = x.std()
if std == 0:
return torch.tensor(0.0)
return ((x - mean) ** 4).mean() / (std ** 4) - 3
@staticmethod
def quantization_friendliness(x: torch.Tensor) -> torch.Tensor:
"""
量化友好度评分 (越低越好)
重尾分布 → 高分 → 量化不友好
正态分布 → 低分 → 量化友好
"""
skew = ActivationStatistics.compute_skewness(x)
kurt = ActivationStatistics.compute_kurtosis(x)
return torch.abs(skew) + torch.abs(kurt) * 0.5
class TopologyAwareRegularizer(nn.Module):
"""
拓扑感知激活正则化器
根据网络层的位置(浅层/深层)调整正则化强度
浅层: 特征提取, 激活分布较均匀 → 轻正则化
深层: 任务特化, 激活分布可能极端 → 重正则化
"""
def __init__(self, n_layers: int = 10, base_weight: float = 0.01):
super().__init__()
self.n_layers = n_layers
self.base_weight = base_weight
self.layer_weights = torch.linspace(0.5, 2.0, n_layers)
def forward(self, activations: list) -> torch.Tensor:
"""
Args:
activations: 各层激活值列表, [tensor, ...]
Returns:
loss: 正则化损失
"""
total_loss = torch.tensor(0.0, device=activations[0].device)
for i, act in enumerate(activations):
layer_weight = self.layer_weights[i].to(act.device)
unfriendliness = ActivationStatistics.quantization_friendliness(act)
total_loss = total_loss + layer_weight * unfriendliness
return self.base_weight * total_loss / len(activations)
class PercentileClipping(nn.Module):
"""
百分位裁剪机制
学习部署就绪的激活裁剪边界
导出后直接用于 PTQ 工作流
"""
def __init__(self, init_percentile: float = 99.5):
super().__init__()
self.clip_percentile_logit = nn.Parameter(
torch.tensor(self._percentile_to_logit(init_percentile))
)
@staticmethod
def _percentile_to_logit(p: float) -> float:
"""将百分位转换为 logit 空间"""
import math
normalized = (p - 90) / (99.99 - 90)
normalized = max(1e-6, min(1-1e-6, normalized))
return math.log(normalized / (1 - normalized))
def get_percentile(self) -> float:
"""获取当前裁剪百分位"""
normalized = torch.sigmoid(self.clip_percentile_logit)
return 90 + normalized * (99.99 - 90)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Args:
x: 激活值 tensor
Returns:
clipped: 裁剪后的激活值
clip_bounds: (min, max) 裁剪边界
"""
p = self.get_percentile()
lower_p = (100 - p) / 2
upper_p = 100 - lower_p
clip_min = torch.quantile(x.flatten().float(), lower_p / 100)
clip_max = torch.quantile(x.flatten().float(), upper_p / 100)
clipped = torch.clamp(x, clip_min, clip_max)
return clipped, torch.stack([clip_min, clip_max])
class SCULPTTrainer:
"""
SCULPT 训练器
在标准 FP32 训练中集成:
1. 拓扑感知正则化
2. 百分位裁剪
训练完成后导出裁剪边界供 PTQ 使用
"""
def __init__(self,
model: nn.Module,
base_lr: float = 1e-3,
reg_weight: float = 0.01,
init_percentile: float = 99.5):
self.model = model
self.reg_weight = reg_weight
self.activations = []
self._register_hooks()
n_conv_layers = sum(1 for m in model.modules() if isinstance(m, (nn.Conv2d, nn.Linear)))
self.regularizer = TopologyAwareRegularizer(n_conv_layers, base_weight=reg_weight)
self.clipper = PercentileClipping(init_percentile)
def _register_hooks(self):
"""注册前向钩子收集激活"""
self.hooks = []
for module in self.model.modules():
if isinstance(module, (nn.Conv2d, nn.Linear)):
def hook(m, inp, out, self_ref=self):
self_ref.activations.append(out)
self.hooks.append(module.register_forward_hook(hook))
def compute_loss(self, task_loss: torch.Tensor) -> Dict[str, torch.Tensor]:
"""计算总损失 = 任务损失 + 正则化损失"""
self.activations = []
_ = self.model(torch.randn(1, 3, 224, 224,
device=next(self.model.parameters()).device))
reg_loss = self.regularizer(self.activations)
total_loss = task_loss + self.reg_weight * reg_loss
return {
'total': total_loss,
'task': task_loss,
'regularization': reg_loss
}
def get_clip_bounds(self, sample_input: torch.Tensor) -> Dict[str, Tuple[float, float]]:
"""获取各层裁剪边界（导出给 PTQ）"""
self.activations = []
_ = self.model(sample_input)
bounds = {}
for i, act in enumerate(self.activations):
_, clip_bounds = self.clipper(act)
bounds[f'layer_{i}'] = (
clip_bounds[0].item(),
clip_bounds[1].item()
)
return bounds
class SimpleDMSModel(nn.Module):
"""简化 DMS 模型用于测试"""
def __init__(self, n_classes: int = 5):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, n_classes)
)
def forward(self, x):
return self.classifier(self.features(x))
def simulate_int8_quantize(x: torch.Tensor, clip_min: float, clip_max: float) -> torch.Tensor:
"""
模拟 INT8 量化
Args:
x: FP32 tensor
clip_min, clip_max: SCULPT 学习的裁剪边界
Returns:
quantized: INT8 量化后反量化的 tensor
"""
x_clipped = torch.clamp(x, clip_min, clip_max)
scale = max(abs(clip_min), abs(clip_max)) / 127
if scale == 0:
return x
x_int = torch.round(x_clipped / scale).clamp(-128, 127)
x_dequant = x_int * scale
return x_dequant
if __name__ == "__main__":
model = SimpleDMSModel(n_classes=5)
trainer = SCULPTTrainer(model, reg_weight=0.01)
print("=== SCULPT 训练测试 ===")
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(3):
x = torch.randn(8, 3, 224, 224)
target = torch.randint(0, 5, (8,))
output = model(x)
task_loss = F.cross_entropy(output, target)
losses = trainer.compute_loss(task_loss)
optimizer.zero_grad()
losses['total'].backward()
optimizer.step()
print(f"Epoch {epoch}: task={losses['task']:.4f}, reg={losses['regularization']:.4f}")
model.eval()
sample = torch.randn(1, 3, 224, 224)
clip_bounds = trainer.get_clip_bounds(sample)
print(f"\n=== SCULPT 导出裁剪边界 ===")
for layer, (lo, hi) in list(clip_bounds.items())[:5]:
print(f"{layer}: [{lo:.4f}, {hi:.4f}]")
print(f"\n=== INT8 量化精度对比 ===")
with torch.no_grad():
output_fp32 = model(sample)
print(f"FP32 输出: {output_fp32[0, :3].tolist()}")
first_conv = list(model.features.children())[0]
weight = first_conv.weight.data
bounds = list(clip_bounds.values())[0]
weight_q = simulate_int8_quantize(weight, bounds[0], bounds[1])
quant_error = F.mse_loss(weight, weight_q)
print(f"INT8 权重量化 MSE: {quant_error:.6f}")
print(f"当前裁剪百分位: {trainer.clipper.get_percentile():.2f}%")
|