norm.py 927 B

12345678910111213141516171819202122232425
  1. import torch.nn as nn
  2. import numpy as np
  3. import torch
  4. class SpecificNorm(nn.Module):
  5. def __init__(self, epsilon=1e-8):
  6. """
  7. @notice: avoid in-place ops.
  8. https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
  9. """
  10. super(SpecificNorm, self).__init__()
  11. self.mean = np.array([0.485, 0.456, 0.406])
  12. self.mean = torch.from_numpy(self.mean).float().cuda()
  13. self.mean = self.mean.view([1, 3, 1, 1])
  14. self.std = np.array([0.229, 0.224, 0.225])
  15. self.std = torch.from_numpy(self.std).float().cuda()
  16. self.std = self.std.view([1, 3, 1, 1])
  17. def forward(self, x):
  18. mean = self.mean.expand([1, 3, x.shape[2], x.shape[3]])
  19. std = self.std.expand([1, 3, x.shape[2], x.shape[3]])
  20. x = (x - mean) / std
  21. return x