train_FLIR_Align_Dual_FCOS.py

from model.Align_Dual_FCOS_PVT import Align_Dual_FCOSDetector
import torch
from dataset.FLIR_Aligned_dataset import FLIRDataset, FLIRGenerator
import math,time
from dataset.augment import Transforms
import os
import numpy as np
import random
import torch.backends.cudnn as cudnn
import argparse
import coco_eval
from torch.cuda.amp import autocast, GradScaler

torch.manual_seed(0)
torch.cuda.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
cudnn.benchmark = False
cudnn.deterministic = True

random.seed(0)

def seed_worker(worker_id):
    np.random.seed(0)
    random.seed(0)

def main(opt):
    os.environ["CUDA_VISIBLE_DEVICES"]=opt.n_gpu
    init_distributed_mode(opt)
    print(opt)

    transform = Transforms()
    train_dataset=FLIRDataset(
                            transform=transform  
                            )

    test_dataset=FLIRGenerator()
    # 初始化
    scaler = GradScaler()  # 用于梯度缩放，防止下溢
    model = Align_Dual_FCOSDetector(num_class = opt.num_class)



    if torch.cuda.is_available():
        model.cuda()
        model_without_ddp = model
        if opt.distributed:
            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[opt.gpu], find_unused_parameters=True)
            model_without_ddp = model.module
    else:
        raise ValueError('CUDA is not available!')

    #print(model)
   
    BATCH_SIZE=opt.batch_size

    EPOCHS=opt.epochs


    if opt.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
        train_loader = torch.utils.data.DataLoader(
            train_dataset, batch_size=opt.batch_size, sampler=train_sampler, num_workers=opt.n_cpu, drop_last=True, worker_init_fn=seed_worker, generator=torch.Generator().manual_seed(0),
            collate_fn=train_dataset.collate_fn) 

    else:
        train_loader=torch.utils.data.DataLoader(
            train_dataset, batch_size=opt.batch_size, num_workers=opt.n_cpu, shuffle=True, drop_last=True, worker_init_fn=seed_worker, generator=torch.Generator().manual_seed(0),
            collate_fn=train_dataset.collate_fn)

            # 用来保存coco_info的文件
    results_file = "{}_{}.txt".format(opt.dataset, opt.model_mode)


    WARMUP_STEPS=500
    WARMUP_FACTOR = 1.0 / 3.0
    GLOBAL_STEPS=0
    LR_INIT=1e-4

    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
    

    def lr_func(epoch, step):
        lr = LR_INIT

        if step < WARMUP_STEPS:
            alpha = float(step) / WARMUP_STEPS
            warmup_factor = WARMUP_FACTOR * (1.0 - alpha) + alpha
            lr = lr*warmup_factor

        elif epoch >= opt.lr_drop:
            lr = lr*0.1

        return float(lr)

    import collections
    loss_hist = collections.deque(maxlen=5000)
    mAP = 0

    total_time = 0
    
    for epoch in range(EPOCHS):
        if opt.distributed:
            train_sampler.set_epoch(epoch)
    
        model.train()
        epoch_start = time.time()  # 单个epoch计时                                                                                                                                                                                                                                                                                                                                                                                                                 
        # if opt.distributed: 
        #     model.module.fcos_body.freeze_bn()
        # else:
        #     model.fcos_body.freeze_bn()
        
        for epoch_step, data in enumerate(train_loader):

            batch_imgs_rgb = data['img_rgb']
            batch_imgs_ir = data['img_ir']
            batch_boxes = data['boxes']
            batch_classes = data['classes']

            batch_imgs_rgb=batch_imgs_rgb.cuda()
            batch_imgs_ir=batch_imgs_ir.cuda()
            batch_boxes=batch_boxes.cuda()
            batch_classes=batch_classes.cuda()

            lr = lr_func(epoch, GLOBAL_STEPS)
            for param in optimizer.param_groups:
                param['lr']=lr
            
            start_time=time.time()

            optimizer.zero_grad()
            # 前向传播（自动混合精度）
            #with autocast():  # 自动选择FP16/FP32
            losses=model(batch_imgs_rgb, batch_imgs_ir, batch_boxes, batch_classes)
            loss=losses[-1].mean()
            loss.backward()
            # 反向传播（自动梯度缩放）
            # scaler.scale(loss).backward()  # 1. 缩放梯度
            # scaler.step(optimizer)         # 2. 更新参数（自动反向缩放）
            # scaler.update()                # 3. 调整缩放系数
            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)

            optimizer.step()

            loss_hist.append(float(loss))

            end_time=time.time()
            cost_time=int((end_time-start_time)*1000)
            if epoch_step % 50 == 0:
                print('Epoch: {:0>3} | Iteration: {:0>3} | lr: {:1.6f} | cls_loss: {:1.5f} | cnt_loss: {:1.5f} | reg_loss: {:1.5f} | cost_time:{:0>4}ms | Total_loss: {:1.5f}'.format(
                    epoch, epoch_step, lr, losses[0].mean(),losses[1].mean(),losses[2].mean(), cost_time, np.mean(loss_hist)
                ))
                #print("global_steps:%d epoch:%d steps:%d cls_loss:%.4f cnt_loss:%.4f reg_loss:%.4f cost_time:%dms lr=%.4e total_loss:%.4f"%\
                #    (GLOBAL_STEPS,epoch,epoch_step,losses[0].mean(),losses[1].mean(),losses[2].mean(),cost_time,lr, loss.mean()))


            GLOBAL_STEPS+=1

        epoch_time = time.time() - epoch_start
        total_time = total_time + epoch_time
        # 测试时仅使用主进程
        if not opt.distributed or opt.rank in [-1, 0]:
            if epoch >= 10:
                print('Evaluating dataset')
                
                coco_info = coco_eval.evaluate_coco_multispectral(test_dataset, model.eval())
                
                #scheduler.step() 
                # write into txt
                with open(results_file, "a") as f:
                    # 写入的数据包括coco指标还有loss和learning rate
                    result_info = [str(round(i, 5)) for i in coco_info + [np.mean(loss_hist), optimizer.param_groups[0]["lr"]]]
                    txt = "epoch:{} {}".format(epoch, '  '.join(result_info))
                    f.write(txt + "\n")

                if coco_info[1] > mAP:
                    mAP = coco_info[1]
                    torch.save(model_without_ddp.state_dict(), os.path.join(opt.save_path, '{}_FCOS_{}.pth'.format(opt.dataset, mAP)))

                print("当前模型最高检测精度为: " + str(mAP))
                torch.save(model_without_ddp.state_dict(), os.path.join(opt.save_path, '{}_last.pth'.format(opt.dataset, mAP)))
        
    # 原始时间（秒）→ 转换为小时制
    hours = int(total_time // 3600)
    minutes = int((total_time % 3600) // 60)
    seconds = int(total_time % 60)
    print(f'\nTotal training time: {hours}h {minutes}m {seconds}s')






import torch.distributed as dist
def is_dist_avail_and_initialized():
    """检查是否支持分布式环境"""
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True

def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()

def reduce_dict(input_dict, average=True):
    """
    Args:
        input_dict (dict): all the values will be reduced
        average (bool): whether to do average or sum
    Reduce the values in the dictionary from all processes so that all processes
    have the averaged results. Returns a dict with the same fields as
    input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:  # 单GPU的情况
        return input_dict
    with torch.no_grad():  # 多GPU的情况
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.all_reduce(values)
        if average:
            values /= world_size

        reduced_dict = {k: v for k, v in zip(names, values)}
        return reduced_dict


def init_distributed_mode(args):
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.gpu = int(os.environ['LOCAL_RANK'])
    elif 'SLURM_PROCID' in os.environ:
        args.rank = int(os.environ['SLURM_PROCID'])
        args.gpu = args.rank % torch.cuda.device_count()
    else:
        print('Not using distributed mode')
        args.distributed = False
        return

    args.distributed = True

    torch.cuda.set_device(args.gpu)
    args.dist_backend = 'nccl'
    print('| distributed init (rank {}): {}'.format(
        args.rank, args.dist_url), flush=True)
    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                         world_size=args.world_size, rank=args.rank)
    torch.distributed.barrier()
    setup_for_distributed(args.rank == 0)

def setup_for_distributed(is_master):
    """
    This function disables when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, default='FLIR', help="dataset name")
    parser.add_argument("--num_class", type=int, default=3, help="dataset name")
    parser.add_argument('--model_mode', default='baseline_dual_align_head_align_PVT', type=str, help='model mode, such as baseline+head_align')
    parser.add_argument("--save_path", type=str, default='/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/FLIR/baseline_dual_align_head_align_PVT/', help='weights save path')
    parser.add_argument("--epochs", type=int, default=40, help="number of epochs")
    parser.add_argument('--lr_drop', help='epoch number for learning rate reduction', default=30, type=int)
    parser.add_argument("--batch_size", type=int, default=8, help="size of each image batch")
    parser.add_argument("--n_cpu", type=int, default=4, help="number of cpu threads to use during batch generation") 
    parser.add_argument("--n_gpu", type=str, default='3', help="number of cpu threads to use during batch generation")
    
    # 开启的进程数(注意不是线程)
    parser.add_argument('--world-size', default=4, type=int,
                        help='number of distributed processes')
    parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
    # 训练设备类型
    parser.add_argument('--device', default='cuda', help='device')
    opt = parser.parse_args()

    main(opt)