使用argparse模块来定义和解析命令行参数
创建一个ArgumentParser对象
parser = argparse.ArgumentParser()
训练的轮数,每批图像的大小,更新模型参数之前累积梯度的次数,模型定义文件的路径。
parser.add_argument("--epochs", type=int, default=100, help="number of epochs")
parser.add_argument("--batch_size", type=int, default=1, help="size of each image batch")
parser.add_argument("--gradient_accumulations", type=int, default=2, help="number of gradient accums before step")
parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
数据配置文件的路径,从预训练的模型权重开始训练,生成批次数据时使用的CPU线程数。
parser.add_argument("--data_config", type=str, default="config/coco.data", help="path to data config file")
parser.add_argument("--pretrained_weights", type=str, help="if specified starts from checkpoint model")
parser.add_argument("--n_cpu", type=int, default=0, help="number of cpu threads to use during batch generation")
每张图像的尺寸,每隔多少个epoch保存一次模型权重,每隔多少个epoch在验证集上进行一次评估,每十批计算一次平均精度(mAP),是否允许多尺度训练,
parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
parser.add_argument("--checkpoint_interval", type=int, default=20, help="interval between saving model weights")
parser.add_argument("--evaluation_interval", type=int, default=20, help="interval evaluations on validation set")
parser.add_argument("--compute_map", default=False, help="if True computes mAP every tenth batch")
parser.add_argument("--multiscale_training", default=True, help="allow for multi-scale training")
使用parse_args方法解析命令行参数
opt = parser.parse_args()
使用TensorFlow 2.0以上版本中的tf.summary模块创建日志记录器(Logger)
def __init__(self, log_dir):
Create a summary writer logging to log_dir.
这个类的构造函数接受一个参数log_dir,它表示日志文件将要保存的目录。
函数的作用是创建一个日志记录器,用于记录TensorFlow的摘要信息(例如训练过程中的损失、准确率等)。
self.writer = tf.summary.create_file_writer(log_dir)
调用Logger,创建目录
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("output", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)
定义 parse_data_config 的函数,它用于解析数据配置文件(文件内容分类种类,训练集路径,测试集路径,文件名称路径)

def parse_data_config(path):
"""Parses the data configuration file"""
options['gpus'] = '0,1,2,3'
options['num_workers'] = '10'
with open(path, 'r') as fp:
if line == '' or line.startswith('#'):
key, value = line.split('=')
options[key.strip()] = value.strip()
调用函数
data_config = parse_data_config(opt.data_config)
train_path = data_config["train"]
valid_path = data_config["valid"]
定义了一个名为 load_classes 的函数,它用于从指定路径加载类别标签
Loads class labels at 'path'
names = fp.read().split("\n")[:-1]
调用函数
class_names = load_classes(data_config["names"])
定义了一个名为 parse_model_config 的函数,它用于解析 YOLOv3 模型的配置文件

读取文件划分如卷积、池化、上采样、路由、快捷连接和 YOLO 层
def parse_model_config(path):
"""Parses the yolo-v3 layer configuration file and returns module definitions"""
lines = file.read().split('\n')
lines = [x for x in lines if x and not x.startswith('#')]
lines = [x.rstrip().lstrip() for x in lines]
module_defs[-1]['type'] = line[1:-1].rstrip()
if module_defs[-1]['type'] == 'convolutional':
module_defs[-1]['batch_normalize'] = 0
key, value = line.split("=")
module_defs[-1][key.rstrip()] = value.strip()
这个函数的目的是将 YOLOv3 模型配置文件中的文本描述转换成 PyTorch 可以理解的网络层模块。它首先处理超参数,然后逐个处理每个模块定义,根据模块的类型(如卷积、池化、上采样、路由、快捷连接和 YOLO 层)创建相应的 PyTorch 层,并添加到 module_list 中。
def create_modules(module_defs):
Constructs module list of layer blocks from module configuration in module_defs
hyperparams = module_defs.pop(0)
output_filters = [int(hyperparams["channels"])]
module_list = nn.ModuleList()
for module_i, module_def in enumerate(module_defs):
modules = nn.Sequential()
if module_def["type"] == "convolutional":
bn = int(module_def["batch_normalize"])
filters = int(module_def["filters"])
kernel_size = int(module_def["size"])
pad = (kernel_size - 1) // 2
in_channels=output_filters[-1],
stride=int(module_def["stride"]),
modules.add_module(f"batch_norm_{module_i}", nn.BatchNorm2d(filters, momentum=0.9))
if module_def["activation"] == "leaky":
modules.add_module(f"leaky_{module_i}", nn.LeakyReLU(0.1))
elif module_def["type"] == "maxpool":
kernel_size = int(module_def["size"])
stride = int(module_def["stride"])
if kernel_size == 2 and stride == 1:
modules.add_module(f"_debug_padding_{module_i}", nn.ZeroPad2d((0, 1, 0, 1)))
maxpool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=int((kernel_size - 1) // 2))
modules.add_module(f"maxpool_{module_i}", maxpool)
elif module_def["type"] == "upsample":
upsample = Upsample(scale_factor=int(module_def["stride"]), mode="nearest")
modules.add_module(f"upsample_{module_i}", upsample)
elif module_def["type"] == "route":
layers = [int(x) for x in module_def["layers"].split(",")]
filters = sum([output_filters[1:][i] for i in layers])
modules.add_module(f"route_{module_i}", EmptyLayer())
elif module_def["type"] == "shortcut":
filters = output_filters[1:][int(module_def["from"])]
modules.add_module(f"shortcut_{module_i}", EmptyLayer())
elif module_def["type"] == "yolo":
anchor_idxs = [int(x) for x in module_def["mask"].split(",")]
anchors = [int(x) for x in module_def["anchors"].split(",")]
anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
anchors = [anchors[i] for i in anchor_idxs]
num_classes = int(module_def["classes"])
img_size = int(hyperparams["height"])
yolo_layer = YOLOLayer(anchors, num_classes, img_size)
modules.add_module(f"yolo_{module_i}", yolo_layer)
module_list.append(modules)
output_filters.append(filters)
return hyperparams, module_list
定义了一个名为 Darknet 的类,它是用于构建 YOLOv3 目标检测模型的 PyTorch 神经网络类。
class Darknet(nn.Module):
"""YOLOv3 object detection model"""
def __init__(self, config_path, img_size=416):
super(Darknet, self).__init__()
self.module_defs = parse_model_config(config_path)
self.hyperparams, self.module_list = create_modules(self.module_defs)
self.yolo_layers = [layer[0] for layer in self.module_list if hasattr(layer[0], "metrics")]
self.header_info = np.array([0, 0, 0, self.seen, 0], dtype=np.int32)
调用函数
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)