来自segmentation_models_pytorch的Unet在训练中停滞

Question

我一直在关注有关在自定义数据集上训练分割模型的教程，但它拒绝在模型训练方面取得任何进展。

这是我的模型设置

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),
])

还有时代跑步者

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,
)

而且，当我进行训练时，模型只是停留在第一个时期，没有任何进展

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!')

结果：

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

我就这样放置了3个小时，它没有一点变化

我在具有 32 GB RAM 的 CPU (i7-10710U) 上运行（我知道它比 GPU 慢得多，但我的 GPU (GeForce 1650mq) 不支持 cuda），并且我之前运行过类似的模型，没有任何问题。

有人可以帮我吗？也许我错过了什么？也许有一个更轻的模型可以在我的系统上运行？

我已经尝试过其他一些设置和模型，YOLOv8和YOLOv3也拒绝训练。

Answer 1

我也面临着类似的问题。我在 80GB A100 GPU 上运行我的代码，执行速度非常慢。