(self, x, gamma)
| 11 | self.block_size = block_size |
| 12 | |
| 13 | def forward(self, x, gamma): |
| 14 | # shape: (bsize, channels, height, width) |
| 15 | |
| 16 | if self.training: |
| 17 | batch_size, channels, height, width = x.shape |
| 18 | bernoulli = Bernoulli(gamma) |
| 19 | mask = bernoulli.sample((batch_size, channels, height - (self.block_size - 1), width - (self.block_size - 1))) |
| 20 | if torch.cuda.is_available(): |
| 21 | mask = mask.cuda() |
| 22 | block_mask = self._compute_block_mask(mask) |
| 23 | countM = block_mask.size()[0] * block_mask.size()[1] * block_mask.size()[2] * block_mask.size()[3] |
| 24 | count_ones = block_mask.sum() |
| 25 | |
| 26 | return block_mask * x * (countM / count_ones) |
| 27 | else: |
| 28 | return x |
| 29 | |
| 30 | def _compute_block_mask(self, mask): |
| 31 | left_padding = int((self.block_size-1) / 2) |
nothing calls this directly
no test coverage detected