Мобильная версияПроблема: Я столкнулся с проблемой сходимости. Потери начинаются с высоких значений и слегка уменьшаются (например, с 60 до 10) в первые несколько эпох, но затем выходят на плато и отказываются снижаться дальше. Модель не обучается, даже когда я пытаюсь переобучить ее в обучающем наборе.
Мой подход:
- Я загружаю предварительно обученную модель с помощью YOLO('yolov8n.pt').model.
- Я замораживаю первые 10 слоев магистрали.
- Я обновляю конфигурацию cfg.nc = 7.
- Я использую v8DetectionLoss со стандартным циклом обучения PyTorch.
from ultralytics import YOLO
from torch.utils.data import DataLoader
from dataset import WaterDataset
import torch.optim as opt
from torchvision import transforms
from tqdm import tqdm
import torch
from torchvision.transforms import RandomAffine, ColorJitter
import argparse
from ultralytics.utils.loss import v8DetectionLoss
from torch.optim.lr_scheduler import CosineAnnealingLR
from ultralytics.cfg import get_cfg
from ultralytics.utils import DEFAULT_CFG_PATH
def args():
parser = argparse.ArgumentParser()
parser.add_argument('--root', type=str, default='aquarium_pretrain')
return parser.parse_args()
def collate_fn(batch):
imgs, labs = zip(*batch) # imgs: tuple of tensors (C,H,W)
imgs = torch.stack([img if torch.is_tensor(img) else torch.from_numpy(img).permute(2,0,1) for img in imgs], dim=0) # [B,3,H,W]
# labs should already be tensors of shape (n,5) or empty; keep as list
return imgs, list(labs)
def train():
model_yl = YOLO('yolov8n.pt')
model = model_yl.model
model.train()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
freeze_idx = 10
for k, v in model.named_parameters():
v.requires_grad = True
# Sau đó đi khóa lại các lớp Backbone
for i, child in enumerate(model.model):
if i < freeze_idx:
for param in child.parameters():
param.requires_grad = False
print(f"Layer {i}: FROZEN")
else:
print(f"Layer {i}: TRAINABLE")
cfg = get_cfg(DEFAULT_CFG_PATH)
cfg.nc = 7
model.args = cfg
transform_train = transforms.Compose([
transforms.Resize((640, 640)),
# RandomAffine(
# degrees=(-5, 5),
# translate=(0.15, 0.15),
# scale=(0.85, 1.15),
# shear=10
# ),
# ColorJitter(
# brightness=0.125,
# contrast=0.5,
# saturation=0.5,
# hue=0.05
# ),
transforms.ToTensor(),
])
transform_val = transforms.Compose([
transforms.Resize((640, 640)),
transforms.ToTensor(),
])
train_dataset = WaterDataset(type='train', transform=transform_train, root=args.root)
val_dataset = WaterDataset(type='valid', transform= transform_val ,root=args.root)
model.names = {i: name for i, name in enumerate(train_dataset.categories)}
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=2, collate_fn=collate_fn)
val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False, num_workers=2, collate_fn=collate_fn)
optimizer = opt.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)
criterion = v8DetectionLoss(model)
for epoch in range(0,100):
progress_bar = tqdm(train_loader, desc="Training")
for images, labels in progress_bar:
images = images.to(device)
lb ={
"batch_idx":[],
"cls":[],
"bboxes":[]
}
for i, label in enumerate(labels):
for l in label:
labell = l.to(device)
lb["batch_idx"].append(i)
lb["cls"].append(labell[0])
lb["bboxes"].append(labell[1:])
lb["batch_idx"] = torch.tensor(lb["batch_idx"]).to(device)
lb["cls"] = torch.tensor(lb["cls"]).to(device)
lb["bboxes"] = torch.stack(lb["bboxes"]).to(device)
preds = model(images)
loss, _ = criterion(preds, lb)
optimizer.zero_grad()
loss.sum().backward()
optimizer.step()
progress_bar.set_description(
"Epoch {}/{}. Loss {:0.4f}".format(epoch + 1, 100, loss.sum().mean()))
scheduler.step()
if __name__ == '__main__':
args = args()
train()
Подробнее здесь: https://stackoverflow.com/questions/798 ... -custom-da
