| 226 | return sim.mean(axis=0) |
| 227 | |
| 228 | def search(mod, inputs, outputs, where): |
| 229 | |
| 230 | mod._forward_hooks.clear() |
| 231 | |
| 232 | normal_in = [_[:batch_size] for _ in inputs] |
| 233 | fakequant_in = [_[batch_size:] for _ in inputs] |
| 234 | |
| 235 | disable_fake_quant(mod) |
| 236 | normal_out = mod(*normal_in) |
| 237 | enable_fake_quant(mod) |
| 238 | |
| 239 | ob = getattr(mod, where) |
| 240 | if ob is None: |
| 241 | return |
| 242 | |
| 243 | orig_scale = ob.orig_scale |
| 244 | cosine = optimal = 0 |
| 245 | for scale in np.linspace(start * orig_scale, stop * orig_scale, num): |
| 246 | ob.scale = scale |
| 247 | fakequant_out = mod(*fakequant_in) |
| 248 | dis = get_cosine(normal_out, fakequant_out) |
| 249 | if dis > cosine: |
| 250 | cosine = dis |
| 251 | optimal = scale |
| 252 | if optimal == 0: |
| 253 | logger.warning("EasyQuant finds no better scale") |
| 254 | else: |
| 255 | ob.scale = optimal |
| 256 | |
| 257 | fakequant_out = outputs[batch_size:] |
| 258 | return concat([normal_out, fakequant_out]) |
| 259 | |
| 260 | data = concat([data, data]) |
| 261 | |