Hi all!
I have an implementation of Virtual Batch Normalization (VBN) that I’m using on each convolution of a discriminator network similar to DCGAN, more specifically, this is the discriminator network on the SEGAN paper.
Let’s assume disc is the state of the Discriminator after after a forward pass and that d_loss is the loss of disc after the same forward pass.
We then back-propagate the loss and take a peek at the norm of the gradients to find out that p.grad is equal to None
The code works normally if I remove the VBN after each conv layeror OR if in VBN’s forward method I set x.data = out and return x.
Any help is appreciated!!!
d_loss.backward()
d_grad_norm = 0
for k, p in disc.named_parameters():
d_grad_norm += p.grad.data.norm()
class Discriminator(nn.Module):
def __init__(self, ndf, kernel_size):
super(Discriminator, self).__init__()
self.encoder = nn.ModuleList([
# 16
nn.Conv1d(1, ndf, kernel_size, 2, 15, 1)
VBN(ndf),
# 32
nn.Conv1d(ndf, ndf * 2, kernel_size, 2, 15, 1),
VBN(ndf*2),
# more convolutions... and finally
# Linear for output
nn.Linear(8, 1)
])
class VBN(Module):
"""
Virtual Batch Normalization
"""
def __init__(self, n_features, epsilon=1e-5):
super(VBN, self).__init__()
assert isinstance(epsilon, float)
# batch statistics
self.epsilon = epsilon
self.mean = torch.zeros(1, n_features, 1)
self.mean_sq = torch.zeros(1, n_features, 1)
self.batch_size = None
# reference output
self.reference_output = None
def initialize(self, x):
# compute batch statistics
# self.mean = torch.mean(x, [0, 2], keepdim=True)
# self.mean_sq = torch.mean(x**2, [0, 2], keepdim=True)
self.mean = x.data.mean(2).mean(0).resize_(1, x.size(1), 1)
self.mean_sq = (x.data**2).mean(2).mean(0).resize_(1, x.size(1), 1)
self.batch_size = x.size(0)
assert x is not None
assert self.mean is not None
assert self.mean_sq is not None
# compute reference output
out = self._normalize(x, self.mean, self.mean_sq)
self.reference_output = out
def forward(self, x):
if self.reference_output is None:
self.initialize(x)
new_coeff = 1. / (self.batch_size + 1.)
old_coeff = 1. - new_coeff
# new_mean = torch.mean(x, [0, 2], keep_dims=True)
# new_mean_sq = torch.mean(x**2, [0, 2], keep_dims=True)
new_mean = x.data.mean(2).mean(0).resize_as_(self.mean)
new_mean_sq = (x.data**2).mean(2).mean(0).resize_as_(self.mean_sq)
mean = new_coeff * new_mean + old_coeff * self.mean
mean_sq = new_coeff * new_mean_sq + old_coeff * self.mean_sq
out = self._normalize(x, mean, mean_sq)
return Variable(out)
def _normalize(self, x, mean, mean_sq):
assert self.epsilon is not None
assert mean_sq is not None
assert mean is not None
assert len(x.size()) == 3
gamma = torch.normal(means=torch.ones(1, x.size(1), 1), std=0.02)
gamma = gamma.float().cuda(async=True)
beta = torch.cuda.FloatTensor(1, x.size(1), 1).fill_(0)
std = torch.sqrt(self.epsilon + mean_sq - mean**2)
out = x.data - mean
out = out / std
out = out * gamma
out = out + beta
return out
def __repr__(self):
return ('{name}({num_features}, eps={eps}, mean={mean}, mean_sq={mean_sq}'.format(
name=self.__class__.__name__, **self.__dict__))
