lcnn/models/line_detect/trainer.py

import os
import time
from datetime import datetime

import torch
from torch.utils.tensorboard import SummaryWriter

from models.base.base_model import BaseModel
from models.base.base_trainer import BaseTrainer
from models.config.config_tool import read_yaml
from models.line_detect.dataset_LD import WirePointDataset
from models.line_detect.postprocess import box_line_, show_
from utils.log_util import show_line, save_last_model, save_best_model
from tools import utils


def _loss(losses):
    total_loss = 0
    for i in losses.keys():
        if i != "loss_wirepoint":
            total_loss += losses[i]
        else:
            loss_labels = losses[i]["losses"]
    loss_labels_k = list(loss_labels[0].keys())
    for j, name in enumerate(loss_labels_k):
        loss = loss_labels[0][name].mean()
        total_loss += loss

    return total_loss


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def move_to_device(data, device):
    if isinstance(data, (list, tuple)):
        return type(data)(move_to_device(item, device) for item in data)
    elif isinstance(data, dict):
        return {key: move_to_device(value, device) for key, value in data.items()}
    elif isinstance(data, torch.Tensor):
        return data.to(device)
    else:
        return data  # 对于非张量类型的数据不做任何改变


class Trainer(BaseTrainer):
    def __init__(self, model=None,
                 dataset=None,
                 device='cuda',
                 **kwargs):

        super().__init__(model, dataset, device, **kwargs)

    def move_to_device(self, data, device):
        if isinstance(data, (list, tuple)):
            return type(data)(self.move_to_device(item, device) for item in data)
        elif isinstance(data, dict):
            return {key: self.move_to_device(value, device) for key, value in data.items()}
        elif isinstance(data, torch.Tensor):
            return data.to(device)
        else:
            return data  # 对于非张量类型的数据不做任何改变

    def load_best_model(self, model, optimizer, save_path, device):
        if os.path.exists(save_path):
            checkpoint = torch.load(save_path, map_location=device)
            model.load_state_dict(checkpoint['model_state_dict'])
            if optimizer is not None:
                optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            epoch = checkpoint['epoch']
            loss = checkpoint['loss']
            print(f"Loaded best model from {save_path} at epoch {epoch} with loss {loss:.4f}")
        else:
            print(f"No saved model found at {save_path}")
        return model, optimizer

    def writer_loss(self, writer, losses, epoch):
        try:
            for key, value in losses.items():
                if key == 'loss_wirepoint':
                    for subdict in losses['loss_wirepoint']['losses']:
                        for subkey, subvalue in subdict.items():
                            writer.add_scalar(f'loss/{subkey}',
                                              subvalue.item() if hasattr(subvalue, 'item') else subvalue,
                                              epoch)
                elif isinstance(value, torch.Tensor):
                    writer.add_scalar(f'loss/{key}', value.item(), epoch)
        except Exception as e:
            print(f"TensorBoard logging error: {e}")

    def train_cfg(self, model: BaseModel, cfg):
        # cfg = r'./config/wireframe.yaml'
        cfg = read_yaml(cfg)
        print(f'cfg:{cfg}')
        # print(cfg['n_dyn_negl'])
        self.train(model, **cfg)

    # def train(self, model, **kwargs):
    #     dataset_train = WirePointDataset(dataset_path=kwargs['io']['datadir'], dataset_type='train')
    #     train_sampler = torch.utils.data.RandomSampler(dataset_train)
    #     # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
    #     train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=64, drop_last=True)
    #     train_collate_fn = utils.collate_fn_wirepoint
    #     data_loader_train = torch.utils.data.DataLoader(
    #         dataset_train, batch_sampler=train_batch_sampler, num_workers=8, collate_fn=train_collate_fn
    #     )
    #
    #     dataset_val = WirePointDataset(dataset_path=kwargs['io']['datadir'], dataset_type='val')
    #     val_sampler = torch.utils.data.RandomSampler(dataset_val)
    #     # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
    #     val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=64, drop_last=True)
    #     val_collate_fn = utils.collate_fn_wirepoint
    #     data_loader_val = torch.utils.data.DataLoader(
    #         dataset_val, batch_sampler=val_batch_sampler, num_workers=8, collate_fn=val_collate_fn
    #     )
    #
    #     train_result_ptath = os.path.join('train_results', datetime.now().strftime("%Y%m%d_%H%M%S"))
    #     wts_path = os.path.join(train_result_ptath, 'weights')
    #     tb_path = os.path.join(train_result_ptath, 'logs')
    #     writer = SummaryWriter(tb_path)
    #
    #     optimizer = torch.optim.Adam(model.parameters(), lr=kwargs['optim']['lr'])
    #     # writer = SummaryWriter(kwargs['io']['logdir'])
    #     model.to(device)
    #
    #
    #
    #     # # 加载权重
    #     # save_path = 'logs/pth/best_model.pth'
    #     # model, optimizer = self.load_best_model(model, optimizer, save_path, device)
    #
    #     # logdir_with_pth = os.path.join(kwargs['io']['logdir'], 'pth')
    #     # os.makedirs(logdir_with_pth, exist_ok=True)  # 创建目录（如果不存在）
    #     last_model_path = os.path.join(wts_path, 'last.pth')
    #     best_model_path = os.path.join(wts_path, 'best.pth')
    #     global_step = 0
    #
    #     for epoch in range(kwargs['optim']['max_epoch']):
    #         print(f"epoch:{epoch}")
    #         total_train_loss = 0.0
    #
    #         model.train()
    #
    #         for imgs, targets in data_loader_train:
    #             imgs = move_to_device(imgs, device)
    #             targets=move_to_device(targets,device)
    #             # print(f'imgs:{len(imgs)}')
    #             # print(f'targets:{len(targets)}')
    #             losses = model(imgs, targets)
    #             loss = _loss(losses)
    #             total_train_loss += loss.item()
    #             optimizer.zero_grad()
    #             loss.backward()
    #             optimizer.step()
    #             self.writer_loss(writer, losses, global_step)
    #             global_step+=1
    #
    #
    #         avg_train_loss = total_train_loss / len(data_loader_train)
    #         if epoch == 0:
    #             best_loss = avg_train_loss;
    #
    #         writer.add_scalar('loss/train', avg_train_loss, epoch)
    #
    #
    #         if os.path.exists(f'{wts_path}/last.pt'):
    #             os.remove(f'{wts_path}/last.pt')
    #         # torch.save(model.state_dict(), f'{wts_path}/last.pt')
    #         save_last_model(model,last_model_path,epoch,optimizer)
    #         best_loss = save_best_model(model,best_model_path,epoch,avg_train_loss,best_loss,optimizer)
    #
    #         model.eval()
    #         with torch.no_grad():
    #             for batch_idx, (imgs, targets) in enumerate(data_loader_val):
    #                 t_start = time.time()
    #                 print(f'start to predict:{t_start}')
    #                 pred = model(self.move_to_device(imgs, self.device))
    #                 t_end = time.time()
    #                 print(f'predict used:{t_end - t_start}')
    #                 if batch_idx == 0:
    #                     show_line(imgs[0], pred, epoch, writer)
    #                 break

    def train(self, model, **kwargs):
        default_params = {
            'io': {
                'logdir': 'logs /',
                'datadir': '/ root / autodl - tmp / wirenet_rgb_gray',
                'num_workers': 8,
                'tensorboard_port': 6000,
                'validation_interval': 300,
                'batch_size': 4,
                'batch_size_eval': 2,
            },
            'optim':{
                'name': 'Adam',
                'lr': 4.0e-4,
                'amsgrad': True,
                'weight_decay': 1.0e-4,
                'max_epoch': 90000000,
                'lr_decay_epoch': 10,
            },
        }

        # 更新默认参数
        for key, value in kwargs.items():
            if key in default_params:
                default_params[key] = value
            else:
                raise ValueError(f"Unknown argument: {key}")

        # 解析参数
        dataset_path = default_params['io']['datadir']
        num_workers = default_params['io']['num_workers']
        batch_size_train = default_params['io']['batch_size']
        batch_size_eval = default_params['io']['batch_size_eval']
        epochs = default_params['optim']['max_epoch']
        lr = default_params['optim']['lr']

        dataset_train = WirePointDataset(dataset_path=dataset_path, dataset_type='train')
        train_sampler = torch.utils.data.RandomSampler(dataset_train)
        # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
        train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=batch_size_train, drop_last=True)
        train_collate_fn = utils.collate_fn_wirepoint
        data_loader_train = torch.utils.data.DataLoader(
            dataset_train, batch_sampler=train_batch_sampler, num_workers=num_workers, collate_fn=train_collate_fn
        )

        dataset_val = WirePointDataset(dataset_path=dataset_path, dataset_type='val')
        val_sampler = torch.utils.data.RandomSampler(dataset_val)
        # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
        val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=batch_size_eval, drop_last=True)
        val_collate_fn = utils.collate_fn_wirepoint
        data_loader_val = torch.utils.data.DataLoader(
            dataset_val, batch_sampler=val_batch_sampler, num_workers=num_workers, collate_fn=val_collate_fn
        )

        train_result_ptath = os.path.join('train_results', datetime.now().strftime("%Y%m%d_%H%M%S"))
        wts_path = os.path.join(train_result_ptath, 'weights')
        tb_path = os.path.join(train_result_ptath, 'logs')
        writer = SummaryWriter(tb_path)

        optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        # writer = SummaryWriter(kwargs['io']['logdir'])
        model.to(device)

        # # 加载权重
        # save_path = 'logs/pth/best_model.pth'
        # model, optimizer = self.load_best_model(model, optimizer, save_path, device)

        # logdir_with_pth = os.path.join(kwargs['io']['logdir'], 'pth')
        # os.makedirs(logdir_with_pth, exist_ok=True)  # 创建目录（如果不存在）
        last_model_path = os.path.join(wts_path, 'last.pth')
        best_model_path = os.path.join(wts_path, 'best.pth')
        global_step = 0

        for epoch in range(epochs):
            print(f"epoch:{epoch}")
            total_train_loss = 0.0

            model.train()

            for imgs, targets in data_loader_train:
                imgs = move_to_device(imgs, device)
                targets = move_to_device(targets, device)
                # print(f'imgs:{len(imgs)}')
                # print(f'targets:{len(targets)}')
                losses = model(imgs, targets)
                loss = _loss(losses)
                total_train_loss += loss.item()
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                self.writer_loss(writer, losses, global_step)
                global_step += 1

            avg_train_loss = total_train_loss / len(data_loader_train)
            if epoch == 0:
                best_loss = avg_train_loss;

            writer.add_scalar('loss/train', avg_train_loss, epoch)

            if os.path.exists(f'{wts_path}/last.pt'):
                os.remove(f'{wts_path}/last.pt')
            # torch.save(model.state_dict(), f'{wts_path}/last.pt')
            save_last_model(model, last_model_path, epoch, optimizer)
            best_loss = save_best_model(model, best_model_path, epoch, avg_train_loss, best_loss, optimizer)

            model.eval()
            with torch.no_grad():
                for batch_idx, (imgs, targets) in enumerate(data_loader_val):
                    t_start = time.time()
                    print(f'start to predict:{t_start}')
                    pred = model(self.move_to_device(imgs, self.device))
                    t_end = time.time()
                    print(f'predict used:{t_end - t_start}')
                    if batch_idx == 0:
                        show_line(imgs[0], pred, epoch, writer)
                    break