train.py

import torch
import torch.nn as nn
from pytorch_lightning import LightningModule, Trainer, seed_everything
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint, TQDMProgressBar
from pytorch_lightning.loggers import TensorBoardLogger
import torchmetrics
from tqdm import tqdm

import prettytable
import numpy as np
import argparse
from rscd.models.build_model import build_model
from rscd.datasets import build_dataloader
from rscd.optimizers import build_optimizer
from rscd.losses import build_loss
from utils.config import Config

from torch.autograd import Variable

import sys
sys.path.append('rscd')

seed_everything(1234, workers=True)

import numpy as np
import os

import time  # 用于计时

def resize_label(label, size):

    label = np.expand_dims(label,axis=0)
    label_resized = np.zeros((1,label.shape[1],size[0],size[1]))
    interp = nn.Upsample(size=(size[0], size[1]),mode='bilinear')

    labelVar = Variable(torch.from_numpy(label).float())  
    label_resized[:, :,:,:] = interp(labelVar).data.numpy()
    label_resized = np.array(label_resized, dtype=np.int32)
    return torch.from_numpy(np.squeeze(label_resized,axis=0)).float()

def get_args():
    parser = argparse.ArgumentParser('description=Change detection of remote sensing images')
    parser.add_argument("-c", "--config", type=str, default="configs/cdlamba.py")
    return parser.parse_args()

class myTrain(LightningModule):
    def __init__(self, cfg, log_dir = None):
        super(myTrain, self).__init__()
        
        self.cfg = cfg
        self.log_dir = log_dir
        self.net = build_model(cfg.model_config)
        self.loss = build_loss(cfg.loss_config)

        self.loss.to('cuda:{}'.format(cfg.gpus[0]))
        
        metric_cfg1 = cfg.metric_cfg1
        metric_cfg2 = cfg.metric_cfg2
        
        self.tr_oa=torchmetrics.Accuracy(**metric_cfg1)
        self.tr_prec = torchmetrics.Precision(**metric_cfg2)
        self.tr_recall = torchmetrics.Recall(**metric_cfg2)
        self.tr_f1 = torchmetrics.F1Score(**metric_cfg2)
        self.tr_iou=torchmetrics.JaccardIndex(**metric_cfg2)

        self.val_oa=torchmetrics.Accuracy(**metric_cfg1)
        self.val_prec = torchmetrics.Precision(**metric_cfg2)
        self.val_recall = torchmetrics.Recall(**metric_cfg2)
        self.val_f1 = torchmetrics.F1Score(**metric_cfg2)
        self.val_iou=torchmetrics.JaccardIndex(**metric_cfg2)

        self.test_oa=torchmetrics.Accuracy(**metric_cfg1)
        self.test_prec = torchmetrics.Precision(**metric_cfg2)
        self.test_recall = torchmetrics.Recall(**metric_cfg2)
        self.test_f1 = torchmetrics.F1Score(**metric_cfg2)
        self.test_iou=torchmetrics.JaccardIndex(**metric_cfg2)

        self.test_max_f1 = [0 for _ in range(10)]

        self.test_loader = build_dataloader(cfg.dataset_config, mode='test')

    def forward(self, x1, x2) :
        pred = self.net(x1, x2)
        return pred

    def configure_optimizers(self):
        optimizer, scheduler = build_optimizer(self.cfg.optimizer_config, self.net)
        return {'optimizer':optimizer,'lr_scheduler':scheduler, 'monitor': self.cfg.monitor_val}

    def train_dataloader(self):
        loader = build_dataloader(self.cfg.dataset_config, mode='train')
        return loader

    def val_dataloader(self):
        loader = build_dataloader(self.cfg.dataset_config, mode='val')
        return loader

    def output(self, metrics, total_metrics, mode, test_idx=0, test_value=None):
        result_table = prettytable.PrettyTable()
        result_table.field_names = ['Class', 'OA', 'Precision', 'Recall', 'F1_Score', 'IOU']

        for i in range(len(metrics[0])):
            item = [i, '--']
            for j in range(len(metrics)):
                item.append(np.round(metrics[j][i].cpu().numpy(), 4))
            result_table.add_row(item)

        total = list(total_metrics.values())
        total = [np.round(v, 4) for v in total]
        total.insert(0, 'total')
        result_table.add_row(total)

        if mode == 'val' or mode == 'test':
            print(mode)
            print(result_table)

        if self.log_dir:
            base_dir = self.log_dir
        else:
            base_dir = os.path.join('work_dirs', cfg.exp_name)

        if mode == 'test':
            if self.cfg.argmax:
                file_name = os.path.join(base_dir, "test_metrics_{}.txt".format(test_idx)) 
                if metrics[2][1] > self.test_max_f1[test_idx]:
                    self.test_max_f1[test_idx] = metrics[2][1]
                    file_name = os.path.join(base_dir, "test_max_metrics_{}.txt".format(test_idx))
            else:
                file_name = os.path.join(base_dir, "test_metrics_{}_{}.txt".format(test_idx, str(test_value))) 
                if metrics[2][1] > self.test_max_f1[test_idx]:
                    self.test_max_f1[test_idx] = metrics[2][1]
                    file_name = os.path.join(base_dir, "test_max_metrics_{}_{}.txt".format(test_idx, '%.1f' % test_value))
        else:
            file_name = os.path.join(base_dir, "train_metrics.txt") 
        f = open(file_name,"a")
        f.write('epoch:{}/{} {}\n'.format(self.current_epoch, self.cfg.epoch, mode))
        f.write(str(result_table)+'\n')
        f.close()

    def training_step(self, batch, batch_idx):
        imgA, imgB, mask = batch[0], batch[1], batch[2]
        preds = self(imgA, imgB)

        if self.cfg.net == 'SARASNet':
            mask = Variable(resize_label(mask.data.cpu().numpy(), \
                                    size=preds.data.cpu().numpy().shape[2:]).to('cuda')).long()
            param = 1  # This parameter is balance precision and recall to get higher F1-score
            preds[:,1,:,:] = preds[:,1,:,:] + param 

        if self.cfg.argmax:
            loss = self.loss(preds, mask)
            pred = preds.argmax(dim=1)
        else:
            if self.cfg.net == 'maskcd':
                loss = self.loss(preds[1], mask)
                pred = preds[0]
                pred = pred > 0.5
                pred.squeeze_(1)
            else:
                pred = preds.squeeze(1)
                loss = self.loss(pred, mask)
                pred = pred > 0.5

        self.tr_oa(pred, mask)
        self.tr_prec(pred, mask)
        self.tr_recall(pred, mask)
        self.tr_f1(pred, mask)
        self.tr_iou(pred, mask)

        self.log('tr_loss', loss, on_step=True,on_epoch=True,prog_bar=True)
        return loss

    def on_train_epoch_end(self):
        metrics = [self.tr_prec.compute(),
                   self.tr_recall.compute(),
                   self.tr_f1.compute(),
                   self.tr_iou.compute()]
        
        log = {'tr_oa': float(self.tr_oa.compute().cpu()),
               'tr_prec': np.mean([item.cpu() for item in metrics[0]]),
               'tr_recall': np.mean([item.cpu() for item in metrics[1]]),
               'tr_f1': np.mean([item.cpu() for item in metrics[2]]),
               'tr_miou': np.mean([item.cpu() for item in metrics[3]])}
        
        self.output(metrics, log, 'train')
        
        for key, value in zip(log.keys(), log.values()):
            self.log(key, value, on_step=False,on_epoch=True,prog_bar=True)
        self.log('tr_change_f1', metrics[2][1], on_step=False,on_epoch=True,prog_bar=True)

        self.tr_oa.reset()
        self.tr_prec.reset()
        self.tr_recall.reset()
        self.tr_f1.reset()
        self.tr_iou.reset()

    def validation_step(self, batch, batch_idx):
        imgA, imgB, mask = batch[0], batch[1], batch[2]
        preds = self(imgA, imgB)

        if self.cfg.net == 'SARASNet':
            mask = Variable(resize_label(mask.data.cpu().numpy(), \
                                    size=preds.data.cpu().numpy().shape[2:]).to('cuda')).long()
            param = 1  # This parameter is balance precision and recall to get higher F1-score
            preds[:,1,:,:] = preds[:,1,:,:] + param 

        if self.cfg.argmax:
            loss = self.loss(preds, mask)
            pred = preds.argmax(dim=1)
        else:
            if self.cfg.net == 'maskcd':
                loss = self.loss(preds[1], mask)
                pred = preds[0]
                pred = pred > 0.5
                pred.squeeze_(1)
            else:
                pred = preds.squeeze(1)
                loss = self.loss(pred, mask)
                pred = pred > 0.5

        self.val_oa(pred, mask)
        self.val_prec(pred, mask)
        self.val_recall(pred, mask)
        self.val_f1(pred, mask)
        self.val_iou(pred, mask)

        self.log('val_loss', loss, on_step=True,on_epoch=True,prog_bar=True)
        return loss

    def on_validation_epoch_end(self):
        metrics = [self.val_prec.compute(),
                   self.val_recall.compute(),
                   self.val_f1.compute(),
                   self.val_iou.compute()]

        log = {'val_oa': float(self.val_oa.compute().cpu()),
               'val_prec': np.mean([item.cpu() for item in metrics[0]]),
               'val_recall': np.mean([item.cpu() for item in metrics[1]]),
               'val_f1': np.mean([item.cpu() for item in metrics[2]]),
               'val_miou': np.mean([item.cpu() for item in metrics[3]])}
        
        self.output(metrics, log, 'val')
        
        for key, value in zip(log.keys(), log.values()):
            self.log(key, value, on_step=False,on_epoch=True,prog_bar=True)
        self.log('val_change_f1', metrics[2][1], on_step=False,on_epoch=True,prog_bar=True)

        self.val_oa.reset()
        self.val_prec.reset()
        self.val_recall.reset()
        self.val_f1.reset()
        self.val_iou.reset()

        for idx in range(0, len(self.cfg.monitor_test), 1):
            if self.cfg.argmax:
                self.log(self.cfg.monitor_test[idx], self.test(idx), on_step=False,on_epoch=True,prog_bar=True)
            else:
                t = 0.2 + 0.1 * idx
                self.log(self.cfg.monitor_test[idx], self.test(idx, t), on_step=False,on_epoch=True,prog_bar=True)

    def test(self, idx, value = None):
        for input in tqdm(self.test_loader):
            raw_predictions, mask_test = self(input[0].cuda(cfg.gpus[0]), input[1].cuda(cfg.gpus[0])), input[2].cuda(cfg.gpus[0])

            if self.cfg.net == 'SARASNet':
                mask_test = Variable(resize_label(mask_test.data.cpu().numpy(), \
                                        size=raw_predictions.data.cpu().numpy().shape[2:]).to('cuda')).long()
                param = 1  # This parameter is balance precision and recall to get higher F1-score
                raw_predictions[:,1,:,:] = raw_predictions[:,1,:,:] + param 

            if self.cfg.argmax:
                pred_test = raw_predictions.argmax(dim=1)
            else:
                if self.cfg.net == 'maskcd':
                    raw_prediction = raw_predictions[0]
                    pred_test = raw_prediction > value
                    pred_test.squeeze_(1)
                else:
                    pred_test = raw_predictions.squeeze(1)
                    pred_test = pred_test > 0.5

            self.test_oa(pred_test, mask_test)
            self.test_iou(pred_test, mask_test)
            self.test_prec(pred_test, mask_test)
            self.test_f1(pred_test, mask_test)
            self.test_recall(pred_test, mask_test)

        metrics_test = [self.test_prec.compute(),
                   self.test_recall.compute(),
                   self.test_f1.compute(),
                   self.test_iou.compute()]
        
        log = {'test_oa': float(self.test_oa.compute().cpu()),
            'test_prec': np.mean([item.cpu() for item in metrics_test[0]]),
            'test_recall': np.mean([item.cpu() for item in metrics_test[1]]),
            'test_f1': np.mean([item.cpu() for item in metrics_test[2]]),
            'test_miou': np.mean([item.cpu() for item in metrics_test[3]])}
        
        self.output(metrics_test, log, 'test', idx, value)
        
        self.test_oa.reset()
        self.test_prec.reset()
        self.test_recall.reset()
        self.test_f1.reset()
        self.test_iou.reset()

        return metrics_test[2][1]

if __name__ == "__main__":
    args = get_args()
    cfg = Config.fromfile(args.config)
    logger = TensorBoardLogger(save_dir = "work_dirs",
                               sub_dir = 'log',
                               name = cfg.exp_name,
                               default_hp_metric = False)
    
    log_dir = os.path.dirname(logger.log_dir)

    model = myTrain(cfg, log_dir)
    # —— 在这里插入“推理 FPS 测试”功能 —— #
    device = torch.device(f'cuda:{cfg.gpus[0]}' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model.eval()

    # 从验证集 dataloader 里取一个 batch
    val_loader = model.val_dataloader()
    batch_iter = iter(val_loader)
    try:
        batch = next(batch_iter)
        imgA_batch = batch[0]
        imgB_batch = batch[1]
    except StopIteration:
        raise RuntimeError("验证集 dataloader 为空，请检查数据集配置。")

    # 将输入搬到同一个设备
    imgA_batch = imgA_batch.to(device)
    imgB_batch = imgB_batch.to(device)

    # 热身推理 10 次
    with torch.no_grad():
        for _ in range(10):
            _ = model(imgA_batch, imgB_batch)

    # 正式计时 N 次推理
    N = 100
    torch.cuda.synchronize(device)
    start_time = time.time()
    with torch.no_grad():
        for _ in range(N):
            _ = model(imgA_batch, imgB_batch)
    torch.cuda.synchronize(device)
    elapsed = time.time() - start_time
    fps = N / elapsed
    print(f"[推理 FPS 测试] 输入分辨率 = {imgA_batch.shape[2]}×{imgA_batch.shape[3]}，"
          f"Batch Size = {imgA_batch.shape[0]}，推理 {N} 次总耗时：{elapsed:.4f} 秒，FPS = {fps:.2f}")
    # —— 插入结束 —— #
    
    pbar = TQDMProgressBar(refresh_rate=1)
    lr_monitor=LearningRateMonitor(logging_interval = cfg.logging_interval)
    callbacks = [pbar, lr_monitor]

    ckpt_cb = ModelCheckpoint(dirpath = f'{log_dir}/ckpts/val',
                            filename = '{' + cfg.monitor_val + ':.4f}' + '-{epoch:d}',
                            monitor = cfg.monitor_val,
                            mode = 'max',
                            save_top_k = cfg.save_top_k,
                            save_last=True)
    callbacks.append(ckpt_cb)

    for m_test in cfg.monitor_test:
        ckpt_cb = ModelCheckpoint(dirpath = f'{log_dir}/ckpts/test/{m_test}',
                                filename = '{' + m_test + ':.4f}' + '-{epoch:d}',
                                monitor = m_test,
                                mode = 'max',
                                save_top_k = cfg.save_top_k,
                                save_last=True)
        callbacks.append(ckpt_cb)
    
    trainer = Trainer(max_epochs = cfg.epoch,
                    #   precision='16-mixed',
                      callbacks = callbacks,
                      logger = logger,
                      enable_model_summary = True,
                      accelerator = 'auto',
                      devices = cfg.gpus,
                      num_sanity_val_steps = 2,
                      benchmark = True)
    
    trainer.fit(model, ckpt_path=cfg.resume_ckpt_path)