(m)
| 107 | |
| 108 | |
| 109 | def fuse_model(m): |
| 110 | prev_previous_type = nn.Identity() |
| 111 | prev_previous_name = '' |
| 112 | previous_type = nn.Identity() |
| 113 | previous_name = '' |
| 114 | for name, module in m.named_modules(): |
| 115 | if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type(module) == nn.ReLU: |
| 116 | # print("FUSED ", prev_previous_name, previous_name, name) |
| 117 | torch.quantization.fuse_modules(m, [prev_previous_name, previous_name, name], inplace=True) |
| 118 | elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d: |
| 119 | # print("FUSED ", prev_previous_name, previous_name) |
| 120 | torch.quantization.fuse_modules(m, [prev_previous_name, previous_name], inplace=True) |
| 121 | # elif previous_type == nn.Conv2d and type(module) == nn.ReLU: |
| 122 | # print("FUSED ", previous_name, name) |
| 123 | # torch.quantization.fuse_modules(m, [previous_name, name], inplace=True) |
| 124 | |
| 125 | prev_previous_type = previous_type |
| 126 | prev_previous_name = previous_name |
| 127 | previous_type = type(module) |
| 128 | previous_name = name |
nothing calls this directly
no outgoing calls
no test coverage detected