| 52 | |
| 53 | # Record the mean and std like a BN layer but do no normalization |
| 54 | class BNStatistics(nn.Module): |
| 55 | def __init__(self, num_features): |
| 56 | super(BNStatistics, self).__init__() |
| 57 | shape = (1, num_features, 1, 1) |
| 58 | self.register_buffer('running_mean', torch.zeros(shape)) |
| 59 | self.register_buffer('running_var', torch.zeros(shape)) |
| 60 | self.is_first_batch = True |
| 61 | |
| 62 | def forward(self, x): |
| 63 | if self.running_mean.device != x.device: |
| 64 | self.running_mean = self.running_mean.to(x.device) |
| 65 | self.running_var = self.running_var.to(x.device) |
| 66 | self.running_mean, self.running_var = update_running_mean_var(x, self.running_mean, self.running_var, momentum=0.9, is_first_batch=self.is_first_batch) |
| 67 | self.is_first_batch = False |
| 68 | return x |
| 69 | |
| 70 | # This is designed to insert BNStat layer between Conv2d(without bias) and its bias |
| 71 | class BiasAdd(nn.Module): |
no outgoing calls
no test coverage detected