Same loss value after certain no. of epochs at early stage of training

Pranjali_Singh: <p>I am training UNet with (4,40,40,40) size of input data. I performed some preprocessing and then fed the data into the network that uses Adam optimimzer and weighted CE Loss. However, around 9th epoch (total 90) the loss value is always 23.01 with train and validation F1-accurcy around 25%. I have tried many thing and checked every code multiple times but coldnt find the problem.</p>
<p>Code for layers of UNet:</p>
<pre><code class="lang-auto">class Conv(nn.Module):

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.single_conv = nn.Sequential(
            nn.Conv3d(in_channels, mid_channels, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm3d(mid_channels),
            nn.ReLU(inplace = True), #put inplace = True when using RELU    
            nn.Conv3d(mid_channels, out_channels, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace = True)
        )       #Same kernel , stride and padding throughout

    def forward(self, x):
        return self.single_conv(x)

class Down(nn.Module):

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool3d(2),
            nn.Dropout(p=0.2),
            Conv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):

    def __init__(self, in_channels, out_channels, trilinear=True):
        super().__init__()

        # if trilinear, use the normal convolutions to reduce the number of channels
        if trilinear:
            self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
            self.conv = Conv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose3d(in_channels, in_channels // 2,  kernel_size=2, stride=2)
            self.conv = Conv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        diffZ = x2.size()[2] - x1.size()[2]
        diffY = x2.size()[3] - x1.size()[3]
        diffX = x2.size()[4] - x1.size()[4]
        
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2,
                        diffZ // 2, diffZ - diffZ // 2])
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

</code></pre>
<p>Code for model of UNet:</p>
<pre><code class="lang-auto">from UNet_Layers import *

class UNet_multichannels(nn.Module):
    
    def __init__(self, n_channels, n_classes, feature, trilinear=True):
        super(UNet_multichannels, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.trilinear = trilinear
        self.feature = feature

        self.inp = Conv(self.n_channels,self.feature)
        self.down1 = Down(self.feature, 2*self.feature)
        factor = 2 if self.trilinear else 1
        self.down2 = Down(2*self.feature,4*self.feature // factor)
        self.up1 = Up(4*self.feature, 2*self.feature // factor, self.trilinear)
        self.up2 = Up(2*self.feature, self.feature, self.trilinear)
        self.out = OutConv(self.feature, self.n_classes) 

    def forward(self, x):
        x1 = self.inp(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x = self.up1(x3, x2)
        x = self.up2(x, x1)
        logits = self.out(x)
        return logits
</code></pre>