For every layer of ResNet, I would like to assign the average value of convs to each conv, every time after they’re updated by backward.
class ResNet(nn.Module):
def __init__(self, block, num_blocks, num_classes=10):...
def _make_layer(self, block, planes, num_blocks, stride):...
def forward(self, x):...
def setAvg(self):
for layer in [self.layer1,self.layer2,self.layer3,self.layer4]:
weight_conv_average = torch.mean(
torch.stack([torch.Tensor(p) for n, p in layer.named_parameters() if 'conv' in n and (p.size()[0]==p.size()[1])]), 0)
for n, p in layer.named_parameters():
if 'conv' in n and (p.size()[0]==p.size()[1]):
#print(n,' ', p.size())
p.data.copy_(weight_conv_average)
And in the main part of training, I use
rawnet = ResNet18()
net = rawnet.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
...
def train(epoch):
...
loss.backward()
with torch.no_grad():
rawnet.setAvg()
optimizer.step()
However this will cause Segmentation fault (core dumped) with Linux cuda env.
I’ve tried to invoke setAvg in forward but also failed.
And the reason why I distinguish net and rawnet is that net is DataParallel object, which has no attribute 'setAvg'.
Big Thanks!