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()