pokouqiege/models/line_net/trainer.py

import os
import time
from datetime import datetime

import numpy as np
import torch
from matplotlib import pyplot as plt
from torch.utils.tensorboard import SummaryWriter

from libs.vision_libs.utils import draw_bounding_boxes
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_net.dataset_LD import WirePointDataset
from models.wirenet.postprocess import postprocess
from tools import utils
from torchvision import transforms
import matplotlib as mpl

cmap = plt.get_cmap("jet")
norm = mpl.colors.Normalize(vmin=0.4, vmax=1.0)
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
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
def c(x):
    return sm.to_rgba(x)


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


class Trainer(BaseTrainer):
    def __init__(self, model=None, **kwargs):
        super().__init__(model, device, **kwargs)
        self.model = model
        # print(f'kwargs:{kwargs}')
        self.init_params(**kwargs)

    def init_params(self, **kwargs):
        if kwargs != {}:
            print(f'train_params:{kwargs["train_params"]}')
            self.freeze_config = kwargs['train_params']['freeze_params']
            print(f'freeze_config:{self.freeze_config}')
            self.dataset_path = kwargs['io']['datadir']
            self.batch_size = kwargs['train_params']['batch_size']
            self.num_workers = kwargs['train_params']['num_workers']
            self.logdir = kwargs['io']['logdir']
            self.resume_from = kwargs['train_params']['resume_from']
            self.optim = ''
            self.train_result_ptath = os.path.join(self.logdir, datetime.now().strftime("%Y%m%d_%H%M%S"))
            self.wts_path = os.path.join(self.train_result_ptath, 'weights')
            self.tb_path = os.path.join(self.train_result_ptath, 'logs')
            self.writer = SummaryWriter(self.tb_path)
            self.last_model_path = os.path.join(self.wts_path, 'last.pth')
            self.best_train_model_path = os.path.join(self.wts_path, 'best_train.pth')
            self.best_val_model_path = os.path.join(self.wts_path, 'best_val.pth')
            self.max_epoch = kwargs['train_params']['max_epoch']

    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 freeze_params(self, model):
        """根据配置冻结模型参数"""
        default_config = {
            'backbone': True,  # 冻结 backbone
            'rpn': False,  # 不冻结 rpn
            'roi_heads': {
                'box_head': False,
                'box_predictor': False,
                'line_head': False,
                'line_predictor': {
                    'fc1': False,
                    'fc2': {
                        '0': False,
                        '2': False,
                        '4': False
                    }
                }
            }
        }

        # 更新默认配置
        default_config.update(self.freeze_config)
        config = default_config

        print("\n===== Parameter Freezing Configuration =====")
        for name, module in model.named_children():
            if name in config:
                if isinstance(config[name], bool):
                    for param in module.parameters():
                        param.requires_grad = not config[name]
                    print(f"{'Frozen' if config[name] else 'Trainable'} module: {name}")

                elif isinstance(config[name], dict):
                    for subname, submodule in module.named_children():
                        if subname in config[name]:
                            if isinstance(config[name][subname], bool):
                                for param in submodule.parameters():
                                    param.requires_grad = not config[name][subname]
                                print(
                                    f"{'Frozen' if config[name][subname] else 'Trainable'} submodule: {name}.{subname}")

                            elif isinstance(config[name][subname], dict):
                                for subsubname, subsubmodule in submodule.named_children():
                                    if subsubname in config[name][subname]:
                                        for param in subsubmodule.parameters():
                                            param.requires_grad = not config[name][subname][subsubname]
                                        print(
                                            f"{'Frozen' if config[name][subname][subsubname] else 'Trainable'} sub-submodule: {name}.{subname}.{subsubname}")

        # 打印参数统计
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"\nTotal Parameters: {total_params:,}")
        print(f"Trainable Parameters: {trainable_params:,}")
        print(f"Frozen Parameters: {total_params - trainable_params:,}")

    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_predict_result(self, img, result, epoch):
        img = img.cpu().detach()
        im = img.permute(1, 2, 0)
        self.writer.add_image("z-ori", im, epoch, dataformats="HWC")

        boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), result[0]["boxes"],
                                          colors="yellow", width=1)
        self.writer.add_image("z-boxes", boxed_image.permute(1, 2, 0), epoch, dataformats="HWC")

        PLTOPTS = {"color": "#33FFFF", "s": 15, "edgecolors": "none", "zorder": 5}
        # print(f'pred[1]:{pred[1]}')
        heatmaps = result[-2][0]
        print(f'heatmaps:{heatmaps.shape}')
        jmap = heatmaps[1: 2].cpu().detach()
        lmap = heatmaps[2: 3].cpu().detach()
        self.writer.add_image("z-jmap", jmap, epoch)
        self.writer.add_image("z-lmap", lmap, epoch)
        # plt.imshow(lmap)
        # plt.show()
        H = result[-1]['wires']
        # lines = H["lines"][0].cpu().numpy()
        lines=result[0]["lines"]
        scores =100
        for i in range(1, len(lines)):
            if (lines[i] == lines[0]).all():
                lines = lines[:i]
                scores = scores[:i]
                break

        # postprocess lines to remove overlapped lines
        diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
        nlines, nscores = postprocess(lines, scores, diag * 0.01, 0, False)

        for i, t in enumerate([0]):
            plt.gca().set_axis_off()
            plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
            plt.margins(0, 0)
            for (a, b), s in zip(nlines, nscores):
                if s < t:
                    continue
                plt.plot([a[1], b[1]], [a[0], b[0]], c=c(s), linewidth=2, zorder=s)
                plt.scatter(a[1], a[0], **PLTOPTS)
                plt.scatter(b[1], b[0], **PLTOPTS)
            plt.gca().xaxis.set_major_locator(plt.NullLocator())
            plt.gca().yaxis.set_major_locator(plt.NullLocator())
            plt.imshow(im)
            plt.tight_layout()
            fig = plt.gcf()
            fig.canvas.draw()

            width, height = fig.get_size_inches() * fig.get_dpi()  # 获取图像尺寸
            tmp_img = fig.canvas.tostring_argb()
            tmp_img_np = np.frombuffer(tmp_img, dtype=np.uint8)
            tmp_img_np = tmp_img_np.reshape(int(height), int(width), 4)

            img_rgb = tmp_img_np[:, :, 1:]  # 提取RGB部分，忽略Alpha通道

            # image_from_plot = np.frombuffer(tmp_img[:,:,1:], dtype=np.uint8).reshape(
            #     fig.canvas.get_width_height()[::-1] + (3,))
            plt.close()

            img2 = transforms.ToTensor()(img_rgb)

            self.writer.add_image("z-output", img2, epoch)

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

    def train_from_cfg(self, model: BaseModel, cfg, freeze_config=None):  # 新增：支持传入冻结配置
        cfg = read_yaml(cfg)
        # print(f'cfg:{cfg}')
        # self.freeze_config = freeze_config or {}  # 更新冻结配置

        self.train(model, **cfg)

    def train(self, model, **kwargs):

        self.init_params(**kwargs)

        dataset_train = WirePointDataset(dataset_path=self.dataset_path, dataset_type='train')
        dataset_val = WirePointDataset(dataset_path=self.dataset_path, dataset_type='val')

        train_sampler = torch.utils.data.RandomSampler(dataset_train)
        val_sampler = torch.utils.data.RandomSampler(dataset_val)
        train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=self.batch_size, drop_last=True)
        val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=self.batch_size, drop_last=True)
        train_collate_fn = utils.collate_fn
        val_collate_fn = utils.collate_fn

        data_loader_train = torch.utils.data.DataLoader(
            dataset_train, batch_sampler=train_batch_sampler, num_workers=self.num_workers, collate_fn=train_collate_fn
        )
        data_loader_val = torch.utils.data.DataLoader(
            dataset_val, batch_sampler=val_batch_sampler, num_workers=self.num_workers, collate_fn=val_collate_fn
        )

        model.to(device)

        optimizer = torch.optim.Adam(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=kwargs['train_params']['optim']['lr']
        )

        for epoch in range(self.max_epoch):
            print(f"train epoch:{epoch}")

            model, epoch_train_loss = self.one_epoch(model, data_loader_train, epoch, optimizer)

            # ========== Validation ==========
            with torch.no_grad():
                model, epoch_val_loss = self.one_epoch(model, data_loader_val, epoch, optimizer, phase='val')

            if epoch==0:
                best_train_loss = epoch_train_loss
                best_val_loss = epoch_val_loss

            self.save_last_model(model,self.last_model_path, epoch, optimizer)
            best_train_loss = self.save_best_model(model, self.best_train_model_path, epoch, epoch_train_loss,
                                                   best_train_loss,
                                                   optimizer)
            best_val_loss = self.save_best_model(model, self.best_val_model_path, epoch, epoch_val_loss, best_val_loss,
                                                 optimizer)

    def one_epoch(self, model, data_loader, epoch, optimizer, phase='train'):
        if phase == 'train':
            model.train()
        if phase == 'val':
            model.eval()

        total_loss = 0
        epoch_step = 0
        global_step = epoch * len(data_loader)
        for imgs, targets in data_loader:
            imgs = self.move_to_device(imgs, device)
            targets = self.move_to_device(targets, device)
            if phase== 'val':

                result,losses = model(imgs, targets)
            else:
                losses = model(imgs, targets)

            loss = _loss(losses)
            total_loss += loss.item()
            if phase == 'train':
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            self.writer_loss(losses, global_step, phase=phase)
            global_step += 1

            if epoch_step == 0 and phase == 'val':
                t_start = time.time()
                print(f'start to predict:{t_start}')
                result = model(self.move_to_device(imgs, self.device))
                t_end = time.time()
                print(f'predict used:{t_end - t_start}')
                self.writer_predict_result(img=imgs[0], result=result, epoch=epoch)
                epoch_step+=1

        avg_loss = total_loss / len(data_loader)
        print(f'{phase}/loss epoch{epoch}:{avg_loss:4f}')
        self.writer.add_scalar(f'loss/{phase}', avg_loss, epoch)
        return model, avg_loss

    def save_best_model(self, model, save_path, epoch, current_loss, best_loss, optimizer=None):
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        if current_loss <= best_loss:
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'loss': current_loss
            }
            if optimizer is not None:
                checkpoint['optimizer_state_dict'] = optimizer.state_dict()

            torch.save(checkpoint, save_path)
            print(f"Saved best model at epoch {epoch} with loss {current_loss:.4f}")

            return current_loss

        return best_loss

    def save_last_model(self, model, save_path, epoch, optimizer=None):

        if os.path.exists(f'{self.wts_path}/last.pt'):
            os.remove(f'{self.wts_path}/last.pt')

        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
        }

        if optimizer is not None:
            checkpoint['optimizer_state_dict'] = optimizer.state_dict()

        torch.save(checkpoint, save_path)


if __name__ == '__main__':
    print('')