I am using histogram in tensorboard pytorch to visualize weight in my model. Currently, this is how I use tensorboard to visualize my layers's weight.
for name, weight in model.named_parameters(
tb.add_histogram(name, weight, epoch)
tb.add_histogram(f'{name}.grad', weight.grad, epoch)
How can I add activation (ReLU, tanh) histogram in a similar manner? My model architecture is shown as below. I used UNet++ model. I'm sorry for the display. Actually, the 'def' oobject is inside the 'class', but when I copied my code, I turned out like this. The code is too long for me to fix it. Sorry I'm new to StackOverflow.
This is the image of UNet++: UNet++ architecture
class NestedUNet(nn.Module):
"""
UNET++
"""
def __init__(self, deepsupervision=args.deep_supervision):
super(NestedUNet, self).__init__()
nb_filter = [32, 64, 128, 256, 512]
self.deepsupervision = deepsupervision
self.conv0_0 = VGGBlock(6, nb_filter[0], nb_filter[0])
self.conv1_0 = Down(nb_filter[0], nb_filter[1], nb_filter[1])
self.conv2_0 = Down(nb_filter[1], nb_filter[2], nb_filter[2])
self.conv3_0 = Down(nb_filter[2], nb_filter[3], nb_filter[3])
self.conv4_0 = Down(nb_filter[3], nb_filter[4], nb_filter[4])
self.conv0_1 = Up1(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
self.conv1_1 = Up1(nb_filter[1] + nb_filter[2], nb_filter[1], nb_filter[1])
self.conv2_1 = Up1(nb_filter[2] + nb_filter[3], nb_filter[2], nb_filter[2])
self.conv3_1 = Up1(nb_filter[3] + nb_filter[4], nb_filter[3], nb_filter[3])
self.conv0_2 = Up2(nb_filter[0] * 2 + nb_filter[1], nb_filter[0], nb_filter[0])
self.conv1_2 = Up2(nb_filter[1] * 2 + nb_filter[2], nb_filter[1], nb_filter[1])
self.conv2_2 = Up2(nb_filter[2] * 2 + nb_filter[3], nb_filter[2], nb_filter[2])
self.conv0_3 = Up3(nb_filter[0] * 3 + nb_filter[1], nb_filter[0], nb_filter[0])
self.conv1_3 = Up3(nb_filter[1] * 3 + nb_filter[2], nb_filter[1], nb_filter[1])
self.conv0_4 = Up4(nb_filter[0] * 4 + nb_filter[1], nb_filter[0], nb_filter[0])
if self.deepsupervision:
self.output = OutConv(nb_filter[0], 3, deepsupervision=True)
else:
self.output = OutConv(nb_filter[0], 3, deepsupervision=False)
self.synthesis = SynthesisModule()
def forward(self, x):
x0_0 = self.conv0_0(x)
x1_0 = self.conv1_0(x0_0)
x0_1 = self.conv0_1(x1_0, x0_0)
x2_0 = self.conv2_0(x1_0)
x1_1 = self.conv1_1(x2_0, x1_0)
x0_2 = self.conv0_2(x1_1, x0_0, x0_1)
x3_0 = self.conv3_0(x2_0)
x2_1 = self.conv2_1(x3_0, x2_0)
x1_2 = self.conv1_2(x2_1, x1_0, x1_1)
x0_3 = self.conv0_3(x1_2, x0_0, x0_1, x0_2)
x4_0 = self.conv4_0(x3_0)
x3_1 = self.conv3_1(x4_0, x3_0)
x2_2 = self.conv2_2(x3_1, x2_0, x2_1)
x1_3 = self.conv1_3(x2_2, x1_0, x1_1, x1_2)
x0_4 = self.conv0_4(x1_3, x0_0, x0_1, x0_2, x0_3)
if self.deepsupervision:
output = self.output([x0_1, x0_2, x0_3, x0_4])
else:
output = self.output(x0_4)
final = self.synthesis(output, x)
return final
class VGGBlock(nn.Module):
def __init__(self, in_channels, middle_channels, out_channels):
super(VGGBlock, self).__init__()
self.vgg = nn.Sequential(
nn.Conv2d(in_channels, middle_channels, (3, 3), padding='same', bias=False),
nn.BatchNorm2d(middle_channels),
nn.ReLU(inplace=True),
nn.Dropout(0.4),
nn.Conv2d(middle_channels, out_channels, (3, 3), padding='same', bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Dropout(0.3),
)
self.initialize_weights()
def initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def forward(self, x):
return self.vgg(x)
class Up3(nn.Module):
def __init__(self, in_channels, middle_channels, out_channels, mode='bilinear'):
super().__init__()
self.up = nn.Upsample(scale_factor=2, mode=mode, align_corners=True)
self.conv = VGGBlock(in_channels, middle_channels, out_channels)
def forward(self, x1, x2, x3, x4):
x1 = self.up(x1)
x = torch.cat([x4, x3, x2, x1], dim=1)
return self.conv(x)
class Up4(nn.Module):
def __init__(self, in_channels, middle_channels, out_channels, mode='bilinear'):
super().__init__()
self.up = nn.Upsample(scale_factor=2, mode=mode, align_corners=True)
self.conv = VGGBlock(in_channels, middle_channels, out_channels)
def forward(self, x1, x2, x3, x4, x5):
x1 = self.up(x1)
x = torch.cat([x5, x4, x3, x2, x1], dim=1)
return self.conv(x)
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels, deepsupervision=False):
super(OutConv, self).__init__()
self.deepsupervision = deepsupervision
self.final = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1)),
nn.BatchNorm2d(out_channels),
nn.Tanh()
)
self.intialize_weights()
def intialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.xavier_normal_(m.weight, gain=nn.init.calculate_gain('relu')) # gain=5/3
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def forward(self, x):
if self.deepsupervision:
return [self.final(x[0]), self.final(x[1]), self.final(x[2]), self.final(x[3])]
else:
return self.final(x)
