一、官方源代码
#!/usr/bin/env python3
#
# train an SSD detection model on Pascal VOC or Open Images datasets
# https://github.com/dusty-nv/jetson-inference/blob/master/docs/pytorch-ssd.md
#
import os
import sys
import logging
import argparse
import datetime
import itertools
import torch
from torch.utils.data import DataLoader, ConcatDataset
from torch.utils.tensorboard import SummaryWriter
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
from vision.utils.misc import Timer, freeze_net_layers, store_labels
from vision.ssd.ssd import MatchPrior
from vision.ssd.vgg_ssd import create_vgg_ssd
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite
from vision.datasets.voc_dataset import VOCDataset
from vision.datasets.open_images import OpenImagesDataset
from vision.nn.multibox_loss import MultiboxLoss
from vision.ssd.config import vgg_ssd_config
from vision.ssd.config import mobilenetv1_ssd_config
from vision.ssd.config import squeezenet_ssd_config
from vision.ssd.data_preprocessing import TrainAugmentation, TestTransform
from eval_ssd import MeanAPEvaluator
DEFAULT_PRETRAINED_MODEL='models/mobilenet-v1-ssd-mp-0_675.pth'
parser = argparse.ArgumentParser(
description='Single Shot MultiBox Detector Training With PyTorch')
# Params for datasets
parser.add_argument("--dataset-type", default="open_images", type=str,
help='Specify dataset type. Currently supports voc and open_images.')
parser.add_argument('--datasets', '--data', nargs='+', default=["data"], help='Dataset directory path')
parser.add_argument('--balance-data', action='store_true',
help="Balance training data by down-sampling more frequent labels.")
# Params for network
parser.add_argument('--net', default="mb1-ssd",
help="The network architecture, it can be mb1-ssd, mb1-ssd-lite, mb2-ssd-lite or vgg16-ssd.")
parser.add_argument('--resolution', type=int, default=300,
help="the NxN pixel resolution of the model (can be changed for mb1-ssd only)")
parser.add_argument('--freeze-base-net', action='store_true',
help="Freeze base net layers.")
parser.add_argument('--freeze-net', action='store_true',
help="Freeze all the layers except the prediction head.")
parser.add_argument('--mb2-width-mult', default=1.0, type=float,
help='Width Multiplifier for MobilenetV2')
# Params for loading pretrained basenet or checkpoints.
parser.add_argument('--base-net', help='Pretrained base model')
parser.add_argument('--pretrained-ssd', default=DEFAULT_PRETRAINED_MODEL, type=str, help='Pre-trained base model')
parser.add_argument('--resume', default=None, type=str, help='Checkpoint state_dict file to resume training from')
# Params for SGD
parser.add_argument('--lr', '--learning-rate', default=0.01, type=float,
help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float,
help='Momentum value for optim')
parser.add_argument('--weight-decay', default=5e-4, type=float,
help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.1, type=float,
help='Gamma update for SGD')
parser.add_argument('--base-net-lr', default=0.001, type=float,
help='initial learning rate for base net, or None to use --lr')
parser.add_argument('--extra-layers-lr', default=None, type=float,
help='initial learning rate for the layers not in base net and prediction heads.')
# Scheduler
parser.add_argument('--scheduler', default="cosine", type=str,
help="Scheduler for SGD. It can one of multi-step and cosine")
# Params for Multi-step Scheduler
parser.add_argument('--milestones', default="80,100", type=str,
help="milestones for MultiStepLR")
# Params for Cosine Annealing
parser.add_argument('--t-max', default=100, type=float,
help='T_max value for Cosine Annealing Scheduler.')
# Train params
parser.add_argument('--batch-size', default=4, type=int,
help='Batch size for training')
parser.add_argument('--num-epochs', '--epochs', default=30, type=int,
help='the number epochs')
parser.add_argument('--num-workers', '--workers', default=2, type=int,
help='Number of workers used in dataloading')
parser.add_argument('--validation-epochs', default=1, type=int,
help='the number epochs between running validation')
parser.add_argument('--validation-mean-ap', action='store_true',
help='Perform computation of Mean Average Precision (mAP) during validation')
parser.add_argument('--debug-steps', default=10, type=int,
help='Set the debug log output frequency.')
parser.add_argument('--use-cuda', default=True, action='store_true',
help='Use CUDA to train model')
parser.add_argument('--checkpoint-folder', '--model-dir', default='models/',
help='Directory for saving checkpoint models')
parser.add_argument('--log-level', default='info', type=str,
help='Logging level, one of: debug, info, warning, error, critical (default: info)')
args = parser.parse_args()
logging.basicConfig(stream=sys.stdout, level=getattr(logging, args.log_level.upper(), logging.INFO),
format='%(asctime)s - %(message)s', datefmt="%Y-%m-%d %H:%M:%S")
tensorboard = SummaryWriter(log_dir=os.path.join(args.checkpoint_folder, "tensorboard", f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"))
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")
if args.use_cuda and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logging.info("Using CUDA...")
def train(loader, net, criterion, optimizer, device, debug_steps=100, epoch=-1):
net.train(True)
train_loss = 0.0
train_regression_loss = 0.0
train_classification_loss = 0.0
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num_batches = 0
for i, data in enumerate(loader):
images, boxes, labels = data
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
optimizer.zero_grad()
confidence, locations = net(images)
regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
loss = regression_loss + classification_loss
loss.backward()
optimizer.step()
train_loss += loss.item()
train_regression_loss += regression_loss.item()
train_classification_loss += classification_loss.item()
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
if i and i % debug_steps == 0:
avg_loss = running_loss / debug_steps
avg_reg_loss = running_regression_loss / debug_steps
avg_clf_loss = running_classification_loss / debug_steps
logging.info(
f"Epoch: {epoch}, Step: {i}/{len(loader)}, " +
f"Avg Loss: {avg_loss:.4f}, " +
f"Avg Regression Loss {avg_reg_loss:.4f}, " +
f"Avg Classification Loss: {avg_clf_loss:.4f}"
)
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num_batches += 1
train_loss /= num_batches
train_regression_loss /= num_batches
train_classification_loss /= num_batches
logging.info(
f"Epoch: {epoch}, " +
f"Training Loss: {train_loss:.4f}, " +
f"Training Regression Loss {train_regression_loss:.4f}, " +
f"Training Classification Loss: {train_classification_loss:.4f}"
)
tensorboard.add_scalar('Loss/train', train_loss, epoch)
tensorboard.add_scalar('Regression Loss/train', train_regression_loss, epoch)
tensorboard.add_scalar('Classification Loss/train', train_classification_loss, epoch)
def test(loader, net, criterion, device):
net.eval()
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num = 0
for _, data in enumerate(loader):
images, boxes, labels = data
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
num += 1
with torch.no_grad():
confidence, locations = net(images)
regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
loss = regression_loss + classification_loss
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
return running_loss / num, running_regression_loss / num, running_classification_loss / num
if __name__ == '__main__':
timer = Timer()
logging.info(args)
# make sure that the checkpoint output dir exists
if args.checkpoint_folder:
args.checkpoint_folder = os.path.expanduser(args.checkpoint_folder)
if not os.path.exists(args.checkpoint_folder):
os.mkdir(args.checkpoint_folder)
# select the network architecture and config
if args.net == 'vgg16-ssd':
create_net = create_vgg_ssd
config = vgg_ssd_config
elif args.net == 'mb1-ssd':
create_net = create_mobilenetv1_ssd
config = mobilenetv1_ssd_config
config.set_image_size(args.resolution)
elif args.net == 'mb1-ssd-lite':
create_net = create_mobilenetv1_ssd_lite
config = mobilenetv1_ssd_config
elif args.net == 'sq-ssd-lite':
create_net = create_squeezenet_ssd_lite
config = squeezenet_ssd_config
elif args.net == 'mb2-ssd-lite':
create_net = lambda num: create_mobilenetv2_ssd_lite(num, width_mult=args.mb2_width_mult)
config = mobilenetv1_ssd_config
else:
logging.fatal("The net type is wrong.")
parser.print_help(sys.stderr)
sys.exit(1)
# create data transforms for train/test/val
train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std)
target_transform = MatchPrior(config.priors, config.center_variance,
config.size_variance, 0.5)
test_transform = TestTransform(config.image_size, config.image_mean, config.image_std)
# load datasets (could be multiple)
logging.info("Prepare training datasets.")
datasets = []
for dataset_path in args.datasets:
if args.dataset_type == 'voc':
dataset = VOCDataset(dataset_path, transform=train_transform,
target_transform=target_transform)
label_file = os.path.join(args.checkpoint_folder, "labels.txt")
store_labels(label_file, dataset.class_names)
num_classes = len(dataset.class_names)
elif args.dataset_type == 'open_images':
dataset = OpenImagesDataset(dataset_path,
transform=train_transform, target_transform=target_transform,
dataset_type="train", balance_data=args.balance_data)
label_file = os.path.join(args.checkpoint_folder, "labels.txt")
store_labels(label_file, dataset.class_names)
logging.info(dataset)
num_classes = len(dataset.class_names)
else:
raise ValueError(f"Dataset type {args.dataset_type} is not supported.")
datasets.append(dataset)
# create training dataset
logging.info(f"Stored labels into file {label_file}.")
train_dataset = ConcatDataset(datasets)
logging.info("Train dataset size: {}".format(len(train_dataset)))
train_loader = DataLoader(train_dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True)
# create validation dataset
logging.info("Prepare Validation datasets.")
if args.dataset_type == "voc":
val_dataset = VOCDataset(dataset_path, transform=test_transform,
target_transform=target_transform, is_test=True)
elif args.dataset_type == 'open_images':
val_dataset = OpenImagesDataset(dataset_path,
transform=test_transform, target_transform=target_transform,
dataset_type="test")
logging.info(val_dataset)
logging.info("Validation dataset size: {}".format(len(val_dataset)))
val_loader = DataLoader(val_dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=False)
# create the network
logging.info("Build network.")
net = create_net(num_classes)
min_loss = -10000.0
last_epoch = -1
# prepare eval dataset (for mAP computation)
if args.validation_mean_ap:
if args.dataset_type == "voc":
eval_dataset = VOCDataset(dataset_path, is_test=True)
elif args.dataset_type == 'open_images':
eval_dataset = OpenImagesDataset(dataset_path, dataset_type="test")
eval = MeanAPEvaluator(eval_dataset, net, arch=args.net, eval_dir=os.path.join(args.checkpoint_folder, 'eval_results'))
# freeze certain layers (if requested)
base_net_lr = args.base_net_lr if args.base_net_lr is not None else args.lr
extra_layers_lr = args.extra_layers_lr if args.extra_layers_lr is not None else args.lr
if args.freeze_base_net:
logging.info("Freeze base net.")
freeze_net_layers(net.base_net)
params = itertools.chain(net.source_layer_add_ons.parameters(), net.extras.parameters(),
net.regression_headers.parameters(), net.classification_headers.parameters())
params = [
{'params': itertools.chain(
net.source_layer_add_ons.parameters(),
net.extras.parameters()
), 'lr': extra_layers_lr},
{'params': itertools.chain(
net.regression_headers.parameters(),
net.classification_headers.parameters()
)}
]
elif args.freeze_net:
freeze_net_layers(net.base_net)
freeze_net_layers(net.source_layer_add_ons)
freeze_net_layers(net.extras)
params = itertools.chain(net.regression_headers.parameters(), net.classification_headers.parameters())
logging.info("Freeze all the layers except prediction heads.")
else:
params = [
{'params': net.base_net.parameters(), 'lr': base_net_lr},
{'params': itertools.chain(
net.source_layer_add_ons.parameters(),
net.extras.parameters()
), 'lr': extra_layers_lr},
{'params': itertools.chain(
net.regression_headers.parameters(),
net.classification_headers.parameters()
)}
]
# load a previous model checkpoint (if requested)
timer.start("Load Model")
if args.resume:
logging.info(f"Resuming from the model {args.resume}")
net.load(args.resume)
elif args.base_net:
logging.info(f"Init from base net {args.base_net}")
net.init_from_base_net(args.base_net)
elif args.pretrained_ssd:
logging.info(f"Init from pretrained SSD {args.pretrained_ssd}")
if not os.path.exists(args.pretrained_ssd) and args.pretrained_ssd == DEFAULT_PRETRAINED_MODEL:
os.system(f"wget --quiet --show-progress --progress=bar:force:noscroll --no-check-certificate https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O {DEFAULT_PRETRAINED_MODEL}")
net.init_from_pretrained_ssd(args.pretrained_ssd)
logging.info(f'Took {timer.end("Load Model"):.2f} seconds to load the model.')
# move the model to GPU
net.to(DEVICE)
# define loss function and optimizer
criterion = MultiboxLoss(config.priors, iou_threshold=0.5, neg_pos_ratio=3,
center_variance=0.1, size_variance=0.2, device=DEVICE)
optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
logging.info(f"Learning rate: {args.lr}, Base net learning rate: {base_net_lr}, "
+ f"Extra Layers learning rate: {extra_layers_lr}.")
# set learning rate policy
if args.scheduler == 'multi-step':
logging.info("Uses MultiStepLR scheduler.")
milestones = [int(v.strip()) for v in args.milestones.split(",")]
scheduler = MultiStepLR(optimizer, milestones=milestones,
gamma=0.1, last_epoch=last_epoch)
elif args.scheduler == 'cosine':
logging.info("Uses CosineAnnealingLR scheduler.")
scheduler = CosineAnnealingLR(optimizer, args.t_max, last_epoch=last_epoch)
else:
logging.fatal(f"Unsupported Scheduler: {args.scheduler}.")
parser.print_help(sys.stderr)
sys.exit(1)
# train for the desired number of epochs
logging.info(f"Start training from epoch {last_epoch + 1}.")
for epoch in range(last_epoch + 1, args.num_epochs):
train(train_loader, net, criterion, optimizer, device=DEVICE, debug_steps=args.debug_steps, epoch=epoch)
scheduler.step()
if epoch % args.validation_epochs == 0 or epoch == args.num_epochs - 1:
val_loss, val_regression_loss, val_classification_loss = test(val_loader, net, criterion, DEVICE)
logging.info(
f"Epoch: {epoch}, " +
f"Validation Loss: {val_loss:.4f}, " +
f"Validation Regression Loss {val_regression_loss:.4f}, " +
f"Validation Classification Loss: {val_classification_loss:.4f}"
)
tensorboard.add_scalar('Loss/val', val_loss, epoch)
tensorboard.add_scalar('Regression Loss/val', val_regression_loss, epoch)
tensorboard.add_scalar('Classification Loss/val', val_classification_loss, epoch)
if args.validation_mean_ap:
mean_ap, class_ap = eval.compute()
eval.log_results(mean_ap, class_ap, f"Epoch: {epoch}, ")
tensorboard.add_scalar('Mean Average Precision/val', mean_ap, epoch)
for i in range(len(class_ap)):
tensorboard.add_scalar(f"Class Average Precision/{eval_dataset.class_names[i+1]}", class_ap[i], epoch)
model_path = os.path.join(args.checkpoint_folder, f"{args.net}-Epoch-{epoch}-Loss-{val_loss}.pth")
net.save(model_path)
logging.info(f"Saved model {model_path}")
logging.info("Task done, exiting program.")
tensorboard.close()
使用方法
所有参数使用默认值训练
python3 train_ssd.py --dataset-type=voc --data=data/drone/ --model-dir=models/drone/ --resolution=512
调整一些参数进行训练
python3 train_ssd.py --dataset-type=voc --data=data/drone/ --model-dir=models/drone_512px_lr001_bs4_epochs1000/ --resolution=512 --num-workers=2 --batch-size=4 --lr=0.01 --epochs=1000
在内存比较小的系统中训练，数据加载的工作线程"–num-workers"要减少，甚至等于0，否则会出错！！！
如在jetson nano中
python3 train_ssd.py --dataset-type=voc --data=data/drone/ --model-dir=models/drone/ --resolution=512 --num-workers=0 --batch-size=4 --epochs=100
注意：训练时传入的分辨率（--resolution）要和采样时图片的分辨率一样，也和转换时传入的分辨率一样，否则效果很差！！！
二、解析代码
加上中文注释
下面是加上中文注释的代码：
```python
#!/usr/bin/env python3
#
# 在Pascal VOC或Open Images数据集上训练SSD检测模型
# https://github.com/dusty-nv/jetson-inference/blob/master/docs/pytorch-ssd.md
#
import os
import sys
import logging
import argparse
import datetime
import itertools
import torch
from torch.utils.data import DataLoader, ConcatDataset
from torch.utils.tensorboard import SummaryWriter
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
from vision.utils.misc import Timer, freeze_net_layers, store_labels
from vision.ssd.ssd import MatchPrior
from vision.ssd.vgg_ssd import create_vgg_ssd
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite
from vision.datasets.voc_dataset import VOCDataset
from vision.datasets.open_images import OpenImagesDataset
from vision.nn.multibox_loss import MultiboxLoss
from vision.ssd.config import vgg_ssd_config
from vision.ssd.config import mobilenetv1_ssd_config
from vision.ssd.config import squeezenet_ssd_config
from vision.ssd.data_preprocessing import TrainAugmentation, TestTransform
from eval_ssd import MeanAPEvaluator
# 默认预训练模型路径
DEFAULT_PRETRAINED_MODEL='models/mobilenet-v1-ssd-mp-0_675.pth'
# 定义参数解析器
parser = argparse.ArgumentParser(
description='使用PyTorch训练单发多盒检测器(SSD)')
# 数据集参数
parser.add_argument("--dataset-type", default="open_images", type=str,
help='指定数据集类型。目前支持voc和open_images。')
parser.add_argument('--datasets', '--data', nargs='+', default=["data"], help='数据集目录路径')
parser.add_argument('--balance-data', action='store_true',
help="通过下采样频繁标签来平衡训练数据。")
# 网络参数
parser.add_argument('--net', default="mb1-ssd",
help="网络架构，可以是mb1-ssd, mb1-ssd-lite, mb2-ssd-lite或vgg16-ssd。")
parser.add_argument('--resolution', type=int, default=300,
help="模型的NxN像素分辨率（仅适用于mb1-ssd）。")
parser.add_argument('--freeze-base-net', action='store_true',
help="冻结基础网络层。")
parser.add_argument('--freeze-net', action='store_true',
help="冻结除预测头之外的所有层。")
parser.add_argument('--mb2-width-mult', default=1.0, type=float,
help='MobilenetV2的宽度乘数')
# 加载预训练基础网络或检查点的参数
parser.add_argument('--base-net', help='预训练的基础模型')
parser.add_argument('--pretrained-ssd', default=DEFAULT_PRETRAINED_MODEL, type=str, help='预训练的基础模型')
parser.add_argument('--resume', default=None, type=str, help='从检查点状态字典文件恢复训练')
# SGD参数
parser.add_argument('--lr', '--learning-rate', default=0.01, type=float,
help='初始学习率')
parser.add_argument('--momentum', default=0.9, type=float,
help='优化器的动量值')
parser.add_argument('--weight-decay', default=5e-4, type=float,
help='SGD的权重衰减')
parser.add_argument('--gamma', default=0.1, type=float,
help='SGD的Gamma更新')
parser.add_argument('--base-net-lr', default=0.001, type=float,
help='基础网络的初始学习率，或使用--lr')
parser.add_argument('--extra-layers-lr', default=None, type=float,
help='基础网络和预测头以外层的初始学习率。')
# 学习率调度器
parser.add_argument('--scheduler', default="cosine", type=str,
help="SGD的调度器。可以是multi-step或cosine")
# 多步调度器的参数
parser.add_argument('--milestones', default="80,100", type=str,
help="MultiStepLR的里程碑")
# 余弦退火调度器的参数
parser.add_argument('--t-max', default=100, type=float,
help='余弦退火调度器的T_max值。')
# 训练参数
parser.add_argument('--batch-size', default=4, type=int,
help='训练的批量大小')
parser.add_argument('--num-epochs', '--epochs', default=30, type=int,
help='训练的周期数')
parser.add_argument('--num-workers', '--workers', default=2, type=int,
help='数据加载时使用的工作线程数')
parser.add_argument('--validation-epochs', default=1, type=int,
help='运行验证的周期数')
parser.add_argument('--validation-mean-ap', action='store_true',
help='在验证期间计算平均精度均值(mAP)')
parser.add_argument('--debug-steps', default=10, type=int,
help='设置调试日志输出频率。')
parser.add_argument('--use-cuda', default=True, action='store_true',
help='使用CUDA进行模型训练')
parser.add_argument('--checkpoint-folder', '--model-dir', default='models/',
help='保存检查点模型的目录')
parser.add_argument('--log-level', default='info', type=str,
help='日志级别，可以是：debug, info, warning, error, critical (默认: info)')
# 解析命令行参数
args = parser.parse_args()
# 配置日志
logging.basicConfig(stream=sys.stdout, level=getattr(logging, args.log_level.upper(), logging.INFO),
format='%(asctime)s - %(message)s', datefmt="%Y-%m-%d %H:%M:%S")
# 设置TensorBoard日志记录器
tensorboard = SummaryWriter(log_dir=os.path.join(args.checkpoint_folder, "tensorboard", f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"))
# 检查CUDA是否可用，并选择设备
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")
if args.use_cuda and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logging.info("Using CUDA...")
# 训练函数
def train(loader, net, criterion, optimizer, device, debug_steps=100, epoch=-1):
net.train(True)
train_loss = 0.0
train_regression_loss = 0.0
train_classification_loss = 0.0
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num_batches = 0
for i, data in enumerate(loader):
images, boxes, labels = data
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
optimizer.zero_grad()
confidence, locations = net(images)
regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
loss = regression_loss + classification_loss
loss.backward()
optimizer.step()
train_loss += loss.item()
train_regression_loss += regression_loss.item()
train_classification_loss += classification_loss.item()
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
if i and i % debug_steps == 0:
avg_loss = running_loss / debug_steps
avg_reg_loss = running_regression_loss / debug_steps
avg_clf_loss = running_classification_loss / debug_steps
logging.info(
f"Epoch: {epoch}, Step: {i}/{len(loader)}, " +
f"Avg Loss: {avg_loss:.4f}, " +
f"Avg Regression Loss {avg_reg_loss:.4f}, " +
f"Avg Classification Loss: {avg_clf_loss:.4f}"
)
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num_batches += 1
train_loss /= num_batches
train_regression_loss /= num_batches
train_classification_loss /= num_batches
logging.info(
f"Epoch: {epoch}, " +
f"Training Loss: {train_loss:.4f}, " +
f"Training Regression Loss {train_regression_loss:.4f}, " +
f"Training Classification Loss: {train_classification_loss:.4f}"
)
tensorboard.add_scalar('Loss/train', train_loss, epoch)
tensorboard.add_scalar('Regression Loss/train', train_regression_loss, epoch)
tensorboard.add_scalar('Classification Loss/train', train_classification_loss, epoch)
# 测试函数
def test(loader, net, criterion, device):
net.eval()
running_loss = 0.0
running_regression_loss = 0.0
running_classification
详细解析
下面是对代码每一部分的详细解析，包括中文注释：
导入必要的模块
import os
import sys
import logging
import argparse
import datetime
import itertools
import torch
from torch.utils.data import DataLoader, ConcatDataset
from torch.utils.tensorboard import SummaryWriter
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
from vision.utils.misc import Timer, freeze_net_layers, store_labels
from vision.ssd.ssd import MatchPrior
from vision.ssd.vgg_ssd import create_vgg_ssd
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite
from vision.datasets.voc_dataset import VOCDataset
from vision.datasets.open_images import OpenImagesDataset
from vision.nn.multibox_loss import MultiboxLoss
from vision.ssd.config import vgg_ssd_config
from vision.ssd.config import mobilenetv1_ssd_config
from vision.ssd.config import squeezenet_ssd_config
from vision.ssd.data_preprocessing import TrainAugmentation, TestTransform
from eval_ssd import MeanAPEvaluator
- 导入操作系统接口（
os、sys） - 导入日志处理模块（
logging） - 导入命令行参数解析模块（
argparse） - 导入日期时间处理模块（
datetime） - 导入用于高效组合迭代工具（
itertools） - 导入PyTorch库（
torch） - 导入数据加载和处理工具（
DataLoader、ConcatDataset） - 导入TensorBoard日志记录工具（
SummaryWriter） - 导入学习率调度器（
CosineAnnealingLR、MultiStepLR） - 导入其他工具和SSD模型相关模块
设置默认预训练模型路径
DEFAULT_PRETRAINED_MODEL='models/mobilenet-v1-ssd-mp-0_675.pth'
- 定义预训练模型的默认路径。
参数解析
parser = argparse.ArgumentParser(description='使用PyTorch训练单发多盒检测器(SSD)')
- 使用
argparse定义命令行参数解析器。
数据集参数
parser.add_argument("--dataset-type", default="open_images", type=str, help='指定数据集类型。目前支持voc和open_images。')
parser.add_argument('--datasets', '--data', nargs='+', default=["data"], help='数据集目录路径')
parser.add_argument('--balance-data', action='store_true', help="通过下采样频繁标签来平衡训练数据。")
- 定义数据集类型、路径和数据平衡参数。
--balance-data参数详细解析
参数 parser.add_argument('--balance-data', action='store_true', help="通过下采样频繁标签来平衡训练数据。") 的详细说明如下：
参数的作用 这个参数用于在训练数据集存在类别不平衡的情况下，通过下采样频繁标签来平衡训练数据。类别不平衡是指某些类别的样本数量远多于其他类别，这可能导致模型在训练过程中偏向预测这些频繁出现的类别，影响模型的泛化能力和准确率。
参数详细解释
--balance-data：这是参数的名称。当在命令行中使用这个参数时，argparse会将其值设置为True。如果不使用这个参数，则其值为False。action='store_true'：这意味着当指定--balance-data参数时，args.balance_data的值将被设置为True。如果没有指定，args.balance_data的值将是False。help：这是对参数的简要描述。当使用--help或-h命令行选项时，会显示这段帮助文本。
为什么需要数据平衡 在实际应用中，数据集通常是不平衡的。例如，在目标检测任务中，有些类别的对象可能出现频率很高，而有些类别的对象则很少出现。这种不平衡会导致模型在训练过程中更偏向于预测频繁出现的类别，从而忽视稀有类别，导致模型在稀有类别上的性能较差。
数据平衡的具体方法 通过下采样（down-sampling）频繁标签，可以使每个类别在训练数据中出现的次数更加均衡。具体方法如下：
- 下采样：减少频繁出现的类别样本的数量，使其与稀有类别样本的数量相当。这样可以避免模型对某一类样本过度拟合。
- 上采样（如果需要，可以结合上采样）：增加稀有类别样本的数量，使其与频繁类别样本的数量相当。可以通过数据增强（如旋转、翻转、裁剪等）来增加稀有类别的样本数量。
代码示例 在实际代码中，使用 --balance-data 参数可以在加载数据集时进行数据平衡处理。以下是一个简单的示例：
# 假设 args.balance_data 是由命令行参数解析得到的布尔值
if args.balance_data:
logging.info("Balancing training data by down-sampling more frequent labels.")
for dataset_path in args.datasets:
if args.dataset_type == 'voc':
dataset = VOCDataset(dataset_path, transform=train_transform, target_transform=target_transform)
dataset.balance_data() # 假设 VOCDataset 类有一个 balance_data 方法
elif args.dataset_type == 'open_images':
dataset = OpenImagesDataset(dataset_path, transform=train_transform, target_transform=target_transform, dataset_type="train", balance_data=args.balance_data)
datasets.append(dataset) else:
logging.info("Using original training data without balancing.")
for dataset_path in args.datasets:
if args.dataset_type == 'voc':
dataset = VOCDataset(dataset_path, transform=train_transform, target_transform=target_transform)
elif args.dataset_type == 'open_images':
dataset = OpenImagesDataset(dataset_path, transform=train_transform, target_transform=target_transform, dataset_type="train", balance_data=args.balance_data)
datasets.append(dataset) ```
###### 总结 `--balance-data` 参数用于在训练数据中存在类别不平衡的情况下，通过下采样频繁出现的标签来平衡训练数据。这有助于提高模型的泛化能力，使
其在稀有类别上的表现更好。具体实现方式可能会根据数据集和任务的不同而有所差异。
#### 网络参数
```python
parser.add_argument('--net', default="mb1-ssd", help="网络架构，可以是mb1-ssd, mb1-ssd-lite, mb2-ssd-lite或vgg16-ssd。")
parser.add_argument('--resolution', type=int, default=300, help="模型的NxN像素分辨率（仅适用于mb1-ssd）。")
parser.add_argument('--freeze-base-net', action='store_true', help="冻结基础网络层。")
parser.add_argument('--freeze-net', action='store_true', help="冻结除预测头之外的所有层。")
parser.add_argument('--mb2-width-mult', default=1.0, type=float, help='MobilenetV2的宽度乘数')
- 定义网络架构、分辨率、冻结层设置和宽度乘数。
--net参数详细解析
参数 parser.add_argument('--net', default="mb1-ssd", help="网络架构，可以是mb1-ssd, mb1-ssd-lite, mb2-ssd-lite或vgg16-ssd。") 的详细说明如下：
参数的作用 这个参数用于指定训练过程中所使用的神经网络架构。不同的网络架构有不同的特点和适用场景，通过指定合适的网络架构，可以优化模型的性能、训练速度和资源占用。
参数详细解释
--net：这是参数的名称。在命令行中使用这个参数可以指定所需的网络架构。default="mb1-ssd"：这是默认值，如果在命令行中没有指定--net参数，则默认使用mb1-ssd架构。help：这是对参数的简要描述。当使用--help或-h命令行选项时，会显示这段帮助文本，解释可以使用的网络架构选项。
可选的网络架构 该参数支持以下几种网络架构：
- mb1-ssd：MobileNetV1-SSD，是一种轻量级的目标检测网络，适合在资源受限的设备（如嵌入式设备和移动设备）上运行。
- mb1-ssd-lite：MobileNetV1-SSD Lite，是
mb1-ssd的轻量化版本，进一步减少了计算量和参数量，适合对资源占用有更高要求的场景。 - mb2-ssd-lite：MobileNetV2-SSD Lite，基于MobileNetV2的架构，同样适用于轻量级目标检测任务，具有更好的特征提取能力。
- vgg16-ssd：基于VGG16的SSD模型，具有较高的精度和较大的计算量，适合在计算资源充足的环境中使用。
如何选择网络架构 选择网络架构时，需要根据具体应用场景和需求考虑以下几点：
- 模型精度：在计算资源允许的情况下，选择精度较高的模型（如
vgg16-ssd）。 - 计算资源：在资源受限的设备上，选择轻量级的模型（如
mb1-ssd、mb1-ssd-lite或mb2-ssd-lite）。 - 推理速度：在实时性要求较高的应用中，选择计算量较小、推理速度较快的模型。
- 任务需求：根据具体的目标检测任务，选择适合的模型架构。例如，对于移动设备上的实时目标检测任务，可以选择
mb1-ssd或mb2-ssd-lite。
代码示例 在代码中，使用 --net 参数可以选择不同的网络架构：
# 根据命令行参数选择网络架构和配置 if args.net == 'vgg16-ssd':
create_net = create_vgg_ssd
config = vgg_ssd_config elif args.net == 'mb1-ssd':
create_net = create_mobilenetv1_ssd
config = mobilenetv1_ssd_config
config.set_image_size(args.resolution) elif args.net == 'mb1-ssd-lite':
create_net = create_mobilenetv1_ssd_lite
config = mobilenetv1_ssd_config elif args.net == 'mb2-ssd-lite':
create_net = lambda num: create_mobilenetv2_ssd_lite(num, width_mult=args.mb2_width_mult)
config = mobilenetv1_ssd_config else:
logging.fatal("The net type is wrong.")
parser.print_help(sys.stderr)
sys.exit(1) ```
###### 总结 `--net` 参数用于指定训练过程中所使用的神经网络架构。选择合适的网络架构可以根据具体应用场景和需求来优化模型的性能、训练速度和资源占用。
不同的网络架构适用于不同的目标检测任务，用户可以根据自身的需求选择合适的架构进行训练。
##### --mb2-width-mul参数的详细解析
参数 `parser.add_argument('--mb2-width-mult', default=1.0, type=float, help='MobilenetV2的宽度乘数')` 的详细解析如下：
###### 参数的作用 这个参数用于调整MobileNetV2网络的宽度乘数（width multiplier）。宽度乘数是一个缩放因子，用于控制网络中每一层的卷积核数量。通过调整宽度乘数，可以在计算效率和模型精度之间进行权衡。参数
--mb2-width-mult 只有在 --net=mb2-ssd-lite 时才起作用。这个参数专门用于调整 MobileNetV2 的宽度乘数，而 mb2-ssd-lite 是基于 MobileNetV2 的网络架构。其他网络架构（如 mb1-ssd, mb1-ssd-lite, vgg16-ssd）不会使用这个参数。
###### 参数详细解释
- `--mb2-width-mult`：这是参数的名称。在命令行中使用这个参数可以指定MobileNetV2网络的宽度乘数。
- `default=1.0`：这是默认值，如果在命令行中没有指定 `--mb2-width-mult` 参数，则默认使用 `1.0` 作为宽度乘数。
- `type=float`：表示这个参数的值应为浮点数。
- `help='MobilenetV2的宽度乘数'`：这是对参数的简要描述。当使用 `--help` 或 `-h` 命令行选项时，会显示这段帮助文本，解释参数的用途。
###### 宽度乘数的作用 宽度乘数用于缩放网络的宽度，即卷积核的数量，从而影响模型的计算复杂度和参数数量。具体来说：
- 当宽度乘数为 `1.0` 时，网络的卷积核数量保持不变，即使用默认的卷积核数量。
- 当宽度乘数小于 `1.0` 时，网络的卷积核数量减少，从而降低计算复杂度和参数数量，但可能会导致模型精度下降。
- 当宽度乘数大于 `1.0` 时，网络的卷积核数量增加，从而提高模型精度，但会增加计算复杂度和参数数量。
###### 选择宽度乘数的考虑
- **计算资源**：在计算资源有限的设备上（如嵌入式设备或移动设备），可以选择较小的宽度乘数以减少计算量。
- **模型精度**：在需要较高精度的任务中，可以选择较大的宽度乘数以提高模型的特征提取能力。
- **实时性要求**：在实时性要求高的应用中，可以选择适中的宽度乘数，既保证一定的模型精度，又满足实时性要求。
###### 代码示例 在实际代码中，可以根据命令行参数设置MobileNetV2的宽度乘数：
```python
# 根据命令行参数选择网络架构和配置 if args.net == 'mb2-ssd-lite':
create_net = lambda num: create_mobilenetv2_ssd_lite(num, width_mult=args.mb2_width_mult)
config = mobilenetv1_ssd_config ```
这里的 `create_mobilenetv2_ssd_lite` 函数需要传入 `width_mult` 参数，用于设置网络的宽度乘数。
###### 使用示例 在命令行中使用不同的宽度乘数： ```bash
# 使用默认宽度乘数 1.0 python train_ssd.py --net mb2-ssd-lite --mb2-width-mult 1.0
# 使用宽度乘数 0.75，减少计算量 python train_ssd.py --net mb2-ssd-lite --mb2-width-mult 0.75
# 使用宽度乘数 1.5，增加模型精度 python train_ssd.py --net mb2-ssd-lite --mb2-width-mult 1.5 ```
###### 总结 `--mb2-width-mult` 参数用于调整MobileNetV2网络的宽度乘数，通过改变每一层的卷积核数量来控制模型的计算复杂度和参数数量。选择合适的宽度乘数可以在计算资源、模型精度和实时性要求之间找到最佳平衡。
#### 加载预训练模型或检查点的参数
```python
parser.add_argument('--base-net', help='预训练的基础模型')
parser.add_argument('--pretrained-ssd', default=DEFAULT_PRETRAINED_MODEL, type=str, help='预训练的基础模型')
parser.add_argument('--resume', default=None, type=str, help='从检查点状态字典文件恢复训练')
- 定义加载预训练基础网络或检查点的参数。
--base-net和--pretrained-ssd参数的详细解析
参数 --base-net 和 --pretrained-ssd 都与预训练模型的加载有关，但它们的作用和使用场景有所不同。下面是对这两个参数的详细解析：
--base-net 参数
parser.add_argument('--base-net', help='预训练的基础模型')
作用
- 定义：这个参数用于指定一个预训练的基础模型文件路径。基础模型通常指的是网络架构的主干部分（backbone），例如 VGG、ResNet 或 MobileNet，而不包括 SSD 的检测头部分。
- 用途：主要用于在训练新的 SSD 模型时，利用一个已经预训练好的基础网络来初始化模型的主干部分。这可以加快训练速度并提升模型性能，因为预训练的基础网络已经学到了有用的特征。
使用示例
在命令行中指定 --base-net 参数：
python train_ssd.py --base-net path/to/pretrained/base_net.pth
在代码中处理 --base-net 参数：
if args.base_net:
logging.info(f"Init from base net {args.base_net}")
net.init_from_base_net(args.base_net)
--pretrained-ssd 参数
parser.add_argument('--pretrained-ssd', default=DEFAULT_PRETRAINED_MODEL, type=str, help='预训练的基础模型')
作用
- 定义：这个参数用于指定一个完整的预训练 SSD 模型文件路径，包括基础网络和检测头。默认值为
DEFAULT_PRETRAINED_MODEL。 - 用途：用于加载一个完整的预训练 SSD 模型来进行微调或继续训练。与
--base-net不同，--pretrained-ssd加载的是一个完整的 SSD 模型，包含了基础网络和检测头的权重。
使用示例
在命令行中指定 --pretrained-ssd 参数：
python train_ssd.py --pretrained-ssd path/to/pretrained/ssd.pth
在代码中处理 --pretrained-ssd 参数：
if args.pretrained_ssd:
logging.info(f"Init from pretrained SSD {args.pretrained_ssd}")
if not os.path.exists(args.pretrained_ssd) and args.pretrained_ssd == DEFAULT_PRETRAINED_MODEL:
os.system(f"wget --quiet --show-progress --progress=bar:force:noscroll --no-check-certificate https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O {DEFAULT_PRETRAINED_MODEL}")
net.init_from_pretrained_ssd(args.pretrained_ssd)
二者对比
| 参数 | 用途 | 适用场景 |
|---|---|---|
--base-net | 只加载预训练的基础网络，不包括检测头 | 在已有预训练的基础网络上构建新的 SSD 模型 |
--pretrained-ssd | 加载完整的预训练 SSD 模型，包括基础网络和检测头 | 进行微调或继续训练已有的完整 SSD 模型 |
代码示例
在实际代码中，这两个参数的使用方式如下：
# 解析命令行参数
args = parser.parse_args()
# 初始化模型
if args.resume:
logging.info(f"Resuming from the model {args.resume}")
net.load(args.resume)
elif args.base_net:
logging.info(f"Init from base net {args.base_net}")
net.init_from_base_net(args.base_net)
elif args.pretrained_ssd:
logging.info(f"Init from pretrained SSD {args.pretrained_ssd}")
if not os.path.exists(args.pretrained_ssd) and args.pretrained_ssd == DEFAULT_PRETRAINED_MODEL:
os.system(f"wget --quiet --show-progress --progress=bar:force:noscroll --no-check-certificate https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O {DEFAULT_PRETRAINED_MODEL}")
net.init_from_pretrained_ssd(args.pretrained_ssd)
基础网络和检测头
在目标检测任务中，基础网络（backbone）和检测头（detection head）是两个关键组件。它们在不同的阶段执行不同的功能，结合起来完成目标检测任务。
基础网络（Backbone）
基础网络是一个预训练的卷积神经网络（CNN），用于提取输入图像的特征。基础网络通常在大型图像分类数据集（如ImageNet）上预训练，以学习到丰富的特征表示。常见的基础网络包括VGG、ResNet、MobileNet等。
功能
- 特征提取：基础网络的主要功能是从输入图像中提取特征图（feature maps），这些特征图包含了输入图像的低级和高级特征，如边缘、纹理和对象形状。
- 多层特征：基础网络通常包含多个卷积层和池化层，这些层级联起来形成多层特征表示。较浅的层捕捉低级特征，较深的层捕捉高级特征。
检测头（Detection Head）
检测头是附加在基础网络上的一组层，用于完成具体的目标检测任务。检测头使用基础网络提取的特征图来预测目标的边界框和类别。
功能
- 边界框回归：检测头包含回归层，用于预测目标在图像中的位置，即目标的边界框。
- 类别预测：检测头还包含分类层，用于预测目标的类别，即识别图像中的目标属于哪个类别。
- 多尺度检测：检测头通常设计成多尺度的，以便检测不同大小的目标。SSD（Single Shot MultiBox Detector）就是一个典型的多尺度检测头。
结构示意图
输入图像
|
v
基础网络 (Backbone)
|
v
特征图 (Feature Maps)
|
v
检测头 (Detection Head)
|
v
预测结果 (Predictions: 类别 + 边界框)
示例
以SSD模型为例，描述基础网络和检测头的结构：
基础网络（Backbone）：
- 使用VGG16或MobileNet等预训练的卷积神经网络。
- 提取输入图像的多尺度特征图。
检测头（Detection Head）：
- 在不同尺度的特征图上添加卷积层。
- 每个特征图上的卷积层输出边界框的坐标和类别得分。
- 通过多尺度的预测实现对不同大小目标的检测。
代码示例
以下是如何使用基础网络和检测头初始化SSD模型的示例代码：
# 定义SSD模型
class SSD(nn.Module):
def __init__(self, backbone, num_classes):
super(SSD, self).__init__()
self.backbone = backbone # 基础网络
self.detection_head = DetectionHead(num_classes) # 检测头
def forward(self, x):
features = self.backbone(x) # 提取特征
predictions = self.detection_head(features) # 生成检测结果
return predictions
# 初始化基础网络和检测头
backbone = create_vgg_backbone() # 创建VGG基础网络
ssd_model = SSD(backbone, num_classes=21) # 初始化SSD模型，检测头有21个类别
加载预训练的基础网络和检测头
使用预训练的基础网络和检测头，可以加快模型训练并提升性能：
if args.base_net:
logging.info(f"Init from base net {args.base_net}")
net.init_from_base_net(args.base_net) # 加载预训练的基础网络权重
if args.pretrained_ssd:
logging.info(f"Init from pretrained SSD {args.pretrained_ssd}")
net.init_from_pretrained_ssd(args.pretrained_ssd) # 加载预训练的完整SSD模型
总结
- 基础网络（Backbone）：主要负责从输入图像中提取特征，常用预训练模型如VGG、ResNet、MobileNet等。
- 检测头（Detection Head）：利用基础网络提取的特征，进行边界框回归和类别预测，完成具体的目标检测任务。
- 结合预训练的基础网络和检测头，可以有效提高模型的性能和训练效率。
总结
--base-net参数用于加载预训练的基础网络，在训练新的 SSD 模型时提供一个已经学到有用特征的起点。--pretrained-ssd参数用于加载完整的预训练 SSD 模型，包括基础网络和检测头，用于微调或继续训练。- 这两个参数都可以显著加速训练过程，提高模型性能，但适用的场景有所不同，需要根据具体需求选择使用哪个参数。
SGD参数
parser.add_argument('--lr', '--learning-rate', default=0.01, type=float, help='初始学习率')
parser.add_argument('--momentum', default=0.9, type=float, help='优化器的动量值')
parser.add_argument('--weight-decay', default=5e-4, type=float, help='SGD的权重衰减')
parser.add_argument('--gamma', default=0.1, type=float, help='SGD的Gamma更新')
parser.add_argument('--base-net-lr', default=0.001, type=float, help='基础网络的初始学习率，或使用--lr')
parser.add_argument('--extra-layers-lr', default=None, type=float, help='基础网络和预测头以外层的初始学习率。')
- 定义学习率、动量、权重衰减等参数。
学习率调度器参数
parser.add_argument('--scheduler', default="cosine", type=str, help="SGD的调度器。可以是multi-step或cosine")
parser.add_argument('--milestones', default="80,100", type=str, help="MultiStepLR的里程碑")
parser.add_argument('--t-max', default=100, type=float, help='余弦退火调度器的T_max值。')
- 定义学习率调度器及其相关参数。
训练参数
parser.add_argument('--batch-size', default=4, type=int, help='训练的批量大小')
parser.add_argument('--num-epochs', '--epochs', default=30, type=int, help='训练的周期数')
parser.add_argument('--num-workers', '--workers', default=2, type=int, help='数据加载时使用的工作线程数')
parser.add_argument('--validation-epochs', default=1, type=int, help='运行验证的周期数')
parser.add_argument('--validation-mean-ap', action='store_true', help='在验证期间计算平均精度均值(mAP)')
parser.add_argument('--debug-steps', default=10, type=int, help='设置调试日志输出频率。')
parser.add_argument('--use-cuda', default=True, action='store_true', help='使用CUDA进行模型训练')
parser.add_argument('--checkpoint-folder', '--model-dir', default='models/', help='保存检查点模型的目录')
parser.add_argument('--log-level', default='info', type=str, help='日志级别，可以是：debug, info, warning, error, critical (默认: info)')
- 定义批量大小、周期数、数据加载线程数、验证周期、调试日志输出频率等参数。
--batch-size参数的详细解析
参数 --batch-size 的详细解析如下：
参数的作用
--batch-size 参数用于指定训练过程中每个批次（batch）包含的样本数量。批量大小是深度学习训练过程中的一个关键超参数，直接影响模型训练的效率、内存占用和训练稳定性。
参数详细解释
--batch-size：这是参数的名称。在命令行中使用这个参数可以指定训练时的批量大小。default=4：这是默认值，如果在命令行中没有指定--batch-size参数，则默认使用4作为批量大小。type=int：表示这个参数的值应为整数。help='训练的批量大小'：这是对参数的简要描述。当使用--help或-h命令行选项时，会显示这段帮助文本，解释参数的用途。
批量大小的影响
批量大小对深度学习模型的训练有多方面的影响：
-
训练效率：
- 大批量（Large Batch Size）：较大的批量大小可以更高效地利用硬件资源（如GPU），减少参数更新的频率，从而加快训练过程。
- 小批量（Small Batch Size）：较小的批量大小通常会增加参数更新的频率，使模型参数更快地进行调整，但每次更新计算量较小，可能导致训练时间较长。
-
内存占用：
- 大批量：较大的批量大小需要更多的显存（GPU内存）来存储每批次的样本数据和梯度信息。如果显存不足，可能会导致内存溢出。
- 小批量：较小的批量大小占用较少的显存，适合在显存较小的硬件设备上训练模型。
-
训练稳定性：
- 大批量：较大的批量大小会使每次参数更新更平稳，收敛过程更稳定，但可能会陷入局部最优解。
- 小批量：较小的批量大小会引入更多的噪声，使参数更新更频繁，可能帮助模型跳出局部最优解，但也可能导致训练过程不稳定。
-
模型性能：
- 批量大小的选择可以影响模型的最终性能。通常需要通过实验调整批量大小，以找到最佳的训练效果。
代码示例
在命令行中指定 --batch-size 参数：
# 使用批量大小为 4
python train_ssd.py --batch-size 4
# 使用批量大小为 16
python train_ssd.py --batch-size 16
在代码中处理 --batch-size 参数：
# 解析命令行参数
args = parser.parse_args()
# 配置数据加载器
train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)
val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
总结
--batch-size参数用于指定训练过程中每个批次包含的样本数量。- 批量大小影响训练效率、内存占用、训练稳定性和模型性能。
- 根据硬件资源和具体任务需求选择合适的批量大小，通常需要通过实验调整以找到最佳值。
配置日志和TensorBoard
args = parser.parse_args()
logging.basicConfig(stream=sys.stdout, level=getattr(logging, args.log_level.upper(), logging.INFO),
format='%(asctime)s - %(message)s', datefmt="%Y-%m-%d %H:%M:%S")
tensorboard = SummaryWriter(log_dir=os.path.join(args.checkpoint_folder, "tensorboard", f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"))
- 解析命令行参数，配置日志级别和格式，设置TensorBoard日志记录器。
设备选择
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")
if args.use_cuda and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logging.info("Using CUDA...")
- 检查CUDA是否可用，并选择相应的设备。
训练函数
def train(loader, net, criterion, optimizer, device, debug_steps=100, epoch=-1):
net.train(True)
train_loss = 0.0
train_regression_loss = 0.0
train_classification_loss = 0.0
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num_batches = 0
for i, data in enumerate(loader):
images, boxes, labels = data
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
optimizer.zero_grad()
confidence, locations = net(images)
regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
loss = regression_loss + classification_loss
loss.backward()
optimizer.step()
train_loss += loss.item()
train_regression_loss += regression_loss.item()
train_classification_loss += classification_loss.item()
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
if i and i % debug_steps == 0:
avg_loss = running_loss / debug_steps
avg_reg_loss = running_regression_loss / debug_steps
avg_clf_loss = running_classification_loss / debug_steps
logging.info(
f"Epoch: {epoch}, Step: {i}/{len(loader)}, " +
f"Avg Loss: {avg_loss:.4
定义预训练模型对训练自己的模型有什么用
定义预训练模型对训练自己的模型有以下几个重要的作用：
1. 加速训练
预训练模型已经在大型数据集（如ImageNet）上训练过，因此它们的参数已经优化到某种程度。将预训练模型作为起点，可以显著减少训练时间，因为模型已经具备了一些通用的特征提取能力。
2. 提高模型性能
预训练模型在大型数据集上学到了有用的特征，这些特征可以迁移到新的任务中。这种迁移学习可以提高新任务的模型性能，特别是在目标数据集较小或标注数据较少的情况下。
3. 避免过拟合
当数据集较小时，从头开始训练模型可能导致过拟合，因为模型容易记住训练数据的细节而不能泛化到新数据。预训练模型由于在大量数据上训练过，具备更好的泛化能力，有助于减少过拟合的风险。
4. 适用于不同任务
预训练模型可以用于不同的任务，通过微调（fine-tuning）适应特定任务的需求。例如，可以将预训练的图像分类模型用于目标检测任务，通过在检测数据集上微调模型的参数，使其能够更好地执行目标检测任务。
使用预训练模型的典型流程
- 加载预训练模型：从预训练模型文件中加载权重。
- 微调模型：根据具体任务对模型进行微调。例如，在目标检测任务中，调整模型的最后几层或增加新的层，以适应目标检测的需求。
- 冻结部分层：在训练过程中，可以冻结模型的部分层，使其参数保持不变，只训练新增的层或最后几层。这可以防止对预训练特征的过度调整。
- 训练和评估：在目标数据集上训练微调后的模型，并评估其性能。
示例
在代码中，可以看到如何使用预训练模型进行微调：
# 加载预训练模型
if args.resume:
logging.info(f"Resuming from the model {args.resume}")
net.load(args.resume)
elif args.base_net:
logging.info(f"Init from base net {args.base_net}")
net.init_from_base_net(args.base_net)
elif args.pretrained_ssd:
logging.info(f"Init from pretrained SSD {args.pretrained_ssd}")
if not os.path.exists(args.pretrained_ssd) and args.pretrained_ssd == DEFAULT_PRETRAINED_MODEL:
os.system(f"wget --quiet --show-progress --progress=bar:force:noscroll --no-check-certificate https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O {DEFAULT_PRETRAINED_MODEL}")
net.init_from_pretrained_ssd(args.pretrained_ssd)
args.resume：从之前的训练检查点恢复训练。args.base_net：使用预训练的基础网络初始化模型。args.pretrained_ssd：使用预训练的SSD模型初始化。
通过这些步骤，可以有效利用预训练模型的优势，加快训练过程，提高模型性能，并减少过拟合风险。
预训练模型还可以定义哪些
预训练模型的定义和选择取决于具体任务和使用场景。以下是一些关键点：
可以定义的预训练模型
- 通用图像分类模型：如在ImageNet数据集上训练的ResNet、VGG、Inception等，这些模型可以用于图像分类任务或作为特征提取器用于其他视觉任务。
- 目标检测模型：如在COCO数据集上训练的Faster R-CNN、SSD、YOLO等，这些模型可以用于目标检测任务。
- 语义分割模型：如在PASCAL VOC或Cityscapes数据集上训练的FCN、UNet、DeepLab等，这些模型可以用于语义分割任务。
- 自然语言处理模型：如在大规模语料库上训练的BERT、GPT、RoBERTa等，这些模型可以用于文本分类、生成、问答等任务。
- 其他领域特定模型：如人脸识别、视频处理、语音识别等领域的预训练模型。
如何选择预训练模型
- 任务匹配：选择在与当前任务相似的数据集上预训练的模型。例如，使用在ImageNet上预训练的模型进行图像分类或使用在COCO上预训练的模型进行目标检测。
- 模型架构：根据具体需求选择合适的模型架构。例如，轻量级任务可以选择MobileNet，精度要求高的任务可以选择ResNet或Inception。
- 社区支持：选择有良好社区支持和文档的预训练模型，可以获得更多帮助和资源。
- 可用性：选择容易获得和加载的预训练模型，有些模型可能在特定框架（如PyTorch、TensorFlow）中有官方实现和预训练权重。
不定义预训练模型行不行
可以不使用预训练模型，但需要考虑以下几点：
- 训练时间：从头开始训练模型通常需要更多的时间和计算资源，因为模型参数需要从随机初始化开始进行大量优化。
- 数据需求：从头训练模型通常需要大量的数据，以确保模型能够学到有效的特征。如果数据集较小，可能无法训练出高性能的模型。
- 性能：预训练模型由于在大规模数据上已经训练过，通常具有较好的初始性能和泛化能力。从头训练的模型可能需要更长时间才能达到类似的性能，甚至可能无法达到相同的水平。
示例代码
以下是使用预训练模型与不使用预训练模型的对比：
使用预训练模型：
# 使用预训练的SSD模型初始化
pretrained_model_path = 'models/mobilenet-v1-ssd-mp-0_675.pth'
net = create_mobilenetv1_ssd(num_classes)
net.init_from_pretrained_ssd(pretrained_model_path)
不使用预训练模型（从头训练）：
# 从头初始化模型（没有预训练权重）
net = create_mobilenetv1_ssd(num_classes)
在之前训练的模型基础上继续训练
在您的情况中，如果您想要在之前训练的模型基础上继续训练（即利用已经训练好的权重继续训练以增加新的样本数据的学习），应该使用 --resume 参数来加载您的模型。这个参数允许您从上次训练结束时的状态继续训练，包括模型权重和优化器状态。
参数选择：
-
--resume: 用于继续训练已有模型。例如，如果您已经训练了模型并且有一个检查点文件mb1-ssd-Epoch-886-Loss-0.1785774201154709.pth，则应使用：--resume mb1-ssd-Epoch-886-Loss-0.1785774201154709.pth这会从该检查点继续训练，包括学习率等状态。
-
--base-net: 用于加载一个预训练的基网络模型。如果您正在从头开始训练一个全新的模型并且只想加载网络的基础部分（通常是除了分类器之外的所有层），那么使用这个参数。 -
--pretrained-ssd: 用于加载一个完整的预训练 SSD 模型，这通常包括所有层的预训练权重。这适用于初始训练，但您希望从一个已经针对某任务训练好的模型开始。
是否使用 mobilenet-v1-ssd-mp-0_675.pth：
- 如果您使用
--resume参数继续之前的训练，那么不需要再次使用mobilenet-v1-ssd-mp-0_675.pth文件。因为--resume已经包括了从之前训练中继承的所有权重和训练状态，没有必要重新从预训练模型开始。
示例命令：
假设您已经有了一个训练状态的检查点，并且想要在此基础上继续训练：
python train_ssd.py --dataset-type=voc --data=data/phone/ --model-dir=models/phone_updated/ --resolution=512 --num-workers=2 --batch-size=4 --lr=0.01 --epochs=1000 --resume=models/mb1-ssd-Epoch-886-Loss-0.1785774201154709.pth
这条命令将加载您之前训练的模型，并在此基础上继续训练，不需要再次加载原始的 mobilenet-v1-ssd-mp-0_675.pth 模型。
通过这种方式，您可以确保模型在现有基础上继续学习新的样本，而不是重新开始训练过程。这通常有助于节省时间并改善模型针对特定任务的性能。
总结
预训练模型可以显著加速训练过程、提高模型性能并减少过拟合风险。选择预训练模型时，需要根据任务类型、模型架构、社区支持和可用性等因素进行综合考虑。在某些情况下，可以选择不使用预训练模型，但需要更多的训练时间和数据，且性能可能不如使用预训练模型。
三、增加对小物体的识别能力
要增加对小物体的识别能力，可以从以下几个方面调整训练参数和模型配置：
1. 调整Anchor Box（锚框）的尺寸
对于SSD（Single Shot MultiBox Detector）模型，Anchor Box的大小和比例是影响小物体检测的重要因素。你可以增加更多的小尺寸的Anchor Box，确保模型能够检测到更小的物体。
在 jetson-inference 的实现中，锚框的大小和比例通常在 config 文件中定义。以 mobilenetv1_ssd_config 为例，以下是如何定义和修改这些参数的方法。
修改mobilenetv1_ssd_config.py文件，首先，找到 mobilenetv1_ssd_config.py 文件，并进行修改。通常，这个文件位于 vision/ssd/config 目录下。
# vision/ssd/config/mobilenetv1_ssd_config.py
import numpy as np
from itertools import product
class Config:
def __init__(self):
self.image_size = 300 # 默认图像大小
self.image_mean = np.array([127, 127, 127]) # 图像均值
self.image_std = 128.0 # 图像标准差
self.iou_threshold = 0.5
self.center_variance = 0.1
self.size_variance = 0.2
# 用于计算先验框的参数
self.feature_maps = [19, 10, 5, 3, 2, 1]
self.min_sizes = [30, 60, 111, 162, 213, 264]
self.max_sizes = [60, 111, 162, 213, 264, 315]
self.aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
self.steps = [16, 32, 64, 100, 150, 300]
self.clip = True
self.priors = self._generate_priors()
def _generate_priors(self):
"""生成每个特征图层的先验框"""
priors = []
for k, f in enumerate(self.feature_maps):
for i, j in product(range(f), repeat=2):
f_k = self.image_size / self.steps[k]
cx = (j + 0.5) / f_k
cy = (i + 0.5) / f_k
s_k = self.min_sizes[k] / self.image_size
priors.append([cx, cy, s_k, s_k])
s_k_prime = np.sqrt(s_k * (self.max_sizes[k] / self.image_size))
priors.append([cx, cy, s_k_prime, s_k_prime])
for ar in self.aspect_ratios[k]:
priors.append([cx, cy, s_k * np.sqrt(ar), s_k / np.sqrt(ar)])
priors.append([cx, cy, s_k / np.sqrt(ar), s_k * np.sqrt(ar)])
priors = np.clip(priors, 0.0, 1.0)
return np.array(priors)
config = Config()
在训练代码中使用修改后的配置
在你的训练脚本中，确保你正确地导入和使用修改后的配置：
from vision.ssd.config import mobilenetv1_ssd_config
# 根据需求修改配置参数
config = mobilenetv1_ssd_config.Config()
config.image_size = 512 # 设置图像大小
config.min_sizes = [21, 45, 99, 153, 207, 261]
config.max_sizes = [45, 99, 153, 207, 261, 315]
# 创建数据增强
train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std)
target_transform = MatchPrior(config.priors, config.center_variance, config.size_variance, 0.5)
# 剩下的训练代码保持不变
通过以上方法，你可以修改 min_sizes 和 max_sizes 参数，以更好地检测小物体。在调整这些参数后，确保重新生成和检查先验框，以验证它们是否符合预期。这样可以提高模型对小物体的检测能力。
实例一
图像分辨率为512512，手机占的大小为300140像素，手机上的文字大小为40*10像素，要求能根据不同文字及手机颜色区别不同手机
为了使SSD模型能够检测到图像分辨率为512x512时大小为300x140像素的手机和40x10像素的文字，我们需要调整SSDSpec和SSDBoxSizes的配置。以下是详细步骤：
1. 确定目标物体的尺寸比例
- 手机宽度比例：140 / 512 ≈ 0.2734375
- 手机高度比例：300 / 512 ≈ 0.5859375
- 文字宽度比例：10 / 512 ≈ 0.01953125
- 文字高度比例：40 / 512 ≈ 0.078125
2. 调整SSDSpec和SSDBoxSizes
我们将根据上述比例调整锚框的大小，使得锚框能适应不同尺寸的物体（手机和文字）。
说明
-
调整锚框大小：
- 第一个特征图层的锚框大小设为40到80，用于检测小文字（40x10像素）。
- 第二个特征图层的锚框大小设为80到160，用于检测较小的物体。
- 第三个特征图层的锚框大小设为160到240，用于检测中等大小的物体。
- 第四个特征图层的锚框大小设为240到320，用于检测手机（300x140像素）。
- 第五和第六个特征图层的锚框大小分别设为320到400和400到480，用于检测更大的物体。
-
set_image_size函数：- 函数中
min_ratio和max_ratio的默认值保持为20到90，可以根据需要进一步调整。
- 函数中
总结
通过调整特征图层的锚框大小，使其适应手机（300x140像素）和文字（40x10像素）的实际尺寸，可以提高模型对这些物体的检测精度。你可以根据具体需求进一步微调这些参数以获得最佳效果。
由官方默认源码
import numpy as np
from vision.utils.box_utils import SSDSpec, SSDBoxSizes, generate_ssd_priors
image_size = 300
image_mean = np.array([127, 127, 127]) # RGB layout
image_std = 128.0
iou_threshold = 0.45
center_variance = 0.1
size_variance = 0.2
specs = [
SSDSpec(19, 16, SSDBoxSizes(60, 105), [2, 3]),
SSDSpec(10, 32, SSDBoxSizes(105, 150), [2, 3]),
SSDSpec(5, 64, SSDBoxSizes(150, 195), [2, 3]),
SSDSpec(3, 100, SSDBoxSizes(195, 240), [2, 3]),
SSDSpec(2, 150, SSDBoxSizes(240, 285), [2, 3]),
SSDSpec(1, 300, SSDBoxSizes(285, 330), [2, 3])
]
priors = generate_ssd_priors(specs, image_size)
def set_image_size(size=300, min_ratio=20, max_ratio=90):
global image_size
global specs
global priors
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
import torch
import math
import logging
image_size = size
ssd = create_mobilenetv1_ssd(num_classes=3) # TODO does num_classes matter here?
x = torch.randn(1, 3, image_size, image_size)
feature_maps = ssd(x, get_feature_map_size=True)
steps = [
math.ceil(image_size * 1.0 / feature_map) for feature_map in feature_maps
]
step = int(math.floor((max_ratio - min_ratio) / (len(feature_maps) - 2)))
min_sizes = []
max_sizes = []
for ratio in range(min_ratio, max_ratio + 1, step):
min_sizes.append(image_size * ratio / 100.0)
max_sizes.append(image_size * (ratio + step) / 100.0)
min_sizes = [image_size * (min_ratio / 2) / 100.0] + min_sizes
max_sizes = [image_size * min_ratio / 100.0] + max_sizes
# this update logic makes different boxes than the original for 300x300 (but better for power-of-two)
# for backwards-compatibility, keep the default 300x300 config if that's what's being called for
if image_size != 300:
specs = []
for i in range(len(feature_maps)):
specs.append( SSDSpec(feature_maps[i], steps[i], SSDBoxSizes(min_sizes[i], max_sizes[i]), [2, 3]) ) # ssd-mobilenet-* aspect ratio is [2,3]
logging.info(f'model resolution {image_size}x{image_size}')
for spec in specs:
logging.info(str(spec))
priors = generate_ssd_priors(specs, image_size)
#print(' ')
#print('SSD-Mobilenet-v1 priors:')
#print(priors.shape)
#print(priors)
#print(' ')
#import torch
#torch.save(priors, 'mb1-ssd-priors.pt')
#np.savetxt('mb1-ssd-priors.txt', priors.numpy())
改为
import numpy as np
from vision.utils.box_utils import SSDSpec, SSDBoxSizes, generate_ssd_priors
image_size = 1024
image_mean = np.array([127, 127, 127]) # RGB layout
image_std = 128.0
iou_threshold = 0.45
center_variance = 0.1
size_variance = 0.2
# 调整后的锚框大小以适应手机和文字尺寸
specs = [
SSDSpec(64, 16, SSDBoxSizes(40, 80), [2, 3]), # 用于检测小文字
SSDSpec(32, 32, SSDBoxSizes(80, 160), [2, 3]), # 用于检测较小的物体（包括文字）
SSDSpec(16, 64, SSDBoxSizes(160, 320), [2, 3]), # 用于检测中等大小物体
SSDSpec(8, 128, SSDBoxSizes(320, 480), [2, 3]), # 用于检测中等大小的手机
SSDSpec(4, 256, SSDBoxSizes(480, 640), [2, 3]), # 用于检测较大的手机
SSDSpec(2, 512, SSDBoxSizes(640, 800), [2, 3]), # 用于检测更大的手机
SSDSpec(1, 1024, SSDBoxSizes(800, 960), [2, 3]) # 用于检测最大的物体
]
priors = generate_ssd_priors(specs, image_size)
def set_image_size(size=1024, min_ratio=10, max_ratio=90):
global image_size
global specs
global priors
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
import torch
import math
import logging
image_size = size
ssd = create_mobilenetv1_ssd(num_classes=3) # 根据需要修改num_classes
x = torch.randn(1, 3, image_size, image_size)
feature_maps = ssd(x, get_feature_map_size=True)
steps = [
math.ceil(image_size * 1.0 / feature_map) for feature_map in feature_maps
]
step = int(math.floor((max_ratio - min_ratio) / (len(feature_maps) - 2)))
min_sizes = []
max_sizes = []
for ratio in range(min_ratio, max_ratio + 1, step):
min_sizes.append(image_size * ratio / 100.0)
max_sizes.append(image_size * (ratio + step) / 100.0)
min_sizes = [image_size * (min_ratio / 2) / 100.0] + min_sizes
max_sizes = [image_size * min_ratio / 100.0] + max_sizes
# this update logic makes different boxes than the original for 300x300 (but better for power-of-two)
# for backwards-compatibility, keep the default 300x300 config if that's what's being called for
if image_size != 300:
specs = []
for i in range(len(feature_maps)):
specs.append(SSDSpec(feature_maps[i], steps[i], SSDBoxSizes(min_sizes[i], max_sizes[i]), [2, 3])) # ssd-mobilenet-* aspect ratio is [2,3]
logging.info(f'model resolution {image_size}x{image_size}')
for spec in specs:
logging.info(str(spec))
priors = generate_ssd_priors(specs, image_size)
#print(' ')
#print('SSD-Mobilenet-v1 priors:')
#print(priors.shape)
#print(priors)
#print(' ')
#import torch
#torch.save(priors, 'mb1-ssd-priors.pt')
#np.savetxt('mb1-ssd-priors.txt', priors.numpy())
实例二
在设计SSD (Single Shot Multibox Detector) 的锚框（Anchor Box）配置时，如specs中所列的每个SSDSpec，每个值都有其特定目的和对目标检测性能的影响。以下是如何确定这些值及其对于识别手机模型的帮助的解释：
1. 特征图尺寸（例如：64, 32, 16, 8, 4, 2, 1）
- 目的：不同的特征图尺寸能够捕捉不同尺寸的目标。较大的特征图（如64x64）更适合检测图像中的小目标，因为它们在图像中保留了更多细节。较小的特征图（如1x1）则用于检测图像中的大目标。
- 确定方法：特征图的尺寸由输入图像尺寸逐级降采样确定，每层采样通常降低特征图的空间维度（高和宽）。这个决策基于目标的预期大小以及图像的总尺寸。
- 对手机模型识别的帮助：确保模型可以在不同尺度上有效检测手机屏幕上的不同元素（如图标、文字），从大图标到细小文字。
2. 步长（例如：8, 16, 32, 64, 128, 256, 512）
- 目的：步长决定了特征图中锚框的空间分布，较小的步长在原图中表示锚框更密集，适合捕捉小目标；较大的步长适用于大目标。
- 确定方法：步长通常与特征图的尺寸成反比。例如，原始图像大小为512x512，特征图尺寸为64，则步长为512/64=8。
- 对手机模型识别的帮助：合适的步长设置可以确保在整个图像上均匀地分布锚框，增加检测小细节（如文本）的概率。
3. 锚框尺寸（例如：SSDBoxSizes(20, 50)）
- 目的：定义每个特征图层使用的锚框的物理尺寸范围，以适应不同大小的目标。
- 确定方法：锚框尺寸通常根据目标在图像中的实际大小进行选择。例如，如果目标大小大约在20到50像素之间，那么使用这个尺寸范围的锚框会更有效。
- 对手机模型识别的帮助：通过精确匹配目标的大小，提高模型对手机屏幕上特定元素（如各种大小的图标和文本）的识别准确度。
4. 长宽比（例如：[2, 3]）
- 目的：不同的长宽比可以更好地适应目标的形状，如长条形状或方形。
- 确定方法：长宽比的选择基于目标的典型形状。例如，文本可能更适合长形锚框，而图标可能适合接近正方形的锚框。
- 对手机模型识别的帮助：确保对于不规则形状（如长条的文本或方形的图标）的有效检测。
实例三
为了识别图中的手机模型，并能区分不同手机模型上的文字，我们需要设计合理的锚框设置，使其能够覆盖手机模型的整体区域以及文字区域。以下是针对这种情况的优化锚框设置：
优化锚框设置
- 小锚框：用于识别小尺寸目标，如文字。
- 中等锚框：用于识别中等大小目标，如图标和文字区域。
- 大锚框：用于识别大目标，如整个手机模型。
示例设置
如果手机模型的主要差异在于文字，因此我们需要确保有适合文字的锚框。此外，还需要适合手机模型整体轮廓的锚框。以下是一个包含4个级别特征图的SSD模型设置示例：
specs = [
SSDSpec(64, 8, SSDBoxSizes(20, 40), [1, 2, 1/2]), # 特征图尺寸64x64，适用于小尺寸目标如文字
SSDSpec(32, 16, SSDBoxSizes(40, 80), [1, 2, 1/2]), # 特征图尺寸32x32，适用于中等尺寸目标
SSDSpec(16, 32, SSDBoxSizes(80, 160), [1, 2]), # 特征图尺寸16x16，适用于中到大尺寸目标
SSDSpec(8, 64, SSDBoxSizes(160, 320), [2]) # 特征图尺寸8x8，适用于大尺寸目标，如手机模型
]
解释
- 小锚框：用于检测小目标如文字，尺寸为20x20和40x40，长宽比为1:1, 2:1, 1:2。
- 中等锚框：用于检测中等大小目标，尺寸为40x40和80x80，长宽比为1:1, 2:1, 1:2。
- 中大锚框：用于检测中到大尺寸目标，尺寸为80x80和160x160，长宽比为1:1, 2:1。
- 大锚框：用于检测大目标如手机模型，尺寸为160x320，长宽比为2:1。
总结
通过调整锚框的尺寸和长宽比，使其适合识别手机模型的整体区域和文字区域，可以提高检测精度。结合目标检测和OCR技术，可以实现对不同手机模型上的文字的精确识别。优化后的锚框设置和多层特征图确保模型在不同尺度上有效工作，满足实际应用需求。
每个特征图生成的锚框数量和比例可以设置
锚框越多，识别速度越慢
三个锚框+每种锚框两种比例，总共6个锚框
specs = [
SSDSpec(64, 8, SSDBoxSizes(20, 35, 50), [2, 3]), # 特征图尺寸64x64，三个不同尺寸的锚框
SSDSpec(32, 16, SSDBoxSizes(50, 70, 90), [2, 3]), # 特征图尺寸32x32，三个不同尺寸的锚框
SSDSpec(16, 32, SSDBoxSizes(90, 110, 130), [2, 3]), # 特征图尺寸16x16，三个不同尺寸的锚框
SSDSpec(8, 64, SSDBoxSizes(130, 150, 170), [2, 3]), # 特征图尺寸8x8，三个不同尺寸的锚框
SSDSpec(4, 128, SSDBoxSizes(170, 190, 210), [2, 3]),# 特征图尺寸4x4，三个不同尺寸的锚框
SSDSpec(2, 256, SSDBoxSizes(210, 230, 250), [2, 3]),# 特征图尺寸2x2，三个不同尺寸的锚框
SSDSpec(1, 512, SSDBoxSizes(250, 275, 300), [2, 3]) # 特征图尺寸1x1，三个不同尺寸的锚框
]
三个锚框+每种锚框三种比例，总共9个锚框
specs = [
SSDSpec(64, 8, SSDBoxSizes(20, 35, 50), [1, 2, 3]), # 特征图尺寸64x64，三个不同尺寸和长宽比的锚框
SSDSpec(32, 16, SSDBoxSizes(50, 70, 90), [1, 2, 3]), # 特征图尺寸32x32，三个不同尺寸和长宽比的锚框
SSDSpec(16, 32, SSDBoxSizes(90, 110, 130), [1, 2, 3]), # 特征图尺寸16x16，三个不同尺寸和长宽比的锚框
SSDSpec(8, 64, SSDBoxSizes(130, 150, 170), [1, 2, 3]), # 特征图尺寸8x8，三个不同尺寸和长宽比的锚框
SSDSpec(4, 128, SSDBoxSizes(170, 190, 210), [1, 2, 3]),# 特征图尺寸4x4，三个不同尺寸和长宽比的锚框
SSDSpec(2, 256, SSDBoxSizes(210, 230, 250), [1, 2, 3]),# 特征图尺寸2x2，三个不同尺寸和长宽比的锚框
SSDSpec(1, 512, SSDBoxSizes(250, 275, 300), [1, 2, 3]) # 特征图尺寸1x1，三个不同尺寸和长宽比的锚框
]
总结
这些参数的综合考量和精心设计使得SSD模型可以在不同尺度上有效地识别和定位图像中的目标，对于识别手机屏幕的多样内容尤为重要，包括能够细致地区分小字和图标等细节。这种多尺度和多形状的设计方法是SSD模型强大性能的关键所在。
2. 数据增强
增加数据增强方法，使得模型能更好地学习各种尺度下的物体。例如：
- 随机裁剪（Random Cropping）
- 随机缩放（Random Scaling）
- 随机翻转（Random Flipping）
3. 增加特征金字塔网络（Feature Pyramid Networks, FPN）
FPN可以帮助模型从不同尺度的特征图中提取信息，提高对小物体的检测能力。
4. 调整学习率和批量大小
合适的学习率和批量大小有助于更好地训练模型。
5. 使用预训练模型
使用在大规模数据集上预训练的模型作为初始模型，然后进行微调。
配置示例
以下是一些配置示例，假设你在使用SSD模型进行训练：
import torch
from vision.ssd.config import mobilenetv1_ssd_config
config = mobilenetv1_ssd_config
# 调整Anchor Box的尺寸和比例
config.image_size = 512 # 增加输入图像的分辨率，便于检测小物体
config.min_sizes = [21, 45, 99, 153, 207, 261]
config.max_sizes = [45, 99, 153, 207, 261, 315]
config.aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
# 数据增强
train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std)
target_transform = MatchPrior(config.priors, config.center_variance, config.size_variance, 0.5)
# 调整学习率和批量大小
learning_rate = 0.005 # 可以适当调小学习率
batch_size = 16 # 可以增加批量大小，提高训练稳定性
# 使用预训练模型
pretrained_ssd = 'models/mobilenet-v1-ssd-mp-0_675.pth'
# 训练参数配置
args = Namespace(dataset_type='voc', datasets=['data/phone_1024/'], balance_data=False,
net='mb1-ssd', resolution=1024, freeze_base_net=False, freeze_net=False,
mb2_width_mult=1.0, base_net=None, pretrained_ssd=pretrained_ssd, resume=None,
lr=learning_rate, momentum=0.9, weight_decay=0.0005, gamma=0.1, base_net_lr=0.001,
extra_layers_lr=None, scheduler='cosine', milestones='80,100', t_max=100,
batch_size=batch_size, num_epochs=2000, num_workers=2, validation_epochs=1,
validation_mean_ap=False, debug_steps=10, use_cuda=True,
checkpoint_folder='models/phone_1024px_bs4_epochs2000/', log_level='info')
# 训练
train_ssd(args)
总结
- 调整Anchor Box的大小和比例：更适合小物体的检测。
- 增加数据增强：提高模型的鲁棒性。
- 增加特征金字塔网络（FPN）：增强不同尺度特征的提取。
- 调整学习率和批量大小：找到最佳训练参数。
- 使用预训练模型：利用预训练模型进行微调。
这些调整可以提高模型对小物体的检测能力。请根据你的具体情况和数据集进行调整和优化。
扫一扫
99
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
