Hope somebody can help me, I’ve been stucked on this and there’s no way I can find the origin of my problem…

So I have a model that I have fine-tuned, it’s a resent18 that looks like this (I’m just putting last layers):

      (1): BasicBlock(
        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (7): Sequential(
      (0): BasicBlock(
        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (downsample): Sequential(
          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): BasicBlock(
        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (8): AdaptiveAvgPool2d(output_size=(1, 1))
  )
  (classifier): Linear(in_features=512, out_features=2, bias=True)
  (domain_classifier): Sequential(
    (0): Linear(in_features=512, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=1, bias=True)
  )
)

And I have a lot of unlabeled data that I wanted to apply UDA as simple as possible, this is my code:

class DomainAdversarialNetwork(nn.Module):
    def __init__(self, num_classes=2):
        super(DomainAdversarialNetwork, self).__init__()
        resnet = models.resnet18(pretrained=True)
        self.feature_extractor = nn.Sequential(*list(resnet.children())[:-1])  # Remove the original fully connected layer
        self.classifier = nn.Linear(512, num_classes)  # Assuming ResNet18 outputs 512 features
        self.domain_classifier = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 1)
        )

    def forward(self, x):
        features = self.feature_extractor(x)
        features = features.view(features.size(0), -1)  # Flatten the features
        class_output = self.classifier(features)
        domain_output = self.domain_classifier(features)
        return class_output, domain_output


# Load fine-tuned model's state_dict
state_dict = torch.load('/kaggle/input/fairface.pth') 

# Create an instance of the DomainAdversarialNetwork
dan_model = DomainAdversarialNetwork()

# Load the fine-tuned weights into the feature extractor part of the model
feature_extractor_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items() if k.startswith('conv1') or k.startswith('bn1') or k.startswith('layer')}
dan_model.feature_extractor.load_state_dict(feature_extractor_state_dict, strict=False)

# Load the classifier weights if they exist
if 'fc.weight' in state_dict and 'fc.bias' in state_dict:
    dan_model.classifier.weight.data = state_dict['fc.weight']
    dan_model.classifier.bias.data = state_dict['fc.bias']

# Move the model to GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dan_model.to(device)


# Criterion and optimizer
criterion_domain = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(dan_model.parameters(), lr=0.001)

# Training function for domain adaptation
def train_domain_adaptation(model, unlabeled_loader, criterion_domain, optimizer, device, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        for inputs, _ in unlabeled_loader:
            inputs = inputs.to(device)

            # Generate balanced fake domain labels for the unlabeled data
            batch_size = inputs.size(0)
            half_batch = batch_size // 2
            domain_labels_unlabeled = torch.cat((torch.zeros(half_batch, dtype=torch.float), torch.ones(batch_size - half_batch, dtype=torch.float))).to(device).view(-1, 1)

            optimizer.zero_grad()

            # Forward pass: Unlabeled data
            _, domain_outputs_unlabeled = model(inputs)
            loss_domain_unlabeled = criterion_domain(domain_outputs_unlabeled, domain_labels_unlabeled)

            # Total loss
            loss = loss_domain_unlabeled
            loss.backward()
            optimizer.step()

        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

        # Debugging: Print domain labels and model outputs
        print(f'Domain labels: {domain_labels_unlabeled.detach().cpu().numpy()}')
        print(f'Domain outputs: {domain_outputs_unlabeled.detach().cpu().numpy()}')


train_domain_adaptation(dan_model, unlabeled_loader, criterion_domain, optimizer, device, num_epochs=10)

This is the output:

Epoch [1/10], Loss: 0.6889
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.02605694]
 [0.02605694]
 [0.02605694]]
Epoch [2/10], Loss: 0.7023
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.15232627]
 [0.07930563]
 [0.02575949]]
Epoch [3/10], Loss: 0.6884
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.02364981]
 [0.02909829]
 [0.02364981]]
Epoch [4/10], Loss: 0.6894
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.02283685]
 [0.02283685]
 [0.02283685]]
Epoch [5/10], Loss: 0.6895
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.02219055]
 [0.02219055]
 [0.02219055]]
Epoch [6/10], Loss: 0.6896
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.02155423]
 [0.02155423]
 [0.02155423]]
Epoch [7/10], Loss: 0.6897
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.02093792]
 [0.02093792]
 [0.02093792]]
Epoch [8/10], Loss: 0.6898
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.02035251]
 [0.02035251]
 [0.02035251]]
Epoch [9/10], Loss: 0.6899
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.01980559]
 [0.01980559]
 [0.01980559]]
Epoch [10/10], Loss: 0.6900
Domain labels: [[0.]
 [1.]
 [1.]]
Domain outputs: [[0.01930189]
 [0.01930189]
 [0.01930189]]

I’m not understanding the output, I think the model is not learning and I don’t understand why… if somebody can give me advice please, would be very appreciate it! Thanks:)