[Python] My BYOL model is pretraining with AffectNet dataset subset, I ran into some loss = 0...

I'm working on my graduation thesis and I ran into some problems when training my BYOL model on Kaggle.

For those who wants to view the whole thing on Kaggle, here is the link: https://www.kaggle.com/code/nguyenchiminh/byol-raf-db-1

Anyways, my BYOL structure is as below

`class MLP(nn.Module):
def __init__(self, input_dim=512, hidden_dim=512, output_dim=256):
super(MLP, self).__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True),
nn.Dropout(p=0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True),
nn.Linear(hidden_dim, output_dim)
)

def forward(self, x):
return self.net(x)
class BYOL(nn.Module):
def __init__(self, encoder, projector_hidden_size=512, projector_output_size=256):
super(BYOL, self).__init__()

# Encoder
self.encoder = encoder
encoder_feat_dim = encoder.fc.in_features
self.encoder.fc = nn.Identity()
self.projector = MLP(encoder_feat_dim, projector_hidden_size, projector_output_size)
self.predictor = MLP(projector_output_size, projector_hidden_size, projector_output_size)

# Create target network
self.target_encoder = None
self.target_projector = None
self.update_target_network(tau=1.0) # init target network

def update_target_network(self, tau=0.996):
if self.target_encoder is None:
# Initialize target encoder and projector
self.target_encoder = copy.deepcopy(self.encoder)
self.target_projector = copy.deepcopy(self.projector)

# Freeze target encoder parameters
for param in self.target_encoder.parameters():
param.requires_grad = False
for param in self.target_projector.parameters():
param.requires_grad = False

# Set BatchNorm layers in target encoder/projector to eval mode
for module in self.target_encoder.modules():
if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
module.eval()
for module in self.target_projector.modules():
if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
module.eval()

else:
# EMA update
for online_params, target_params in zip(self.encoder.parameters(), self.target_encoder.parameters()):
target_params.data = tau * target_params.data + (1 - tau) * online_params.data

for online_params, target_params in zip(self.projector.parameters(), self.target_projector.parameters()):
target_params.data = tau * target_params.data + (1 - tau) * online_params.data


def forward(self, x1, x2):
# Online network forward
f_o1 = self.encoder(x1)
z_o1 = self.projector(f_o1)
p_o1 = self.predictor(z_o1)

f_o2 = self.encoder(x2)
z_o2 = self.projector(f_o2)
p_o2 = self.predictor(z_o2)

with torch.no_grad():
f_t1 = self.target_encoder(x1)
z_t1 = self.target_projector(f_t1)

f_t2 = self.target_encoder(x2)
z_t2 = self.target_projector(f_t2)

return p_o1, p_o2, z_t1.detach(), z_t2.detach(), f_o1, f_o2

def get_features(self, x):
return self.encoder(x)
class TwoCropsTransform:
def __init__(self, base_transform):
self.base_transform = base_transform

def __call__(self, x):
q = self.base_transform(x)
k = self.base_transform(x)
return [q, k]

def get_byol_transforms(size=224):
byol_transform = transforms.Compose([
transforms.RandomResizedCrop(size, scale=(0.8, 1.0)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomApply([
transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)
], p=0.8),
transforms.RandomGrayscale(p=0.2),
transforms.RandomApply([transforms.GaussianBlur(kernel_size=9, sigma=(0.1, 2.0))], p=0.5),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
return TwoCropsTransform(byol_transform)

class AffectNetDataset(Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.image_paths = []
self.labels = []
self.class_to_idx = {}

# Define AffectNet emotion classes
# 0: Neutral, 1: Happy, 2: Sad, 3: Surprise, 4: Fear, 5: Disgust, 6: Anger, 7: Contempt
emotion_classes = ['Neutral', 'Happy', 'Sad', 'Surprise', 'Fear', 'Disgust', 'Anger', 'Contempt']

# Find all image paths and assign labels
classes = sorted(os.listdir(root_dir))
for i, class_name in enumerate(classes):
if class_name in emotion_classes or class_name.isdigit():
self.class_to_idx[class_name] = i
class_dir = os.path.join(root_dir, class_name)
if os.path.isdir(class_dir):
for img_name in os.listdir(class_dir):
if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
self.image_paths.append(os.path.join(class_dir, img_name))
self.labels.append(i)

print(f"Found {len(self.class_to_idx)} classes: {list(self.class_to_idx.keys())}")
print(f"Total images: {len(self.image_paths)}")

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

def __getitem__(self, idx):
img_path = self.image_paths[idx]
try:
image = Image.open(img_path).convert('RGB')
except Exception as e:
print(f"Error loading image {img_path}: {e}")
# Return a black image as fallback
image = Image.new('RGB', (224, 224), (0, 0, 0))

label = self.labels[idx]

if self.transform:
image = self.transform(image)

return image, label

def byol_loss(p, z):
p = F.normalize(p, dim=-1, p=2)
z = F.normalize(z, dim=-1, p=2)
return 2 - 2 * (p * z).sum(dim=-1).mean()
def train_byol(model, train_loader, optimizer, scheduler, epochs, device, tau=0.99):
model.train()

model.target_encoder.eval()
model.target_projector.eval()

train_losses = []
train_accuracies = []
learning_rates = []

#temporary variable for epoch
epoch = 0
for images, _ in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}"):
print(f"Type of images: {type(images)}")
print(f"Length of images: {len(images)}")
print(f"Shape of images[0]: {images[0].shape}")
print(f"Shape of images[1]: {images[1].shape}")

# Check if views are different
diff = torch.abs(images[0] - images[1]).mean()
print(f"Mean absolute difference between views: {diff:.4f}")

break

for epoch in range(epochs):
running_loss = 0.0
total_cosine_sim = 0.0 # Add this line
num_batches = 0 # Add this line

for images, _ in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}"):
# Get the two views of each image
im_1 = images[0].to(device)
im_2 = images[1].to(device)

# Zero the gradients
optimizer.zero_grad()

# Forward pass
p_1, p_2, z_2, z_1, _, _ = model(im_1, im_2)

# Compute loss
loss = (byol_loss(p_1, z_2) + byol_loss(p_2, z_1)) / 2

# Calculate cosine similarity for accuracy metric (add these lines)
p_1_norm = F.normalize(p_1, dim=-1, p=2)
z_2_norm = F.normalize(z_2, dim=-1, p=2)
p_2_norm = F.normalize(p_2, dim=-1, p=2)
z_1_norm = F.normalize(z_1, dim=-1, p=2)

cosine_sim = ((p_1_norm * z_2_norm).sum(dim=-1).mean() +
(p_2_norm * z_1_norm).sum(dim=-1).mean()) / 2
total_cosine_sim += cosine_sim.item()
num_batches += 1

# Backward pass
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
optimizer.step()

# Only update the target network after a warm-up (from epoch 1 onward)
if epoch > 1:
model.update_target_network(tau=0.99)
else:
model.update_target_network(tau=1.0)

# Update statistics
running_loss += loss.item()

# Step the scheduler after each epoch
scheduler.step()
current_lr = scheduler.get_last_lr()[0]
learning_rates.append(current_lr)

epoch_loss = running_loss / len(train_loader)
epoch_accuracy = (total_cosine_sim / num_batches) * 100 # Convert to percentage, add this line
train_losses.append(epoch_loss)
train_accuracies.append(epoch_accuracy) # Add this line
print(f"Epoch {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%, LR: {current_lr:.6f}") # Modified this line

return train_losses, train_accuracies, learning_rates # Modified return statement`


Accuracy of 1 of the versions, other versions are typically like this too

after it reached 0, it stays at 0

I have tried to modify the hyperparameters like LR, WD, tau, ... and nothing worked. Originally in the paper it is designed with the SGD optimizer, my setting uses the Adam optimizer. Initially my MLP was simpler like this

`self.layer1 = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True)
)
self.layer2 = nn.Linear(hidden_dim, output_dim)`


but the result is typically like the image that I've shown.

Continue reading...