Shapley-NAS: Discovering Operation Contribution for NeuralArchitecture Search代码详解

import os
#导入os模块,用于路径处理
import sys
#导入sys模块,用于系统相关操作
import time
#导入time模块,用于时间相关操作
import glob
#导入glob模块,用于通配符管理文件
import numpy as np
#导入numpy模块,用于科学计算
import torch
#导入pytorch框架
import utils
#导入utils文件
import logging
#导入logging模块,用于日志记录
import argparse
#导入argparse模块,用于命令行参数解析
import torch.nn as nn
#导入torch.nn模块,用于神经网络构建
import torch.utils
#导入torch.utils模块
import torchvision.datasets as dset
#导入torchvision.datasets模块,用于图像数据集加载
import torch.backends.cudnn as cudnn
#导入torch.backends.cudnn模块,用于GPU加速

from torch.autograd import Variable
#从torch.autograd导入Variable,用于神经网络训练
from model_search import Network
#从model_search文件导入Network类,用于模型定义

import time
#再次导入time模块

parser = argparse.ArgumentParser("cifar")
#初始化argparse解析器,名称为'cifar'
parser.add_argument('--data', type=str, default='data', help='location of the data corpus')
#添加命令行参数'--data',类型为字符串,默认值为'data',用于提示数据文件位置
parser.add_argument('--set', type=str, default='cifar10', help='location of the data corpus')
#添加命令行参数'--set',类型为字符串,默认值为'cifar10',用于选择数据集
parser.add_argument('--batch_size', type=int, default=256, help='batch size')
#添加命令行参数'--batch_size',类型为整数,默认值为256,用于设置batchsize
parser.add_argument('--learning_rate', type=float, default=0.1, help='init learning rate')
#添加命令行参数'--learning_rate',类型为浮点数,默认值为0.1,用于设置初始学习率
parser.add_argument('--learning_rate_min', type=float, default=0.0, help='min learning rate')
#添加命令行参数'--learning_rate_min',类型为浮点数,默认值为0.0,用于设置最小学习率
parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
#添加命令行参数'--momentum',类型为浮点数,默认值为0.9,用于设置动量
parser.add_argument('--weight_decay', type=float, default=3e-4, help='weight decay')
#添加命令行参数'--weight_decay',类型为浮点数,默认值为3e-4,用于设置权重衰减
parser.add_argument('--report_freq', type=float, default=50, help='report frequency')
#添加命令行参数'--report_freq',类型为浮点数,默认值为50,用于设置报告频率
parser.add_argument('--gpu', type=int, default=0, help='gpu device id')
#添加命令行参数'--gpu',类型为整数,默认值为0,用于设置GPU设备ID
parser.add_argument('--epochs', type=int, default=50, help='num of training epochs')
#添加命令行参数'--epochs',类型为整数,默认值为50,用于设置训练的epoch数
parser.add_argument('--init_channels', type=int, default=16, help='num of init channels')
#添加命令行参数'--init_channels',类型为整数,默认值为16,用于设置初始通道数
parser.add_argument('--layers', type=int, default=8, help='total number of layers')
#添加命令行参数'--layers',类型为整数,默认值为8,用于设置总层数
parser.add_argument('--model_path', type=str, default='saved_models', help='path to save the model')
#添加命令行参数'--model_path',类型为字符串,默认值为'saved_models',用于设置模型保存路径
parser.add_argument('--cutout', action='store_true', default=False, help='use cutout')
#添加命令行参数'--cutout',类型为布尔值,默认值为False,用于是否使用cutout
parser.add_argument('--cutout_length', type=int, default=16, help='cutout length')
#添加命令行参数'--cutout_length',类型为整数,默认值为16,用于设置cutout长度
parser.add_argument('--drop_path_prob', type=float, default=0.3, help='drop path probability')
#添加命令行参数'--drop_path_prob',类型为浮点数,默认值为0.3,用于设置drop path概率
parser.add_argument('--save', type=str, default='EXP', help='experiment name')
#添加命令行参数'--save',类型为字符串,默认值为'EXP',用于设置实验名称
parser.add_argument('--seed', type=int, default=2, help='random seed')
#添加命令行参数'--seed',类型为整数,默认值为2,用于设置随机种子
parser.add_argument('--grad_clip', type=float, default=5, help='gradient clipping')
#添加命令行参数'--grad_clip',类型为浮点数,默认值为5,用于设置梯度裁剪
parser.add_argument('--celltype', type=str, default='normal', help='experiment name')
#添加命令行参数'--celltype',类型为字符串,默认值为'normal',用于设置细胞类型
parser.add_argument('--resume', type=str, default='', help='resume from pretrained')
#添加命令行参数'--resume',类型为字符串,默认值'',用于设置模型加载路径
parser.add_argument('--train_portion', type=float, default=0.5, help='portion of training data')
#添加命令行参数'--train_portion',类型为浮点数,默认值为0.5,用于设置训练数据比例
parser.add_argument('--unrolled', action='store_true', default=False, help='use one-step unrolled validation loss')
#添加命令行参数'--unrolled',类型为布尔值,默认值为False,用于是否使用unroll验证损失
parser.add_argument('--arch_learning_rate', type=float, default=6e-4, help='learning rate for arch encoding')
#添加命令行参数'--arch_learning_rate',类型为浮点数,默认值为6e-4,用于设置体系结构编码学习率
parser.add_argument('--arch_weight_decay', type=float, default=1e-3, help='weight decay for arch encoding')
#添加命令行参数'--arch_weight_decay',类型为浮点数,默认值为1e-3,用于设置体系结构编码权重衰减
parser.add_argument('--shapley_momentum', type=float, default=0.8, help='momentum for updating shapley')
#添加命令行参数'--shapley_momentum',类型为浮点数,默认值为0.8,用于设置更新Shapley值的动量
parser.add_argument('--step_size', type=float, default=0.1, help='step size for updating shapley')
#添加命令行参数'--step_size',类型为浮点数,默认值为0.1,用于设置更新Shapley值的步长
parser.add_argument('--samples', type=int, default=10, help='number of samples for estimation')
#添加命令行参数'--samples',类型为整数,默认值为10,用于设置估计的样本数
parser.add_argument('--threshold', type=float, default=0.5, help='early truncation threshold')
#添加命令行参数'--threshold',类型为浮点数,默认值为0.5,用于设置提前截断阈值

args = parser.parse_args()
#解析添加的命令行参数

args.save = 'search-{}-{}'.format(args.save, time.strftime("%Y%m%d-%H%M%S"))
#设置保存目录名称,包含实验名称和时间戳
utils.create_exp_dir(args.save, scripts_to_save=glob.glob('*.py'))
#创建实验目录,并保存相关.py文件

log_format = '%(asctime)s %(message)s'
#设置日志格式
logging.basicConfig(stream=sys.stdout, level=logging.INFO, #设置日志级别为INFO,打印到标准输出
format=log_format, datefmt='%m/%d %I:%M:%S %p')
fh = logging.FileHandler(os.path.join(args.save, 'log.txt'))

logging.getLogger().addHandler(fh) CIFAR_CLASSES = 10 if args.set=='cifar100': CIFAR_CLASSES = 100 def main(): if not torch.cuda.is_available(): logging.info('no gpu device available') sys.exit(1) np.random.seed(args.seed) torch.cuda.set_device(args.gpu) cudnn.benchmark = True torch.manual_seed(args.seed) cudnn.enabled=True torch.cuda.manual_seed(args.seed)

logging.info('gpu device = %d' % args.gpu)
#记录GPU设备ID
logging.info("args = %s", args)
#记录命令行参数

criterion = nn.CrossEntropyLoss()
#定义损失函数为交叉熵损失
criterion = criterion.cuda()
#将损失函数移至GPU
model = Network(args.init_channels, CIFAR_CLASSES, args.layers, criterion)
#定义模型,输入参数包括初始通道数、类数、总层数和损失函数
model = model.cuda()
#将模型移至GPU
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
#记录模型参数量

if args.resume: #如果指定了加载路径

model.load_state_dict(torch.load(args.resume))  
#加载指定路径的模型参数
model.show_arch_parameters()
#显示模型体系结构参数  
genotype = model.genotype()
#获得模型的基因型
logging.info('genotype = %s', genotype)

arch_params = list(map(id, model.arch_parameters()))
#获得模型体系结构参数列表
weight_params = filter(lambda p: id(p) not in arch_params,
model.parameters())
#获得模型权重参数列表
optimizer = torch.optim.SGD(
weight_params, #仅对权重参数进行优化
args.learning_rate, #学习率
momentum=args.momentum, #动量
weight_decay=args.weight_decay) #权重衰减

#定义数据转换
train_transform, valid_transform = utils._data_transforms_cifar10(args)
if args.set=='cifar100': #如果选择CIFAR100数据集
train_data = dset.CIFAR100(root=args.data, train=True, download=True, transform=train_transform)
else: #否则选择CIFAR10数据集
train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform)

num_train = len(train_data) #获得训练集大小
indices = list(range(num_train)) #定义索引列表
split = int(np.floor(args.train_portion * num_train)) #根据比例获得分割索引

train_queue = torch.utils.data.DataLoader( #定义训练集data loader
train_data, batch_size=args.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]), #使用SubsetRandomSampler随机采样一部分数据
pin_memory=True)

valid_queue = torch.utils.data.DataLoader( #定义验证集data loader
train_data, batch_size=1024,
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[split:num_train]),
pin_memory=True,num_workers=8)

infer_queue = torch.utils.data.DataLoader( # 定义推理集data loader
train_data, batch_size=args.batch_size//2,
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[split:num_train]),
pin_memory=True)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( #定义学习率scheduler
optimizer, float(args.epochs), eta_min=args.learning_rate_min)

ops = []
#定义操作列表
for cell_type in ['normal','reduce']: #对normal cell和reduce cell
for edge in range(model.num_edges): #在每条边上
ops.append(['{}{}{}'.format(cell_type, edge, i) for i in #添加操作
range(0, model.num_ops)])
ops = np.concatenate(ops) #连接列表