I have been following a tutorial on training a segmentation model on a custom dataset, but it refuses to make any progrees in training the model.

This is my model setup

import segmentation_models_pytorch as smp
import torch

ENCODER = 'efficientnet-b0'
ENCODER_WEIGHTS = 'imagenet'
CLASSES = ['ship']
ACTIVATION = 'sigmoid'
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

model = smp.Unet(
    encoder_name=ENCODER, 
    encoder_weights=ENCODER_WEIGHTS, 
    classes=len(CLASSES), 
    activation=ACTIVATION,
).to(DEVICE)

from segmentation_models_pytorch import utils as smp_utils

loss = smp_utils.losses.DiceLoss()
metrics = [
    smp_utils.metrics.IoU(threshold=0.5),
]

optimizer = torch.optim.Adam([ 
    dict(params=model.parameters(), lr=0.0001),
])

And epochs runners

train_epoch = smp_utils.train.TrainEpoch(
    model, 
    loss=loss, 
    metrics=metrics, 
    optimizer=optimizer,
    device=DEVICE,
    verbose=True,
)

valid_epoch = smp_utils.train.ValidEpoch(
    model, 
    loss=loss, 
    metrics=metrics, 
    device=DEVICE,
    verbose=True,
)

And, when I run the training, the model just sits on firts epoch with no progress

max_score = 0

for i in range(0, 40):
    
    print('nEpoch: {}'.format(i))
    train_logs = train_epoch.run(train_loader)
    valid_logs = valid_epoch.run(valid_loader)
    
    # do something (save model, change lr, etc.)
    if max_score < valid_logs['iou_score']:
        max_score = valid_logs['iou_score']
        torch.save(model, './best_model.pth')
        print('Model saved!')
        
    if i == 25:
        optimizer.param_groups[0]['lr'] = 1e-5
        print('Decrease decoder learning rate to 1e-5!')

Result:

Epoch: 0
train:   0%|                                                                                                                                                                                                | 0/3851 [00:00<?, ?it/s]

I left it for 3 hours like this and it did not change a bit

I am running on CPU (i7-10710U) (I know it is way slower than GPU, but my GPU (GeForce 1650mq) does not support cuda) with 32 gigs of RAM, and I had simmilar models run before without any problems.

Can someone help me out? Maybe I am missing something? Maybe there is a lighter model that can run on my system?

I’ve already tried some other setups and models, YOLOv8 and YOLOv3 also refused to train.