| 34 | |
| 35 | |
| 36 | class SyncTestCase(TorchTestCase): |
| 37 | def _syncParameters(self, bn1, bn2): |
| 38 | bn1.reset_parameters() |
| 39 | bn2.reset_parameters() |
| 40 | if bn1.affine and bn2.affine: |
| 41 | bn2.weight.data.copy_(bn1.weight.data) |
| 42 | bn2.bias.data.copy_(bn1.bias.data) |
| 43 | |
| 44 | def _checkBatchNormResult(self, bn1, bn2, input, is_train, cuda=False): |
| 45 | """Check the forward and backward for the customized batch normalization.""" |
| 46 | bn1.train(mode=is_train) |
| 47 | bn2.train(mode=is_train) |
| 48 | |
| 49 | if cuda: |
| 50 | input = input.cuda() |
| 51 | |
| 52 | self._syncParameters(_find_bn(bn1), _find_bn(bn2)) |
| 53 | |
| 54 | input1 = Variable(input, requires_grad=True) |
| 55 | output1 = bn1(input1) |
| 56 | output1.sum().backward() |
| 57 | input2 = Variable(input, requires_grad=True) |
| 58 | output2 = bn2(input2) |
| 59 | output2.sum().backward() |
| 60 | |
| 61 | self.assertTensorClose(input1.data, input2.data) |
| 62 | self.assertTensorClose(output1.data, output2.data) |
| 63 | self.assertTensorClose(input1.grad, input2.grad) |
| 64 | self.assertTensorClose(_find_bn(bn1).running_mean, _find_bn(bn2).running_mean) |
| 65 | self.assertTensorClose(_find_bn(bn1).running_var, _find_bn(bn2).running_var) |
| 66 | |
| 67 | def testSyncBatchNormNormalTrain(self): |
| 68 | bn = nn.BatchNorm1d(10) |
| 69 | sync_bn = SynchronizedBatchNorm1d(10) |
| 70 | |
| 71 | self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), True) |
| 72 | |
| 73 | def testSyncBatchNormNormalEval(self): |
| 74 | bn = nn.BatchNorm1d(10) |
| 75 | sync_bn = SynchronizedBatchNorm1d(10) |
| 76 | |
| 77 | self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), False) |
| 78 | |
| 79 | def testSyncBatchNormSyncTrain(self): |
| 80 | bn = nn.BatchNorm1d(10, eps=1e-5, affine=False) |
| 81 | sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) |
| 82 | sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) |
| 83 | |
| 84 | bn.cuda() |
| 85 | sync_bn.cuda() |
| 86 | |
| 87 | self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), True, cuda=True) |
| 88 | |
| 89 | def testSyncBatchNormSyncEval(self): |
| 90 | bn = nn.BatchNorm1d(10, eps=1e-5, affine=False) |
| 91 | sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) |
| 92 | sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) |
| 93 |
nothing calls this directly
no outgoing calls
no test coverage detected