main_fixed.py

import os
import sys
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# 先读 config 再设可见 GPU
from config import Config
opt = Config('training.yml')
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(i) for i in opt.GPU])


# 多卡时自动用 DDP 启动，避免 DataParallel 导致 4 卡比 2 卡慢
if "RANK" not in os.environ and "LOCAL_RANK" not in os.environ and len(opt.GPU) > 1:
    import subprocess
    env = os.environ.copy()
    cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node", str(len(opt.GPU)), sys.argv[0]] + sys.argv[1:]
    sys.exit(subprocess.run(cmd, env=env).returncode)

import torch
torch.backends.cudnn.benchmark = True
import utils as utils
from models.encoder2 import Convres
from restormer import ChannelShuffleWithGBPDeep
from torchvision.transforms import transforms
from PIL import Image
from skimage.metrics import peak_signal_noise_ratio
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import random
import time
import numpy as np
from model.common import VGGLoss
from data_RGB import get_training_data, get_validation_data
try:
    from warmup_scheduler import GradualWarmupScheduler
except ImportError:
    GradualWarmupScheduler = None
from tqdm import tqdm

# DDP：用 torchrun 启动时 4 卡会真正均摊负载，比 DataParallel 快且随卡数缩放
use_ddp = "RANK" in os.environ or "LOCAL_RANK" in os.environ
if use_ddp:
    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    world_size = torch.distributed.get_world_size()
    rank = torch.distributed.get_rank()
else:
    local_rank = 0
    world_size = 1
    rank = 0

# The testing dataset files
img_path = './dataset/test/target_smoke'
targeet_path = './dataset/test/target'
img_list = sorted(os.listdir(img_path))
num_img = len(img_list)
gpus = ','.join([str(i) for i in opt.GPU])
######### Set Seeds ###########
random.seed(1234)
np.random.seed(1234)
torch.manual_seed(1234)
torch.cuda.manual_seed_all(1234)
contrast_loss = torch.nn.CrossEntropyLoss().cuda()
# device = torch.device("cuda:0,1")
start_epoch = 1
mode = opt.MODEL.MODE
session = opt.MODEL.SESSION

result_dir = os.path.join(opt.TRAINING.SAVE_DIR, mode, 'results', session)
model_dir = os.path.join(opt.TRAINING.SAVE_DIR, mode, 'models', session)

utils.mkdir(result_dir)
utils.mkdir(model_dir)

train_dir = opt.TRAINING.TRAIN_DIR
val_dir = opt.TRAINING.VAL_DIR

######### Model ###########
model_G1 = ChannelShuffleWithGBPDeep()
if use_ddp:
    model_G1 = model_G1.cuda(local_rank)
    model_G1 = nn.parallel.DistributedDataParallel(
        model_G1, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True
    )
    if rank == 0:
        print("\n==> DDP: world_size={}, 每卡 batch={}（总 batch={}）\n".format(
            world_size, opt.OPTIM.BATCH_SIZE, opt.OPTIM.BATCH_SIZE * world_size))
else:
    model_G1 = model_G1.cuda()
    device_ids = list(range(torch.cuda.device_count()))
    if len(device_ids) > 1:
        model_G1 = nn.DataParallel(model_G1, device_ids=device_ids)
        print("\n" + "!" * 60)
        print("  当前是 DataParallel，{} 卡时通常会比 2 卡更慢。".format(len(device_ids)))
        print("  要让 4 卡比 2 卡快，请用 DDP 启动:  bash run_ddp.sh")
        print("  或:  CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 main.py")
        print("!" * 60 + "\n")

new_lr = opt.OPTIM.LR_INITIAL
optimizer_G1 = optim.Adam(model_G1.parameters(), lr=new_lr, betas=(0.9, 0.999), eps=1e-8)
# optimizer_G2 = optim.Adam(model_G2.parameters(), lr=new_lr, betas=(0.9, 0.999), eps=1e-8)
######### Scheduler ###########
# warmup_epochs = 90
# scheduler_cosineG1 = optim.lr_scheduler.CosineAnnealingLR(optimizer_G1, opt.OPTIM.NUM_EPOCHS - warmup_epochs,
#                                                           eta_min=opt.OPTIM.LR_MIN)
# scheduler_cosineG2 = optim.lr_scheduler.CosineAnnealingLR(optimizer_G2, opt.OPTIM.NUM_EPOCHS - warmup_epochs,
#                                                           eta_min=opt.OPTIM.LR_MIN)
# scheduler_G1 = GradualWarmupScheduler(optimizer_G1, multiplier=1, total_epoch=warmup_epochs,
#                                       after_scheduler=scheduler_cosineG1)
# scheduler_G2 = GradualWarmupScheduler(optimizer_G2, multiplier=1, total_epoch=warmup_epochs,
#                                       after_scheduler=scheduler_cosineG2)
# scheduler_G1.step()
# scheduler_G2.step()
######### Resume ###########
if opt.TRAINING.RESUME:
    path_chk_rest = ''
    utils.load_checkpointG1(model_G1, path_chk_rest, strict=False)  # strict=False 兼容旧 ckpt（无 prior_encoder 等）
    # utils.load_checkpointG2(model_G2, path_chk_rest)
    # start_epoch = utils.load_start_epoch(path_chk_rest) + 1
    start_epoch = 1
    print("start_epoch=",start_epoch)
    utils.load_optimG1(optimizer_G1, path_chk_rest)
    # utils.load_optimG2(optimizer_G2, path_chk_rest)
    # for i in range(1, start_epoch):
    #     scheduler_G1.step()
    #     scheduler_G2.step()
    # new_lr = scheduler_G2.get_lr()[0]
    # print('------------------------------------------------------------------------------')
    # print("==> Resuming Training with learning rate:", new_lr)
    # print('------------------------------------------------------------------------------')

######### Loss ###########
# criterion_char = Deraining.losses.CharbonnierLoss()
# criterion_edge = Deraining.losses.EdgeLoss()
ide_loss = torch.nn.L1Loss().cuda()
# satu = ContrastLoss().cuda()
######### DataLoaders ###########
train_dataset = get_training_data(train_dir, {'patch_size': opt.TRAINING.TRAIN_PS})
if use_ddp:
    # 4 进程 × 16 workers = 64 个 worker 容易抢资源，每进程少一点
    n_workers = max(8, 16 // world_size)
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
    train_loader = DataLoader(
        train_dataset, batch_size=opt.OPTIM.BATCH_SIZE, shuffle=False, sampler=train_sampler,
        num_workers=n_workers, drop_last=False, pin_memory=True,
        persistent_workers=True if n_workers > 0 else False, prefetch_factor=8 if n_workers > 0 else None)
else:
    train_loader = DataLoader(
        dataset=train_dataset, batch_size=opt.OPTIM.BATCH_SIZE, shuffle=True, num_workers=16,
        drop_last=False, pin_memory=True)

val_dataset = get_validation_data(val_dir, {'patch_size': opt.TRAINING.VAL_PS})
val_loader = DataLoader(dataset=val_dataset, batch_size=16, shuffle=False, num_workers=8, drop_last=False, pin_memory=True)

if rank == 0:
    print('===> Start Epoch {} End Epoch {}'.format(start_epoch, opt.OPTIM.NUM_EPOCHS + 1))
    print('===> Loading datasets')

best_psnr = 0
best_epoch = 0
transform = transforms.ToTensor()
# print("The lr is:",scheduler_G1.get_lr()[0])
for epoch in range(start_epoch, opt.OPTIM.NUM_EPOCHS + 1):
    if use_ddp:
        train_sampler.set_epoch(epoch)
    epoch_start_time = time.time()
    epoch_loss = 0
    model_G1.train()

    for i, data in enumerate(tqdm(train_loader, disable=(rank != 0)), 0):
        optimizer_G1.zero_grad(set_to_none=True)
        target = data[0].cuda(local_rank, non_blocking=True)
        input_ = data[1].cuda(local_rank, non_blocking=True)

        # white = data[2].cuda()
        # gray = data[3].cuda()
        # input_ze = data[4].cuda()
        # white1 = data[5].cuda()
        # gray1 = data[6].cuda()

        ## The first stage: the model_G1 is training and the loss_con is the dual degradation loss in the paper.

        # mu1_1, dr1_1 = model_G1(input_ze, gray1, white1, flag='low')
        # mu2,dr2 = model_G1(target,gray,white, flag = 'clean')
        # 8, 384, 32, 32
        # da = ide_loss(mu1,mu1_1)
        # db = ide_loss(mu1,mu2)
        # loss_con = da/(db + 1e-7)

        ## The second stage: the model_G1 is fixed and the grad not backward.
        # with torch.set_grad_enabled(False):
        #     mu1, dr1 = model_G1(input_, gray, white, flag='low')
        output = model_G1(input_)
        # vgg = loss_vgg(output, target)
        ide = ide_loss(output, target)

        ## In the first stage, you should add the loss_con, e.g., DDLoss
        loss = ide

        loss.backward()
        optimizer_G1.step()
        epoch_loss += loss.item()

    if rank == 0:
        print("------------------------------------------------------------------")
        print("Epoch: {}\tTime: {:.4f}\tLoss: {:.4f}".format(epoch, time.time() - epoch_start_time, epoch_loss))
        print("------------------------------------------------------------------")
    if epoch % 1 == 0 and rank == 0:
        state_to_save = model_G1.module.state_dict() if use_ddp else model_G1.state_dict()
        torch.save({'epoch': epoch,
                    'state_dict_G1': state_to_save,
                    # 'state_dict_G2': model_G2.state_dict(),
                    'optimizer_G1': optimizer_G1.state_dict(),
                    # 'optimizer_G2': optimizer_G2.state_dict()
                    }, os.path.join(model_dir,  'model_{}.pth'.format(epoch)))
        print("laileao")
        model_G1.eval()
        # model_G2.eval()
        transform = transforms.ToTensor()
        PSNR = 0
        # testing stage
        for img in img_list:
            image = Image.open(img_path + '/' + img).convert('RGB')
            target = Image.open(targeet_path + '/' + img).convert('RGB')
            image = transform(image).unsqueeze(0).cuda(local_rank)
            target = transform(target).unsqueeze(0).cuda(local_rank)
            # r, g, b = image[0] + 1, image[1] + 1, image[2] + 1
            # lr_gray = 1. - (0.299 * r + 0.587 * g + 0.114 * b) / 2.
            # lr_gray = torch.unsqueeze(lr_gray, 0)
            # lr_white = 1 - lr_gray
            # [A, B, C] = image.shape
            # image = image.reshape([1, A, B, C])
            # [A, B, C] = target.shape
            # target = target.reshape([1, A, B, C])
            # [A, B, C] = lr_gray.shape
            # lr_gray = lr_gray.reshape([1, A, B, C])
            # [A, B, C] = lr_white.shape
            # lr_white = lr_white.reshape([1, A, B, C])
            # gray_test = torch.cat([image, lr_gray], dim=1)
            # white_test = torch.cat([image, lr_white], dim=1)
            with torch.set_grad_enabled(False):
                pre = model_G1(image)
                # pre = model_G2(image, dr1)
            p_numpy = pre.squeeze(0).cpu().detach().numpy()
            label_numpy = target.squeeze(0).cpu().detach().numpy()
            psnr = peak_signal_noise_ratio(label_numpy, p_numpy, data_range=1)
            PSNR += psnr
        PSNR = PSNR / num_img
        print("PSNR =", PSNR)
    if use_ddp:
        torch.distributed.barrier()