Same output for every image input

[sulakshgupta988]

I am trying to use SWIN Transformers for a problem that maps every image to a score of 1 to 100
When in initialize the model, everything works fine. But after training 1 epoch. the model gives the same output for every image.
I tried changing a few parameters.
Any suggestion what could be the reason behind this?

Thanks

[ancientmooner]

Can you please provide more information, for example, your loss function, training configuration, loss curves, etc?

[sulakshgupta988]

This is the model creation
model = swin_transformer.SwinTransformer(img_size=224, patch_size=4, in_chans=3, num_classes=0, embed_dim=96, depths=[2,2,6,2], num_heads=[3,6,12,24], window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0.1, drop_path_rate=0.1, ape=False, patch_norm=True, use_checkpoint=False)

Below is the model configuration:
optimizer = build_optimizer(model)
n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"number of params: {n_parameters}")
num_steps = int(20 * len(train_set))
warmup_steps = int(20 * len(train_set))
decay_steps = int(3 *len(train_set))
lr_scheduler = StepLRScheduler( optimizer, decay_t=decay_steps, decay_rate=0.1, warmup_lr_init=0.001, warmup_t=warmup_steps, t_in_epochs=False, )
criterion = nn.MSELoss()
model.float().cuda()
num_epochs = 50

Training Loop:

for epoch in range(num_epochs):
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(n_parameters)
    for i, (samples, score) in enumerate(train_set):
      optimizer.zero_grad()
      samples = samples.cuda()
      score = score.cuda()
      outputs = model(samples)
      loss = criterion(outputs.squeeze(), score)
      loss.backward()
      optimizer.step()
      lr_scheduler.step_update(epoch * num_steps + i)
      torch.cuda.synchronize()
      lr = optimizer.param_groups[0]['lr']
    #train_one_epoch(model, criterion, train_set, optimizer, epoch, lr_scheduler, num_epochs)
    acc1, acc5, loss, corr, spear = validate(val_set, model)
    print(f"Correlation on {len(val_set)} test images: {corr:.7f}%" + "     " +str(spear))
    f = open("temp.txt", "a+")
    f.write("\n"+str(corr)+"\n")
    f.close()

Loss Function used is mean squared error.
build_optimizer function is taken from the build_optimizer.py file in the repo.
len(train_set) is around 500.
len(val_set) is around 150.