項(xiàng)目實(shí)踐 | 從零開(kāi)始邊緣部署輕量化人臉檢測(cè)模型——訓(xùn)練篇
1簡(jiǎn)介
該模型是針對(duì)邊緣計(jì)算設(shè)備設(shè)計(jì)的輕量人臉檢測(cè)模型。
在模型大小上，默認(rèn)FP32精度下（.pth）文件大小為 1.04~1.1MB，推理框架int8量化后大小為 300KB 左右。 在模型計(jì)算量上，320x240的輸入分辨率下 90~109 MFlops左右。 模型有兩個(gè)版本，version-slim(主干精簡(jiǎn)速度略快)，version-RFB(加入了修改后的RFB模塊，精度更高)。 提供320x240、640x480不同輸入分辨率下使用widerface訓(xùn)練的預(yù)訓(xùn)練模型，更好的工作于不同的應(yīng)用場(chǎng)景。
2數(shù)據(jù)處理
2.1 輸入尺寸的選擇
由于涉及實(shí)際部署時(shí)的推理速度，因此模型輸入尺寸的選擇也是一個(gè)很重要的話題。
在作者的原github中，也提到了一點(diǎn)，如果在實(shí)際部署的場(chǎng)景中大多數(shù)情況為中近距離、人臉大同時(shí)人臉的數(shù)量也比較少的時(shí)候，則可以采用的輸入尺寸；
如果在實(shí)際部署的場(chǎng)景中大多數(shù)情況為中遠(yuǎn)距離、人臉小同時(shí)人臉的數(shù)量也比較多的時(shí)候，則可以采用或者的輸入尺寸；
這里由于使用的是EAIDK310進(jìn)行部署測(cè)試，邊緣性能不是很好，因此選擇原作者推薦的最小尺寸進(jìn)行訓(xùn)練和部署測(cè)試。
注意：過(guò)小的輸入分辨率雖然會(huì)明顯加快推理速度，但是會(huì)大幅降低小人臉的召回率。
2.2 數(shù)據(jù)篩選
由于widerface官網(wǎng)數(shù)據(jù)集中有比較多的低于10像素的人臉照片，因此在這里選擇剔除這些像素長(zhǎng)寬低于10個(gè)pixel的照片；
這樣做的原因是：不清楚的人臉，不太利于高效模型的收斂，所以需要進(jìn)行過(guò)濾訓(xùn)練。
3SSD網(wǎng)絡(luò)結(jié)構(gòu)
SSD是一個(gè)端到端的模型，所有的檢測(cè)過(guò)程和識(shí)別過(guò)程都是在同一個(gè)網(wǎng)絡(luò)中進(jìn)行的；同時(shí)SSD借鑒了Faster R-CNN的Anchor機(jī)制的想法，這樣就像相當(dāng)于在基于回歸的的檢測(cè)過(guò)程中結(jié)合了區(qū)域的思想，可以使得檢測(cè)效果較定制化邊界框的YOLO v1有比較好的提升。
SSD較傳統(tǒng)的檢測(cè)方法使用頂層特征圖的方法選擇了使用多尺度特征圖，因?yàn)樵诒容^淺的特征圖中可以對(duì)于小目標(biāo)有比較好的表達(dá)，隨著特征圖的深入，網(wǎng)絡(luò)對(duì)于比較大特征也有了比較好表達(dá)能力，故SSD選擇使用多尺度特征圖可以很好的兼顧大目標(biāo)和小目標(biāo)。
SSD模型結(jié)構(gòu)如下：
這里關(guān)于SSD不進(jìn)行更多的闡述，想了解的小伙伴可以掃描下方的二維碼查看（是小編在CSDN的記錄，非常詳細(xì)?。。。?/p>
整個(gè)項(xiàng)目模型搭建如下：
# 網(wǎng)絡(luò)的主題結(jié)構(gòu)為SSD模型
class SSD(nn.Module):
def __init__(self, num_classes: int, base_net: nn.ModuleList, source_layer_indexes: List[int],
extras: nn.ModuleList, classification_headers: nn.ModuleList,
regression_headers: nn.ModuleList, is_test=False, config=None, device=None):
"""Compose a SSD model using the given components.
"""
super(SSD, self).__init__()
self.num_classes = num_classes
self.base_net = base_net
self.source_layer_indexes = source_layer_indexes
self.extras = extras
self.classification_headers = classification_headers
self.regression_headers = regression_headers
self.is_test = is_test
self.config = config
# register layers in source_layer_indexes by adding them to a module list
self.source_layer_add_ons = nn.ModuleList([t[1] for t in source_layer_indexes
if isinstance(t, tuple) and not isinstance(t, GraphPath)])
if device:
self.device = device
else:
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if is_test:
self.config = config
self.priors = config.priors.to(self.device)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
confidences = []
locations = []
start_layer_index = 0
header_index = 0
end_layer_index = 0
for end_layer_index in self.source_layer_indexes:
if isinstance(end_layer_index, GraphPath):
path = end_layer_index
end_layer_index = end_layer_index.s0
added_layer = None
elif isinstance(end_layer_index, tuple):
added_layer = end_layer_index[1]
end_layer_index = end_layer_index[0]
path = None
else:
added_layer = None
path = None
for layer in self.base_net[start_layer_index: end_layer_index]:
x = layer(x)
if added_layer:
y = added_layer(x)
else:
y = x
if path:
sub = getattr(self.base_net[end_layer_index], path.name)
for layer in sub[:path.s1]:
x = layer(x)
y = x
for layer in sub[path.s1:]:
x = layer(x)
end_layer_index += 1
start_layer_index = end_layer_index
confidence, location = self.compute_header(header_index, y)
header_index += 1
confidences.append(confidence)
locations.append(location)
for layer in self.base_net[end_layer_index:]:
x = layer(x)
for layer in self.extras:
x = layer(x)
confidence, location = self.compute_header(header_index, x)
header_index += 1
confidences.append(confidence)
locations.append(location)
confidences = torch.cat(confidences, 1)
locations = torch.cat(locations, 1)
if self.is_test:
confidences = F.softmax(confidences, dim=2)
boxes = box_utils.convert_locations_to_boxes(
locations, self.priors, self.config.center_variance, self.config.size_variance
)
boxes = box_utils.center_form_to_corner_form(boxes)
return confidences, boxes
else:
return confidences, locations
def compute_header(self, i, x):
confidence = self.classification_headers[i](x)
confidence = confidence.permute(0, 2, 3, 1).contiguous()
confidence = confidence.view(confidence.size(0), -1, self.num_classes)
location = self.regression_headers[i](x)
location = location.permute(0, 2, 3, 1).contiguous()
location = location.view(location.size(0), -1, 4)
return confidence, location
def init_from_base_net(self, model):
self.base_net.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage), strict=True)
self.source_layer_add_ons.apply(_xavier_init_)
self.extras.apply(_xavier_init_)
self.classification_headers.apply(_xavier_init_)
self.regression_headers.apply(_xavier_init_)
def init_from_pretrained_ssd(self, model):
state_dict = torch.load(model, map_location=lambda storage, loc: storage)
state_dict = {k: v for k, v in state_dict.items() if not (k.startswith("classification_headers") or k.startswith("regression_headers"))}
model_dict = self.state_dict()
model_dict.update(state_dict)
self.load_state_dict(model_dict)
self.classification_headers.apply(_xavier_init_)
self.regression_headers.apply(_xavier_init_)
def init(self):
self.base_net.apply(_xavier_init_)
self.source_layer_add_ons.apply(_xavier_init_)
self.extras.apply(_xavier_init_)
self.classification_headers.apply(_xavier_init_)
self.regression_headers.apply(_xavier_init_)
def load(self, model):
self.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage))
def save(self, model_path):
torch.save(self.state_dict(), model_path)
4損失函數(shù)
損失函數(shù)作者選擇使用的依舊是SSD的Smooth L1 Loss以及Cross Entropy Loss，其中Smooth L1 Loss用于邊界框的回歸，而Cross Entropy Loss則用于分類(lèi)。
具體pytorch實(shí)現(xiàn)如下：
class MultiboxLoss(nn.Module):
def __init__(self, priors, neg_pos_ratio,
center_variance, size_variance, device):
"""Implement SSD Multibox Loss.
Basically, Multibox loss combines classification loss
and Smooth L1 regression loss.
"""
super(MultiboxLoss, self).__init__()
self.neg_pos_ratio = neg_pos_ratio
self.center_variance = center_variance
self.size_variance = size_variance
self.priors = priors
self.priors.to(device)
def forward(self, confidence, predicted_locations, labels, gt_locations):
"""Compute classification loss and smooth l1 loss.
Args:
confidence (batch_size, num_priors, num_classes): class predictions.
locations (batch_size, num_priors, 4): predicted locations.
labels (batch_size, num_priors): real labels of all the priors.
boxes (batch_size, num_priors, 4): real boxes corresponding all the priors.
"""
num_classes = confidence.size(2)
with torch.no_grad():
# derived from cross_entropy=sum(log(p))
loss = -F.log_softmax(confidence, dim=2)[:, :, 0]
mask = box_utils.hard_negative_mining(loss, labels, self.neg_pos_ratio)
confidence = confidence[mask, :]
# 分類(lèi)損失函數(shù)
classification_loss = F.cross_entropy(confidence.reshape(-1, num_classes), labels[mask], reduction='sum')
pos_mask = labels > 0
predicted_locations = predicted_locations[pos_mask, :].reshape(-1, 4)
gt_locations = gt_locations[pos_mask, :].reshape(-1, 4)
# 邊界框回歸損失函數(shù)
smooth_l1_loss = F.smooth_l1_loss(predicted_locations, gt_locations, reduction='sum') # smooth_l1_loss
# smooth_l1_loss = F.mse_loss(predicted_locations, gt_locations, reduction='sum') #l2 loss
num_pos = gt_locations.size(0)
return smooth_l1_loss / num_pos, classification_loss / num_pos
5結(jié)果預(yù)測(cè)
輸入為：
輸出為：
輸入為：
輸出為：
6模型轉(zhuǎn)換
由于部署使用的是Tengine邊緣推理框架，由于pytorch輸出的模型無(wú)法直接轉(zhuǎn)換到tmfile模型下，因此還是選擇使用onnx中間件的形式進(jìn)行過(guò)度，具體實(shí)現(xiàn)代碼如下：
model_path = "models/pretrained/version-RFB-320.pth"
net = create_Mb_Tiny_RFB_fd(len(class_names), is_test=True)
net.load(model_path)
net.eval()
net.to("cuda")
model_name = model_path.split("/")[-1].split(".")[0]
model_path = f"models/onnx/{model_name}.onnx"
dummy_input = torch.randn(1, 3, 240, 320).to("cuda")
# dummy_input = torch.randn(1, 3, 480, 640).to("cuda") #if input size is 640*480
torch.onnx.export(net, dummy_input, model_path, verbose=False, input_names=['input'], output_names=['scores', 'boxes'])
得到onnx模型后便可以進(jìn)行Tengine模型的轉(zhuǎn)換和部署，該部分將在下一篇文章繼續(xù)討論。
7參考
[1].https://github.com/Linzaer/Ultra-Light-Fast-Generic-Face-Detector-1MB
[2].https://github.com/onnx/onnx
8推薦閱讀
Google新作 | 詳細(xì)解讀 Transformer那些有趣的特性（建議全文背誦）
極品Trick | 在ResNet與Transformer均適用的Skip Connection解讀
Transformer又一城 | Swin-Unet：首個(gè)純Transformer的醫(yī)學(xué)圖像分割模型解讀
輕量化卷積：TBC，不僅僅是參數(shù)共享組卷積，更具備跨通道建模
最快ViT | FaceBook提出LeViT，0.077ms的單圖處理速度卻擁有ResNet50的精度(文末附論文與源碼)
本文論文原文獲取方式，掃描下方二維碼
回復(fù)【UltraFace】即可獲取項(xiàng)目代碼
長(zhǎng)按掃描下方二維碼添加小助手。
可以一起討論遇到的問(wèn)題
聲明：轉(zhuǎn)載請(qǐng)說(shuō)明出處
掃描下方二維碼關(guān)注【集智書(shū)童】公眾號(hào)，獲取更多實(shí)踐項(xiàng)目源碼和論文解讀，非常期待你我的相遇，讓我們以夢(mèng)為馬，砥礪前行！
