Image Binary Classification not converging (outputs only 0s, or 1s)

Hello everyone;

I hope that you’re safe and doing well.

I have an issue making my network converging, all inputs are either 0s or 1s, even with clipping gradient while training. Please find my code, any help would very appreciated.

class Clf(nn.Module):
    def __init__(self, pretrained=None, channels_img=3, features_d=64):
        super(Clf, self).__init__()
            # Input: N x channels_img x 256 x 256
        self.pretrained = pretrained

        if self.pretrained is None:
            self.conv1 = nn.Conv2d(channels_img, features_d, kernel_size=4, stride=2, padding=1) # 64 x 128 x 128
            self.relu = nn.LeakyReLU(0.2)
            self.block1 = self._block(features_d, features_d * 2, 4, 2, 1)  # 64 x 128 x 32 x 32
            self.block2 = self._block(features_d * 2, features_d * 4, 4, 2, 1)  # 64 x 256 x 8 x 8
            self.block3 = self._block(features_d * 4, features_d * 8, 4, 2, 1) # 64 x 512 x 2 x 2
            self.reshape = Reshape() # => (64, -1)
            self.fc1 = nn.Linear(features_d * 8 * 2 * 2, 1024) # =>
            self.fc2 = nn.Linear(1024, 1)



        if self.pretrained is not None:
            self.resnetfc1 = nn.Linear(1000, 256)
            self.resnetfc2 = nn.Linear(256, 1)



            switcher = {
                'resnet18': models.resnet18(pretrained=True),
                'resnet34': models.resnet34(pretrained=True),
                'resnet50': models.resnet50(pretrained=True),
                'mobilenetv2':models.mobilenet_v2(pretrained=True),
                'vgg16': models.vgg16_bn(pretrained=True)
            }

            self.trainedNN = switcher.get(self.pretrained, 'The pretrained you typed doesnt exist')

    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2,2)
        )

    def forward(self, x):
        if self.pretrained is None:
            x = self.conv1(x)
            x = self.relu(x)
            x = self.block1(x)
            x = self.block2(x)
            x = self.block3(x)
            x = self.reshape(x)
            x1 = self.fc1(x)
            x = self.fc2(x1)
            #print(f'mon x => {x}')
            #print(f'Sa Sigmoid => {torch.sigmoid(x)}')
            return torch.sigmoid(x), x1
        else:
            x = self.trainedNN(x)
            x1 = self.resnetfc1(x)
            x = self.resnetfc2(x1)
            return torch.sigmoid(x), x1

A helper.py which has training function, and accuracy computation:

def weights_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.xavier_uniform(m.weight.data)

def check_accuracy(loader, model, classif=None, adv=None):
    num_correct = 0
    num_samples = 0
    model.eval()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device=device)
            y = y.to(device=device)
            if adv is not None and classif is not None:
                _,x = classif(x)
                scores = model(x)
            else:
                scores,_ = model(x)
            predictions = (scores >= 0.5)
            num_correct += (predictions == y).sum()
            num_samples += x.size(0)
        print(f' Accuray num correct / num samples {num_correct/num_samples}')
    model.train()

def train_model(model, criterion, optimizer, loader, test_loader, num_epoch, myCLf=None, adv=None):
    t = trange(num_epoch, desc='Progress Bar')
    for epoch in t:
        train_one_epoch(model, loader, optimizer, criterion, myCLf, adv, desc = t)

def train_one_epoch(model, loader, optimizer, criterion, myClf=None, adv=None, desc=None):
    losses = []
    longueur_data = 0
    for batch_idx, (x, y) in enumerate(loader):
        x = x.to(device=device)

        if adv == None:
            scores,_ = model(x)
            y = y.to(device=device).to(torch.float32).unsqueeze(1)
        else:
            y = y.to(device=device).to(torch.long)
            p_y = myClf(x)[1].detach()
            scores = model(p_y)

        optimizer.zero_grad()
        loss = criterion(scores, y)
        losses.append(loss.item())
        # calcul des gradients
        loss.backward()
        # optimisation des weights
        optimizer.step()

        longueur_data += x.size(0)

    Loss = sum(losses) / longueur_data
    desc.set_postfix(Loss = Loss)

My Custom DataSet.py:

class myDataset(Dataset):
    def __init__(self, csv_file, root_dir, target, length, adv = None, transform=None):
        self.annotations = pd.read_csv(csv_file).iloc[:length,:]
        self.root_dir = root_dir
        self.transform = transform
        self.target = target
        self.length = length
        self.adv = adv

    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, index):
        img_path = os.path.join(self.root_dir, self.annotations.loc[index, 'image_id'])
        image = Image.open(img_path)
        image = np.array(image)

        if self.transform:
            image = self.transform(image=image)["image"]

        image = np.transpose(image, (2, 0, 1)).astype(np.float32)
        image = torch.Tensor(image)

        y_label = torch.tensor(int(self.annotations.loc[index, str(self.target)]))

        if self.adv is None:
            return image, y_label

        if self.adv :
            z_label = torch.tensor(int(self.annotations.loc[index, 'origin']))
            return image, y_label, z_label

And finally, all the calls:

device = 'cuda' if torch.cuda.is_available() else 'cpu'

    batch_size = 64
    learning_rate = 1e-3
    num_epochs = 10

    pretrained = None
    length = 5000

    root_dir = '../../AdvDebias/dataset/img_align_celeba/img_align_celeba'


    aug = al.Compose([
        #al.RandomResizedCrop(256, 256, p=0.2),
        al.Resize(256, 256),
        al.Transpose(p=0.2),
        al.HorizontalFlip(p=0.5),
        al.VerticalFlip(p=0.2),
        al.augmentations.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], p=1.0)
    ])

    dataset = myDataset(csv_file='../input/all_labels.csv',
                        root_dir=root_dir,
                        target='gender',
                        length=length,
                        transform=aug)

    train_set, test_set = torch.utils.data.random_split(dataset,[int(length*0.8), length - int(length*0.8)])


    train_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=True)

    model = Clf(pretrained=pretrained)
    model.apply(weights_init)
    model.to(device)

    #for (name, module) in model.named_modules():
     #   if name.startswith('trainedNN.layer4.2') or name.startswith('trainedNN.layer4.3'):
      #      for layer in module.children():
       #         for param in layer.parameters():
        #            param.requires_grad = False

    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    print(f'Beginning of Training mode ')
    train_model(model, criterion, optimizer, train_loader, test_loader, num_epochs)

    print('Checking accuracy on Training set')
    check_accuracy(train_loader, model)

    print('Checking accuracy on Test set')
    check_accuracy(test_loader, model)

    torch.save(model.state_dict(), '../models/myClf')

Please Consider here, that pretrained = None, and Adv = None.

I can execute MyClassif.py to show you the exploding values, which makes the Sigmoid either 0s or 1s, just tell me.
Thank you very much.
Habib