修改后的代码如下：

import os import argparse import json import torch import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist from torch.cuda.amp import autocast, GradScaler from torch.nn.parallel import DataParallel, DistributedDataParallel from torch.utils.data import DataLoader, RandomSampler, DistributedSampler 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('--num_workers', default=16, type=int, help='number of workers')#工作进程数 parser.add_argument('--use_mp', action='store_true', default=False, help='use Mixed Precision')# 是否使用混合精度训练（混合精度训练就是f为32位或16位） parser.add_argument('--use_ddp', action='store_true', default=False, help='use Distributed Data Parallel')# 是否使用分布式训练 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 environment训练环境

if args.use_ddp: torch.distributed.init_process_group(backend='nccl', init_method='env://')# 初始化分布式训练 world_size=16# 分布式训练时设备总数 local_rank = dist.get_rank()# 获取当前进程的rank torch.cuda.set_device(local_rank)# 设置当前进程使用的GPU设备 if local_rank == 0: print('==> Using DDP.')# 仅进程0输出信息 else: world_size=16 # 设备总数

training config加载配置文件

with open(os.path.join('configs', args.exp, 'base.json'), 'r') as f: b_setup = json.load(f)# 读取基本配置

variant = 'UNet' config_name = 'model_'+variant+'.json'

模型配置文件名称

with open(os.path.joi