这段代码是一个用于配置 Swin Transformer 模型训练和推理的 Python 脚本，主要功能是定义默认配置、从 YAML 文件更新配置以及根据命令行参数进一步调整配置。 ### 代码功能模块分析 #### 1. 导入必要的库 ```python import os import yaml from yacs.config import CfgNode as CN ``` 导入了 `os` 用于文件和路径操作，`yaml` 用于解析 YAML 文件，`CfgNode` 用于创建和管理配置节点。 #### 2. 创建配置节点 ```python _C = CN() _C.MODEL = CN() _C.TRAIN = CN() _C.DATA = CN() _C.CHANNEL = CN() ``` 创建了主配置节点 `_C`，并为模型、训练、数据和信道分别创建子配置节点。 #### 3. 定义默认配置 ```python # 模型配置 _C.MODEL.PATCH_SIZE = 1 _C.MODEL.IN_CHANS = 1 _C.MODEL.EMBED_DIMS = [1] # 其他模型相关配置... # 保存路径配置 _C.MODEL.ENCODER_PATH = "" _C.MODEL.DECODER_PATH = "" # 其他保存路径相关配置... # 训练环境配置 _C.TRAIN.EPOCHS = 1 _C.TRAIN.SAVE_FRE = 1 # 其他训练环境相关配置... # 数据配置 _C.DATA.DATASET = "DIV2K" _C.DATA.INFDATA = "CIFAR10" # 其他数据相关配置... # 信道配置 _C.CHANNEL.TRAIN_SNR = 1 _C.CHANNEL.EVAL_SNR = 1 # 其他信道相关配置... # 固定配置 _C.LEARNING_RATE = 1e-4 ``` 为模型、保存路径、训练环境、数据和信道分别定义了默认配置。 #### 4. 从 YAML 文件更新配置 ```python def _update_config_from_file(config, cfg_file): config.defrost() with open(cfg_file, "r") as f: yaml_cfg = yaml.load(f, Loader=yaml.FullLoader) for cfg in yaml_cfg.setdefault("BASE", [""]): if cfg: _update_config_from_file(config, os.path.join(os.path.dirname(cfg_file), cfg)) print(1) print("=> merge config from {}".format(cfg_file)) config.merge_from_file(cfg_file) # config.freeze() ``` 该函数用于从 YAML 文件中加载配置，并将其合并到现有配置中。如果 YAML 文件中包含 `BASE` 字段，则递归地加载基础配置文件。 #### 5. 根据命令行参数更新配置 ```python def update_config(config, args): _update_config_from_file(config, args.model_config_path) _update_config_from_file(config, args.train_config_path) ``` 该函数用于根据命令行参数中的模型和训练配置文件路径更新配置。 #### 6. 获取完整配置 ```python def get_config(args): config = _C.clone() model_config_path = f"./configs/models/{args.model}_{args.depths}.yaml" train_config_path = f"./configs/train/{args.dataset}.yaml" _update_config_from_file(config, model_config_path) _update_config_from_file(config, train_config_path) if args.mode == "train_inf": train_config_path = f"./configs/train/{args.interference_dataset}.yaml" _update_config_from_file(config, train_config_path) config.DATA.DATASET = args.dataset config.DATA.INFDATA = args.interference_dataset config.TRAIN.LOSS = args.loss config.TRAIN.EVAL_LOSS = args.eval_loss # 其他根据命令行参数更新的配置... return config ``` 该函数根据命令行参数获取完整的配置。首先克隆默认配置，然后从模型和训练配置文件中更新配置。如果训练模式为 `train_inf`，则加载干扰数据集的训练配置文件。最后，根据命令行参数进一步调整配置。 #### 7. 将配置节点转换为字典 ```python def convert_to_dict(cfg_node, key_list=[]): _VALID_TYPES = {tuple, list, str, int, float, bool} if not isinstance(cfg_node, CN): if type(cfg_node) not in _VALID_TYPES: print( "Key {} with value {} is not a valid type; valid types: {}".format( ".".join(key_list), type(cfg_node), _VALID_TYPES ), ) return cfg_node else: cfg_dict = dict(cfg_node) for k, v in cfg_dict.items(): cfg_dict[k] = convert_to_dict(v, key_list + [k]) return cfg_dict ``` 该函数用于将配置节点递归地转换为字典，以便于后续处理。 ###