Merge pre BN to reduce inference runtime.
(module, pre_bn_1, pre_bn_2=None)
| 26 | # function from Next_ViT/classification/utils.py : merge_pre_bn |
| 27 | # copied here to get rid of Next_ViT repo dependancy |
| 28 | def merge_pre_bn(module, pre_bn_1, pre_bn_2=None): |
| 29 | """ Merge pre BN to reduce inference runtime. |
| 30 | """ |
| 31 | weight = module.weight.data |
| 32 | if module.bias is None: |
| 33 | zeros = torch.zeros(module.out_channels, device=weight.device).type(weight.type()) |
| 34 | module.bias = nn.Parameter(zeros) |
| 35 | bias = module.bias.data |
| 36 | if pre_bn_2 is None: |
| 37 | assert pre_bn_1.track_running_stats is True, "Unsupport bn_module.track_running_stats is False" |
| 38 | assert pre_bn_1.affine is True, "Unsupport bn_module.affine is False" |
| 39 | |
| 40 | scale_invstd = pre_bn_1.running_var.add(pre_bn_1.eps).pow(-0.5) |
| 41 | extra_weight = scale_invstd * pre_bn_1.weight |
| 42 | extra_bias = pre_bn_1.bias - pre_bn_1.weight * pre_bn_1.running_mean * scale_invstd |
| 43 | else: |
| 44 | assert pre_bn_1.track_running_stats is True, "Unsupport bn_module.track_running_stats is False" |
| 45 | assert pre_bn_1.affine is True, "Unsupport bn_module.affine is False" |
| 46 | |
| 47 | assert pre_bn_2.track_running_stats is True, "Unsupport bn_module.track_running_stats is False" |
| 48 | assert pre_bn_2.affine is True, "Unsupport bn_module.affine is False" |
| 49 | |
| 50 | scale_invstd_1 = pre_bn_1.running_var.add(pre_bn_1.eps).pow(-0.5) |
| 51 | scale_invstd_2 = pre_bn_2.running_var.add(pre_bn_2.eps).pow(-0.5) |
| 52 | |
| 53 | extra_weight = scale_invstd_1 * pre_bn_1.weight * scale_invstd_2 * pre_bn_2.weight |
| 54 | extra_bias = scale_invstd_2 * pre_bn_2.weight *(pre_bn_1.bias - pre_bn_1.weight * pre_bn_1.running_mean * scale_invstd_1 - pre_bn_2.running_mean) + pre_bn_2.bias |
| 55 | |
| 56 | if isinstance(module, nn.Linear): |
| 57 | extra_bias = weight @ extra_bias |
| 58 | weight.mul_(extra_weight.view(1, weight.size(1)).expand_as(weight)) |
| 59 | elif isinstance(module, nn.Conv2d): |
| 60 | assert weight.shape[2] == 1 and weight.shape[3] == 1 |
| 61 | weight = weight.reshape(weight.shape[0], weight.shape[1]) |
| 62 | extra_bias = weight @ extra_bias |
| 63 | weight.mul_(extra_weight.view(1, weight.size(1)).expand_as(weight)) |
| 64 | weight = weight.reshape(weight.shape[0], weight.shape[1], 1, 1) |
| 65 | bias.add_(extra_bias) |
| 66 | |
| 67 | module.weight.data = weight |
| 68 | module.bias.data = bias |
| 69 | |
| 70 | |
| 71 |