| 39 | |
| 40 | |
| 41 | class CategoricalConditionalBatchNorm2d(ConditionalBatchNorm2d): |
| 42 | |
| 43 | def __init__(self, num_classes, num_features, eps=1e-5, momentum=0.1, |
| 44 | affine=False, track_running_stats=True): |
| 45 | super(CategoricalConditionalBatchNorm2d, self).__init__( |
| 46 | num_features, eps, momentum, affine, track_running_stats |
| 47 | ) |
| 48 | self.weights = nn.Embedding(num_classes, num_features) |
| 49 | self.biases = nn.Embedding(num_classes, num_features) |
| 50 | |
| 51 | self._initialize() |
| 52 | |
| 53 | def _initialize(self): |
| 54 | init.ones_(self.weights.weight.data) |
| 55 | init.zeros_(self.biases.weight.data) |
| 56 | |
| 57 | def forward(self, input, c, **kwargs): |
| 58 | weight = self.weights(c) |
| 59 | bias = self.biases(c) |
| 60 | |
| 61 | return super(CategoricalConditionalBatchNorm2d, self).forward(input, weight, bias) |
| 62 | |
| 63 | |
| 64 | if __name__ == '__main__': |
no outgoing calls
no test coverage detected