修改后的代码如下：

import os import argparse import json import torch import torch.nn as nn import torch.nn.functional as F from torch.cuda.amp import autocast, GradScaler from torch.nn.parallel import DataParallel from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm

from utils import AverageMeter, CosineScheduler, pad_img from datasets import PairLoader from models import *

定义命令行参数

parser = argparse.ArgumentParser() parser.add_argument('--model', default='gunet_t', type=str, help='model name') parser.add_argument('--num_workers', default=16, type=int, help='number of workers') parser.add_argument('--use_mp', action='store_true', default=False, help='use Mixed Precision') parser.add_argument('--save_dir', default='./saved_models/', type=str, help='path to models saving') parser.add_argument('--data_dir', default='./data/', type=str, help='path to dataset') parser.add_argument('--log_dir', default='./logs/', type=str, help='path to logs') parser.add_argument('--train_set', default='SOTS-OUT', type=str, help='train dataset name') parser.add_argument('--val_set', default='SOTS-IN/SOTS-IN', type=str, help='valid dataset name') parser.add_argument('--exp', default='reside-in', type=str, help='experiment setting') args = parser.parse_args()

training config

with open(os.path.join('configs', args.exp, 'base.json'), 'r') as f: b_setup = json.load(f)

variant = args.model.split('')[-1] config_name = 'model'+variant+'.json' if variant in ['t', 's', 'b', 'd'] else 'default.json' with open(os.path.join('configs', args.exp, config_name), 'r') as f: m_setup = json.load(f) print(m_setup)

def train(train_loader, network, criterion, optimizer, scaler, frozen_bn=False): losses = AverageMeter()

torch.cuda.empty_cache()

network.eval() if frozen_bn else network.train()

for batch in train_loader:
    source_img = batch['source'].cuda()
    target_img = batch['target'].cuda()

    with autocast(args.use_mp):
        output = network(source_img)
        loss = criterion(output, target_img)

    optimizer.zero_grad()
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    losses.update(loss.item())

return losses.avg

def valid(val_loader, network): PSNR = AverageMeter()

torch.cuda.empty_cache()

network.eval()

for batch in val_loader:
    source_img = batch['source'].cuda()
    target_img = batch['target'].cuda()

    with torch.no_grad():
        H, W = source_img.shape[2:]
        source_img = pad_img(source_img, network.module.patch_size if hasattr(network.module, 'patch_size') else 16)
        output = network(source_img).clamp_(-1,