睿智的目标檢測65——Pytorch搭建DETR目标檢測平台
- 學習前言
- 源碼下載
- DETR實現思路
- 一、整體結構解析
- 二、網絡結構解析
- 1、主幹網絡Backbone介紹
- a、什麽是殘差網絡
- b、什麽是ResNet50模型
- c、位置編碼
- 2、編碼網絡Encoder網絡介紹
- a、Transformer Encoder的構建
- b、Self-attention結構解析
- 3、解碼網絡Decoder網絡介紹
- 4、預測頭的構建
- 三、預測結果的解碼
- 四、訓練部分
- 1、計算Loss所需内容
- 2、正樣本的匹配過程
- 3、計算Loss
- 訓練自己的DETR模型
- 一、數據集的準備
- 二、數據集的處理
- 三、開始網絡訓練
- 四、訓練結果預測
學習前言
基于Transformer的目标檢測一直沒弄，補上一下。
源碼下載
https://github.com/bubbliiiing/detr-pytorch
喜歡的可以點個star噢。
DETR實現思路
一、整體結構解析
在學習DETR之前，我們需要對DETR所做的工作有一定的了解，這有助于我們後面去了解網絡的細節。上面這幅圖是論文裏的Fig. 2，比較好的展示了整個DETR的工作原理。整個DETR可以分爲四個部分，分别是：backbone、encoder、decoder以及prediction heads。
backbone是DETR的主幹特征提取網絡，輸入的圖片首先會在主幹網絡裏面進行特征提取，提取到的特征可以被稱作特征層，是輸入圖片的特征集合。在主幹部分，我們獲取了一個特征層進行下一步網絡的構建，這一個特征層我稱它爲有效特征層。
encoder是Transformer的編碼網絡-特征加強，在主幹部分獲得的一個有效特征層會首先在高寬維度進行平鋪，成爲一個特征序列，然後會在這一部分繼續使用Self-Attension進行加強特征提取，獲得一個加強後的有效特征層。它屬于Transformer的編碼網絡，編碼的下一步是解碼。
decoder是Transformer的解碼網絡-特征查詢，在encoder部分獲得的一個加強後的有效特征層會在這一部分進行解碼，解碼需要使用到一個非常重要的可學習模塊，即上圖呈現的object queries。在decoder部分，我們使用一個可學習的查詢向量q對加強後的有效特征層進行查詢，獲得預測結果。
prediction heads是DETR的分類器與回歸器，其實就是對decoder獲得的預測結果進行全連接，兩次全連接分别代表種類和回歸參數。圖上畫了4個FFN，源碼中是2個FFN。
因此，整個DETR網絡所作的工作就是 特征提取-特征加強-特征查詢-預測結果。
二、網絡結構解析
1、主幹網絡Backbone介紹
DETR可以采用多種的主幹特征提取網絡，論文中用的是Resnet，本文以Resnet50網絡爲例子來給大家演示一下。
a、什麽是殘差網絡
Residual net(殘差網絡)：
将靠前若幹層的某一層數據輸出直接跳過多層引入到後面數據層的輸入部分。
意味着後面的特征層的内容會有一部分由其前面的某一層線性貢獻。
其結構如下：
深度殘差網絡的設計是爲了克服由于網絡深度加深而産生的學習效率變低與準确率無法有效提升的問題。
b、什麽是ResNet50模型
ResNet50有兩個基本的塊，分别名爲Conv Block和Identity Block，其中Conv Block輸入和輸出的維度是不一樣的，所以不能連續串聯，它的作用是改變網絡的維度；Identity Block輸入維度和輸出維度相同，可以串聯，它的作用是加深網絡的。
Conv Block的結構如下，由圖可以看出，Conv Block可以分爲兩個部分，左邊部分爲主幹部分，存在兩次卷積、标準化、激活函數和一次卷積、标準化；右邊部分爲殘差邊部分，存在一次卷積、标準化，由于殘差邊部分存在卷積，所以我們可以利用Conv Block改變輸出特征層的寬高和通道數：
Identity Block的結構如下，由圖可以看出，Identity Block可以分爲兩個部分，左邊部分爲主幹部分，存在兩次卷積、标準化、激活函數和一次卷積、标準化；右邊部分爲殘差邊部分，直接與輸出相接，由于殘差邊部分不存在卷積，所以Identity Block的輸入特征層和輸出特征層的shape是相同的，可用于加深網絡：
Conv Block和Identity Block都是殘差網絡結構。
總的網絡結構如下：
在DETR中，假設輸入爲[batch_size, 3, 800, 800]，此時輸出爲[batch_size, 2048, 25, 25]，代碼直接使用了torchvision庫中自帶的resnet，因此整個主幹實現代碼爲：
Pythonclass FrozenBatchNorm2d(torch.nn.Module): """ 凍結固定的BatchNorm2d。 """ def __init__(self, n): super(FrozenBatchNorm2d, self).__init__() self.register_buffer("weight", torch.ones(n)) self.register_buffer("bias", torch.zeros(n)) self.register_buffer("running_mean", torch.zeros(n)) self.register_buffer("running_var", torch.ones(n)) def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): num_batches_tracked_key = prefix + 'num_batches_tracked' if num_batches_tracked_key in state_dict: del state_dict[num_batches_tracked_key] super(FrozenBatchNorm2d, self)._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) def forward(self, x): w = self.weight.reshape(1, -1, 1, 1) b = self.bias.reshape(1, -1, 1, 1) rv = self.running_var.reshape(1, -1, 1, 1) rm = self.running_mean.reshape(1, -1, 1, 1) eps = 1e-5 scale = w * (rv + eps).rsqrt() bias = b - rm * scale return x * scale + bias class BackboneBase(nn.Module): """ 用于指定返回哪個層的輸出 這裏返回的是最後一層 """ def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool): super().__init__() for name, parameter in backbone.named_parameters(): if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: parameter.requires_grad_(False) if return_interm_layers: return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} else: return_layers = {'layer4': "0"} # 用于指定返回的層 self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) self.num_channels = num_channels def forward(self, tensor_list: NestedTensor): xs = self.body(tensor_list.tensors) out: Dict[str, NestedTensor] = {} for name, x in xs.items(): m = tensor_list.mask assert m is not None mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0] out[name] = NestedTensor(x, mask) return out class Backbone(BackboneBase): """ ResNet backbone with frozen BatchNorm. """ def __init__(self, name: str, train_backbone: bool, return_interm_layers: bool,dilation: bool): # 首先利用torchvision裏面的model創建一個backbone模型 backbone = getattr(torchvision.models, name)( replace_stride_with_dilation = [False, False, dilation], pretrained = is_main_process(), norm_layer = FrozenBatchNorm2d ) # 根據選擇的模型，獲得通道數 num_channels = 512 if name in ('resnet18', 'resnet34') else 2048 super().__init__(backbone, train_backbone, num_channels, return_interm_layers)c、位置編碼
除去利用主幹進行特征提取外，因爲要傳入Transformer進行特征提取與特征查詢，主幹獲得的特征還需要進行位置編碼。在圖片上不屬于backbone，但是在backbone.py裏實現的，所以一起簡單解析一下。
其實就是原Transformer的position embedding的思想，爲所有特征添加上位置信息，這樣網絡才有區分不同區域的能力。
DETR是爲resnet輸出的特征圖在pos_x和pos_y方向各自計算了一個位置編碼，每個維度的位置編碼長度爲num_pos_feats，默認爲Transformer的特征長度的一半，爲128。對pos_x和pos_y，在奇數位置計算正弦，在偶數位置計算餘弦，然後将計算結果進行拼接。得到一個[batch_size, h, w, 256]的向量。最後進行轉置，獲得[batch_size, 256, h, w]的向量。
代碼如下：
Pythonclass PositionEmbeddingSine(nn.Module): """ 這是一個更标準的位置嵌入版本，按照sine進行分布 """ def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): super().__init__() self.num_pos_feats = num_pos_feats self.temperature = temperature self.normalize = normalize if scale is not None and normalize is False: raise ValueError("normalize should be True if scale is passed") if scale is None: scale = 2 * math.pi self.scale = scale def forward(self, tensor_list: NestedTensor): x = tensor_list.tensors mask = tensor_list.mask assert mask is not None not_mask = ~mask y_embed = not_mask.cumsum(1, dtype=torch.float32) x_embed = not_mask.cumsum(2, dtype=torch.float32) if self.normalize: eps = 1e-6 y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) pos_x = x_embed[:, :, :, None] / dim_t pos_y = y_embed[:, :, :, None] / dim_t pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) return pos2、編碼網絡Encoder網絡介紹
a、Transformer Encoder的構建
在上文中，我們獲得了兩個矩陣，一個矩陣是輸入圖片的特征矩陣，一個是特征矩陣對應的位置編碼。它們的shape分别爲[batch_size, 2048, 25, 25]、[batch_size, 256, 25, 25]。
在編碼網絡部分，DETR使用Transformer的Encoder部分進行特征提取。我們需要首先對特征矩陣進行通道的縮放，如果直接對特征矩陣進行transformer的特征提取的話，由于網絡的通道數太大（2048），會直接導緻顯存不足。利用一個1x1的nn.Conv2d進行通道的壓縮，壓縮後的通道爲256，即Transformer用到的特征長度。此時我們獲得了一個shape爲[batch_size, 256, 25, 25]的特征矩陣。
然後我們對特征矩陣與位置編碼的高寬維度進行平鋪獲得兩個shape爲[batch_size, 256, 625]的矩陣，由于我們使用的是Pytorch自帶的nn.MultiheadAttention，該模塊要求batch_size位于第1維，序列長度位于第0維，所以我們将特征矩陣與位置編碼進行轉置，轉置後的兩個矩陣爲[625, batch_size, 256]。
我們此時可以将其輸入到Encoder當中進行特征提取。Encoder并不會改變輸入的shape，因此經過Encoder進行特征提取的加強後的特征序列shape也爲[625, batch_size, 256]。
由于在DETR中，Transformer的Encoder直接使用了Pytorch的MultiheadAttention，我們不必太貴糾結原理，簡單了解一下就可以，在DETR中，整個Transformer Encoder的實現代碼爲：
Pythonclass TransformerEncoder(nn.Module): def __init__(self, encoder_layer, num_layers, norm=None): super().__init__() self.layers = _get_clones(encoder_layer, num_layers) self.num_layers = num_layers self.norm = norm def forward(self, src, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): output = src # 625, batch_size, 256 => ...(x6)... => 625, batch_size, 256 for layer in self.layers: output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos) if self.norm is not None: output = self.norm(output) return output class TransformerEncoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu", normalize_before=False): super().__init__() # Self-Attention模塊 self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) # FFN模塊 # Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before def with_pos_embed(self, tensor, pos: Optional[Tensor]): return tensor if pos is None else tensor + pos def forward_post(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): # 添加位置信息 # 625, batch_size, 256 => 625, batch_size, 256 q = k = self.with_pos_embed(src, pos) # 使用自注意力機制模塊 # 625, batch_size, 256 => 625, batch_size, 256 src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] # 添加殘差結構 # 625, batch_size, 256 => 625, batch_size, 256 src = src + self.dropout1(src2) # 添加FFN結構 # 625, batch_size, 256 => 625, batch_size, 2048 => 625, batch_size, 256 src = self.norm1(src) src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) # 添加殘差結構 # 625, batch_size, 256 => 625, batch_size, 256 src = src + self.dropout2(src2) src = self.norm2(src) return src def forward_pre(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): src2 = self.norm1(src) q = k = self.with_pos_embed(src2, pos) src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] src = src + self.dropout1(src2) src2 = self.norm2(src) src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) src = src + self.dropout2(src2) return src def forward(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): if self.normalize_before: return self.forward_pre(src, src_mask, src_key_padding_mask, pos) return self.forward_post(src, src_mask, src_key_padding_mask, pos)b、Self-attention結構解析
在這裏可以簡單了解一下多頭注意力機制的原理，多頭注意力機制的計算原理如下：
看懂Self-attention結構，其實看懂下面這個動圖就可以了，動圖中存在一個序列的三個單位輸入，每一個序列單位的輸入都可以通過三個處理（比如全連接）獲得Query、Key、Value，Query是查詢向量、Key是鍵向量、Value值向量。
如果我們想要獲得input-1的輸出，那麽我們進行如下幾步：
1、利用input-1的查詢向量，分别乘上input-1、input-2、input-3的鍵向量，此時我們獲得了三個score。
2、然後對這三個score取softmax，獲得了input-1、input-2、input-3各自的重要程度。
3、然後将這個重要程度乘上input-1、input-2、input-3的值向量，求和。
4、此時我們獲得了input-1的輸出。
如圖所示，我們進行如下幾步：
1、input-1的查詢向量爲[1, 0, 2]，分别乘上input-1、input-2、input-3的鍵向量，獲得三個score爲2，4，4。
2、然後對這三個score取softmax，獲得了input-1、input-2、input-3各自的重要程度，獲得三個重要程度爲0.0，0.5，0.5。
3、然後将這個重要程度乘上input-1、input-2、input-3的值向量，求和，即
0.0 ∗ [ 1 , 2 , 3 ] + 0.5 ∗ [ 2 , 8 , 0 ] + 0.5 ∗ [ 2 , 6 , 3 ] = [ 2.0 , 7.0 , 1.5 ] 0.0 * [1, 2, 3] + 0.5 * [2, 8, 0] + 0.5 * [2, 6, 3] = [2.0, 7.0, 1.5] 0.0∗[1,2,3]+0.5∗[2,8,0]+0.5∗[2,6,3]=[2.0,7.0,1.5]。
4、此時我們獲得了input-1的輸出 [2.0, 7.0, 1.5]。
上述的例子中，序列長度僅爲3，每個單位序列的特征長度僅爲3，在DETR的Transformer Encoder中，序列長度爲625，每個單位序列的特征長度爲256 // num_heads。但計算過程是一樣的。在實際運算時，我們采用矩陣進行運算。
實際的矩陣運算過程如下圖所示。我以實際矩陣爲例子給大家解析：
輸入的Query、Key、Value如下圖所示：
首先利用 查詢向量query 叉乘 轉置後的鍵向量key，這一步可以通俗的理解爲，利用查詢向量去查詢序列的特征，獲得序列每個部分的重要程度score。
輸出的每一行，都代表input-1、input-2、input-3，對當前input的貢獻，我們對這個貢獻值取一個softmax。
然後利用 score 叉乘 value，這一步可以通俗的理解爲，将序列每個部分的重要程度重新施加到序列的值上去。
這個矩陣運算的代碼如下所示，各位同學可以自己試試。
Pythonimport numpy as np def soft_max(z): t = np.exp(z) a = np.exp(z) / np.expand_dims(np.sum(t, axis=1), 1) return a Query = np.array([ [1,0,2], [2,2,2], [2,1,3] ]) Key = np.array([ [0,1,1], [4,4,0], [2,3,1] ]) Value = np.array([ [1,2,3], [2,8,0], [2,6,3] ]) scores = Query @ Key.T print(scores) scores = soft_max(scores) print(scores) out = scores @ Value print(out)3、解碼網絡Decoder網絡介紹
通過上述第二步，我們可以獲得一個利用Encoder加強特征提取後的特征矩陣，它的shape爲[625, batch_size, 256]。
在encoder部分獲得的一個加強後的有效特征層會在這一部分進行解碼，解碼需要使用到一個非常重要的可學習模塊，即上圖呈現的object queries。在decoder部分，我們使用一個可學習的查詢向量q對加強後的有效特征層進行查詢，獲得預測結果。
在實際構建時，我們首先利用nn.Embedding(num_queries, hidden_dim)創建一個Embedding類别，然後利用.weight獲取這個Embedding的權重作爲可學習的查詢向量query_embed。默認的num_queries值爲100，hidden_dim值爲256。因此查詢向量query_embed本質上是一個[100, 256]的矩陣。加上batch維度後變成[100, batch_size, 256]。
Pythonself.query_embed = nn.Embedding(num_queries, hidden_dim) self.query_embed.weight另外，我們還通過tgt = torch.zeros_like(query_embed)創建了一個與查詢向量一樣shape的矩陣，作爲輸入。
參考下方右邊Transformer Decoder的結構，tgt作爲下圖的Output Embedding輸入到Decoder中，query_embed作爲Positional Encoding輸入到Decoder中。
首先自我進行一個Self-Attention的結構，輸入是[100, batch_size, 256]，輸出也是[100, batch_size, 256]。
然後再次利用另一個Self-Attention，将剛剛獲得的[100, batch_size, 256]輸出作爲Self-Attention的q，Encoder加強特征提取後的特征矩陣作爲Self-Attention的k、v，進行特征提取。這個過程可以理解爲使用查詢向量對Self-Attention的k、v進行查詢。由于查詢向量q的序列長度爲100，無論k、v的序列長度爲多少，最終輸出的序列長度都爲100。
因此對于解碼網絡Decoder而言，輸出的序列shape爲[100, batch_size, 256]。
實現代碼爲：
Pythonclass TransformerDecoder(nn.Module): def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False): super().__init__() self.layers = _get_clones(decoder_layer, num_layers) self.num_layers = num_layers self.norm = norm self.return_intermediate = return_intermediate def forward(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): output = tgt intermediate = [] for layer in self.layers: output = layer(output, memory, tgt_mask=tgt_mask, memory_mask=memory_mask, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask, pos=pos, query_pos=query_pos) if self.return_intermediate: intermediate.append(self.norm(output)) if self.norm is not None: output = self.norm(output) if self.return_intermediate: intermediate.pop() intermediate.append(output) if self.return_intermediate: return torch.stack(intermediate) return output.unsqueeze(0) class TransformerDecoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu", normalize_before=False): super().__init__() # q自己做一個self-attention self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) # q、k、v聯合做一個self-attention self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) # FFN模塊 # Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before def with_pos_embed(self, tensor, pos: Optional[Tensor]): return tensor if pos is None else tensor + pos def forward_post(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): #---------------------------------------------# # q自己做一個self-attention #---------------------------------------------# # tgt + query_embed # 100, batch_size, 256 => 100, batch_size, 256 q = k = self.with_pos_embed(tgt, query_pos) # q = k = v = 100, batch_size, 256 => 100, batch_size, 256 tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0] # 添加殘差結構 # 100, batch_size, 256 => 100, batch_size, 256 tgt = tgt + self.dropout1(tgt2) tgt = self.norm1(tgt) #---------------------------------------------# # q、k、v聯合做一個self-attention #---------------------------------------------# # q = 100, batch_size, 256, k = 625, batch_size, 256, v = 625, batch_size, 256 # 輸出的序列長度以q爲準 => 100, batch_size, 256 tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), key=self.with_pos_embed(memory, pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0] # 添加殘差結構 # 100, batch_size, 256 => 100, batch_size, 256 tgt = tgt + self.dropout2(tgt2) tgt = self.norm2(tgt) #---------------------------------------------# # 做一個FFN #---------------------------------------------# # 100, batch_size, 256 => 100, batch_size, 2048 => 100, batch_size, 256 tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt def forward_pre(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): tgt2 = self.norm1(tgt) q = k = self.with_pos_embed(tgt2, query_pos) tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0] tgt = tgt + self.dropout1(tgt2) tgt2 = self.norm2(tgt) tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), key=self.with_pos_embed(memory, pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0] tgt = tgt + self.dropout2(tgt2) tgt2 = self.norm3(tgt) tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) tgt = tgt + self.dropout3(tgt2) return tgt def forward(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): if self.normalize_before: return self.forward_pre(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) return self.forward_post(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)4、預測頭的構建
解碼網絡Decoder的輸出爲[100, batch_size, 256]，在實際使用時，爲了方便，我們再次把batch_size放回第0維度，得到的矩陣爲：[batch_size, 100, 256]
prediction heads是DETR的分類器與回歸器，其實就是對decoder獲得的預測結果進行全連接，兩次全連接分别代表種類和回歸參數。圖上畫了4個FFN，源碼中是2個FFN。
其中輸出分類信息的頭，它最終的全連接神經元個數爲num_classes + 1，num_classes代表要區分的類别數量，+1代表背景類。
如果使用的是voc訓練集，類則爲20種，最後的維度應該爲21。
如果使用的是coco訓練集，類則爲80種，不過中間有一些爲空的種類，空種類有11個，最後的維度應該爲80+11+1=92。
因此分類信息頭的輸出shape爲[batch_size, 100, num_classes + 1]
其中輸出回歸信息的頭，它最終的全連接神經元個數爲4。輸出時會取一個sigmoid。
前兩個系數代表中心點坐标，後兩個系數代表預測框寬高。
因此分類信息頭的輸出shape爲[batch_size, 100, 4]
實現代碼如下：
Python# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved import torch import torch.nn.functional as F from torch import nn from . import ops from .backbone import build_backbone from .ops import NestedTensor, nested_tensor_from_tensor_list from .transformer import build_transformer class MLP(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim, num_layers): super().__init__() self.num_layers = num_layers h = [hidden_dim] * (num_layers - 1) self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) def forward(self, x): for i, layer in enumerate(self.layers): x = F.relu(layer(x)) if i batch_size, 2048, 25, 25 features, pos = self.backbone(samples) # 将網絡的結果進行分割，把特征和mask進行分開 # batch_size, 2048, 25, 25, batch_size, 25, 25 src, mask = features[-1].decompose() assert mask is not None # 将主幹的結果進行一個映射，然後和查詢向量和位置向量傳入transformer。 # batch_size, 2048, 25, 25 => batch_size, 256, 25, 25 => 6, batch_size, 100, 256 hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0] # 輸出分類信息 # 6, batch_size, 100, 256 => 6, batch_size, 100, 21 outputs_class = self.class_embed(hs) # 輸出回歸信息 # 6, batch_size, 100, 256 => 6, batch_size, 100, 4 outputs_coord = self.bbox_embed(hs).sigmoid() # 隻輸出transformer最後一層的内容 # batch_size, 100, 21, batch_size, 100, 4 out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]} if self.aux_loss: out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord) return out @torch.jit.unused def _set_aux_loss(self, outputs_class, outputs_coord): return [{'pred_logits': a, 'pred_boxes': b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])] def freeze_bn(self): for m in self.modules(): if isinstance(m, nn.BatchNorm2d): m.eval()三、預測結果的解碼
由第二步我們可以獲得預測結果，shape分别爲[batch_size, 100, num_classes + 1]，[batch_size, 100, 4]的數據。
在DETR中，并不存在先驗框，也就不需要對先驗框進行調整獲得預測框。
回歸預測結果前兩個系數代表中心點坐标，後兩個系數代表預測框寬高。由于回歸預測結果取了sigmoid，所以值在0-1之間，直接乘上輸入圖片的寬高就是預測框在原圖上的寬高了。
分類預測結果代表這個預測框的種類。前num_classes個系數代表所區分類别的概率，1代表爲背景概率。解碼過程非常簡單，下面代碼中輸出的output就是預測結果。
Pythonimport numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torchvision.ops import nms class DecodeBox(nn.Module): """ This module converts the model's output into the format expected by the coco api""" def box_cxcywh_to_xyxy(self, x): x_c, y_c, w, h = x.unbind(-1) b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] return torch.stack(b, dim=-1) @torch.no_grad() def forward(self, outputs, target_sizes, confidence): out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes'] assert len(out_logits) == len(target_sizes) assert target_sizes.shape[1] == 2 prob = F.softmax(out_logits, -1) scores, labels = prob[..., :-1].max(-1) # convert to [x0, y0, x1, y1] format boxes = self.box_cxcywh_to_xyxy(out_bbox) # and from relative [0, 1] to absolute [0, height] coordinates img_h, img_w = target_sizes.unbind(1) scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) boxes = boxes * scale_fct[:, None, :] outputs = torch.cat([ torch.unsqueeze(boxes[:, :, 1], -1), torch.unsqueeze(boxes[:, :, 0], -1), torch.unsqueeze(boxes[:, :, 3], -1), torch.unsqueeze(boxes[:, :, 2], -1), torch.unsqueeze(scores, -1), torch.unsqueeze(labels, -1), ], -1) results = [] for output in outputs: results.append(output[output[:, 4] > confidence]) # results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)] return results四、訓練部分
1、計算Loss所需内容
計算loss實際上是網絡的預測結果和網絡的真實結果的對比。
和網絡的預測結果一樣，網絡的損失也由兩個部分組成，分别是Reg部分、Cls部分。Reg部分是特征點的回歸參數判斷、Cls部分是特征點包含的物體的種類。
2、正樣本的匹配過程
在DETR中，訓練時正樣本的匹配過程基于匈牙利算法，名字很高級，但是别被吓到，它其實隻是做一個匹配。
不管這個算法叫啥，本身它的作用都是用來進行匹配的，我們看看網絡的輸出和真實框的情況，去掉batch_size維度後，網絡的輸出爲[100, 4]和[100, num_classes + 1]。真實框的shape爲[num_gt, 5]，5中前4個系數爲真實框的坐标，最後一個系數爲真實框的種類。
匹配算法的工作隻是将100個預測結果和num_gt個真實框進行匹配就可以。一個真實框隻匹配一個預測結果，其它的預測結果作爲背景進行拟合。因此，匹配算法的工作是去找到最适合預測num_gt個真實框的num_gt個預測結果。因此我們需要去計算一個代價矩陣（Cost矩陣），用于代表100個預測結果和num_gt個真實框的關系。這是一個[100, gt]的矩陣。
這個代價矩陣由三個部分組成：
a、計算分類成本。獲得預測結果中，該真實框類别對應的預測值，如果預測值越大代表這個預測框預測的越準确，它的成本就越低。
b、計算預測框和真實框之間的L1成本。獲得預測結果中，預測框的坐标，将預測框的坐标和真實框的坐标做一個l1距離，預測的越準，它的成本就越低。
c、計算預測框和真實框之間的IOU成本。獲得預測結果中，預測框的坐标，将預測框的坐标和真實框的坐标做一個IOU距離，預測的越準，它的成本就越低。
三個按照一定的權重相加，就獲得了代價矩陣，這是一個[100, gt]的矩陣。
然後根據代價矩陣，使用匈牙利算法計算最低代價的情況。爲什麽不直接根據代價矩陣選擇真實框最接近的預測結果負責預測呢？因爲有可能一個預測結果同時最接近兩個真實框。匈牙利算法所做的工作其實隻是在代價最小的情況下，将預測結果都匹配上真實框。
Pythonclass HungarianMatcher(nn.Module): """ 此Matcher計算真實框和網絡預測之間的分配 因爲預測多于目标，對最佳預測進行1對1匹配。 """ def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1): super().__init__() # 這是Cost中分類錯誤的相對權重 self.cost_class = cost_class # 這是Cost中邊界框坐标L1誤差的相對權重 self.cost_bbox = cost_bbox # 這是Cost中邊界框giou損失的相對權重 self.cost_giou = cost_giou assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, "all costs cant be 0" @torch.no_grad() def forward(self, outputs, targets): # 獲得輸入的batch_size和query數量 bs, num_queries = outputs["pred_logits"].shape[:2] # 将預測結果的batch維度進行平鋪 # [batch_size * num_queries, num_classes] out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1) # [batch_size * num_queries, 4] out_bbox = outputs["pred_boxes"].flatten(0, 1) # 将真實框進行concat tgt_ids = torch.cat([v["labels"] for v in targets]) tgt_bbox = torch.cat([v["boxes"] for v in targets]) # 計算分類成本。預測越準值越小。 cost_class = -out_prob[:, tgt_ids] # 計算預測框和真實框之間的L1成本。預測越準值越小。 cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) # 計算預測框和真實框之間的IOU成本。預測越準值越小。 cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) # 最終的成本矩陣 C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou C = C.view(bs, num_queries, -1).cpu() sizes = [len(v["boxes"]) for v in targets] # 對每一張圖片進行指派任務，也就是找到真實框對應的num_queries裏面最接近的預測結果，也就是指派num_queries裏面一個預測框去預測某一個真實框 indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))] # 返回指派的結果 return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]3、計算Loss
在完成預測結果和真實框的匹配後，将匹配上的預測結果和真實框進行損失計算。
由第一部分可知，DETR的損失由兩個部分組成：
1、Reg部分，由第2部分可知道每個真實框對應的預測框，獲取到每個真實框對應的預測框後，利用預測框和真實框計算l1的損失和giou損失。
2、Cls部分，由第2部分可知道每個真實框對應的預測框，獲取到每個真實框對應的預測框後，取出該預測框的種類預測結果，根據真實框的種類計算交叉熵損失。沒有匹配上真實框的預測框作爲背景。
Pythonclass SetCriterion(nn.Module): """ 計算DETR的損失。該過程分爲兩個步驟： 1、計算了真實框和模型輸出之間的匈牙利分配 2、根據分配結果計算損失 """ def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses): super().__init__() # 類别數量 self.num_classes = num_classes # 用于匹配的匹配類HungarianMatcher self.matcher = matcher # 損失的權值分配 self.weight_dict = weight_dict # 背景的權重 self.eos_coef = eos_coef # 需要計算的損失 self.losses = losses # 種類的權重 empty_weight = torch.ones(self.num_classes + 1) empty_weight[-1] = self.eos_coef self.register_buffer('empty_weight', empty_weight) def forward(self, outputs, targets): # 首先計算不屬于輔助頭的損失 outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'} # 通過matcher計算每一個圖片，預測框和真實框的對應情況 indices = self.matcher(outputs_without_aux, targets) # 計算這個batch中所有圖片的總的真實框數量 # 計算所有節點的目标框的平均數量，以實現标準化 num_boxes = sum(len(t["labels"]) for t in targets) num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device) if is_dist_avail_and_initialized(): torch.distributed.all_reduce(num_boxes) num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item() # 計算所有的損失 losses = {} for loss in self.losses: losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes)) # 在輔助損失的情況下，我們對每個中間層的輸出重複此過程。 if 'aux_outputs' in outputs: for i, aux_outputs in enumerate(outputs['aux_outputs']): indices = self.matcher(aux_outputs, targets) for loss in self.losses: if loss == 'masks': continue kwargs = {} if loss == 'labels': kwargs = {'log': False} l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs) l_dict = {k + f'_{i}': v for k, v in l_dict.items()} losses.update(l_dict) return losses def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs): # 根據名稱計算損失 loss_map = { 'labels' : self.loss_labels, 'cardinality' : self.loss_cardinality, 'boxes' : self.loss_boxes, } assert loss in loss_map, f'do you really want to compute {loss} loss?' return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs) def loss_labels(self, outputs, targets, indices, num_boxes, log=True): assert 'pred_logits' in outputs # 獲得輸出中的分類部分 src_logits = outputs['pred_logits'] # 找到預測結果中有對應真實框的預測框 idx = self._get_src_permutation_idx(indices) # 獲得整個batch所有框的類别 target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)]) target_classes = torch.full(src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device) # 将其中對應的預測框設置爲目标類别，否則爲背景 target_classes[idx] = target_classes_o # 計算交叉熵 loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight) losses = {'loss_ce': loss_ce} if log: # TODO this should probably be a separate loss, not hacked in this one here losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0] return losses @torch.no_grad() def loss_cardinality(self, outputs, targets, indices, num_boxes): pred_logits = outputs['pred_logits'] device = pred_logits.device # 計算每個batch真實框的數量 tgt_lengths = torch.as_tensor([len(v["labels"]) for v in targets], device=device) # 計算不是背景的預測數 card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1) # 然後将不是背景的預測數和真實情況做一個l1損失 card_err = F.l1_loss(card_pred.float(), tgt_lengths.float()) losses = {'cardinality_error': card_err} return losses def loss_boxes(self, outputs, targets, indices, num_boxes): assert 'pred_boxes' in outputs # 找到預測結果中有對應真實框的預測框 idx = self._get_src_permutation_idx(indices) # 将預測結果中有對應真實框的預測框取出 src_boxes = outputs['pred_boxes'][idx] # 取出真實框 target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0) # 預測框和所有的真實框計算l1的損失 loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none') # 計算giou損失 loss_giou = 1 - torch.diag(ops.generalized_box_iou(ops.box_cxcywh_to_xyxy(src_boxes), ops.box_cxcywh_to_xyxy(target_boxes))) # 返回兩個損失 losses = {} losses['loss_bbox'] = loss_bbox.sum() / num_boxes losses['loss_giou'] = loss_giou.sum() / num_boxes return losses def _get_src_permutation_idx(self, indices): batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)]) src_idx = torch.cat([src for (src, _) in indices]) return batch_idx, src_idx def _get_tgt_permutation_idx(self, indices): batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]) tgt_idx = torch.cat([tgt for (_, tgt) in indices]) return batch_idx, tgt_idx訓練自己的DETR模型
首先前往Github下載對應的倉庫，下載完後利用解壓軟件解壓，之後用編程軟件打開文件夾。
注意打開的根目錄必須正确，否則相對目錄不正确的情況下，代碼将無法運行。
一定要注意打開後的根目錄是文件存放的目錄。
一、數據集的準備
本文使用VOC格式進行訓練，訓練前需要自己制作好數據集，如果沒有自己的數據集，可以通過Github連接下載VOC12+07的數據集嘗試下。
訓練前将标簽文件放在VOCdevkit文件夾下的VOC2007文件夾下的Annotation中。
訓練前将圖片文件放在VOCdevkit文件夾下的VOC2007文件夾下的JPEGImages中。
此時數據集的擺放已經結束。
二、數據集的處理
在完成數據集的擺放之後，我們需要對數據集進行下一步的處理，目的是獲得訓練用的2007_train.txt以及2007_val.txt，需要用到根目錄下的voc_annotation.py。
voc_annotation.py裏面有一些參數需要設置。
分别是annotation_mode、classes_path、trainval_percent、train_percent、VOCdevkit_path，第一次訓練可以僅修改classes_path
Python''' annotation_mode用于指定該文件運行時計算的内容 annotation_mode爲0代表整個标簽處理過程，包括獲得VOCdevkit/VOC2007/ImageSets裏面的txt以及訓練用的2007_train.txt、2007_val.txt annotation_mode爲1代表獲得VOCdevkit/VOC2007/ImageSets裏面的txt annotation_mode爲2代表獲得訓練用的2007_train.txt、2007_val.txt ''' annotation_mode = 0 ''' 必須要修改，用于生成2007_train.txt、2007_val.txt的目标信息 與訓練和預測所用的classes_path一緻即可 如果生成的2007_train.txt裏面沒有目标信息 那麽就是因爲classes沒有設定正确 僅在annotation_mode爲0和2的時候有效 ''' classes_path = 'model_data/voc_classes.txt' ''' trainval_percent用于指定(訓練集+驗證集)與測試集的比例，默認情況下 (訓練集+驗證集):測試集 = 9:1 train_percent用于指定(訓練集+驗證集)中訓練集與驗證集的比例，默認情況下 訓練集:驗證集 = 9:1 僅在annotation_mode爲0和1的時候有效 ''' trainval_percent = 0.9 train_percent = 0.9 ''' 指向VOC數據集所在的文件夾 默認指向根目錄下的VOC數據集 ''' VOCdevkit_path = 'VOCdevkit'classes_path用于指向檢測類别所對應的txt，以voc數據集爲例，我們用的txt爲：
訓練自己的數據集時，可以自己建立一個cls_classes.txt，裏面寫自己所需要區分的類别。
三、開始網絡訓練
通過voc_annotation.py我們已經生成了2007_train.txt以及2007_val.txt，此時我們可以開始訓練了。
訓練的參數較多，大家可以在下載庫後仔細看注釋，其中最重要的部分依然是train.py裏的classes_path。
classes_path用于指向檢測類别所對應的txt，這個txt和voc_annotation.py裏面的txt一樣！訓練自己的數據集必須要修改！
修改完classes_path後就可以運行train.py開始訓練了，在訓練多個epoch後，權值會生成在logs文件夾中。
其它參數的作用如下：
Python#---------------------------------# # Cuda 是否使用Cuda # 沒有GPU可以設置成False #---------------------------------# Cuda = True #---------------------------------------------------------------------# # distributed 用于指定是否使用單機多卡分布式運行 # 終端指令僅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定顯卡。 # Windows系統下默認使用DP模式調用所有顯卡，不支持DDP。 # DP模式： # 設置 distributed = False # 在終端中輸入 CUDA_VISIBLE_DEVICES=0,1 python train.py # DDP模式： # 設置 distributed = True # 在終端中輸入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py #---------------------------------------------------------------------# distributed = False #---------------------------------------------------------------------# # fp16 是否使用混合精度訓練 # 可減少約一半的顯存、需要pytorch1.7.1以上 #---------------------------------------------------------------------# fp16 = False #---------------------------------------------------------------------# # classes_path 指向model_data下的txt，與自己訓練的數據集相關 # 訓練前一定要修改classes_path，使其對應自己的數據集 #---------------------------------------------------------------------# classes_path = 'model_data/voc_classes.txt' #----------------------------------------------------------------------------------------------------------------------------# # 權值文件的下載請看README，可以通過網盤下載。模型的 預訓練權重 對不同數據集是通用的，因爲特征是通用的。 # 模型的 預訓練權重 比較重要的部分是 主幹特征提取網絡的權值部分，用于進行特征提取。 # 預訓練權重對于99%的情況都必須要用，不用的話主幹部分的權值太過随機，特征提取效果不明顯，網絡訓練的結果也不會好 # # 如果訓練過程中存在中斷訓練的操作，可以将model_path設置成logs文件夾下的權值文件，将已經訓練了一部分的權值再次載入。 # 同時修改下方的 凍結階段 或者 解凍階段 的參數，來保證模型epoch的連續性。 # # 當model_path = ''的時候不加載整個模型的權值。 # # 此處使用的是整個模型的權重，因此是在train.py進行加載的，下面的pretrain不影響此處的權值加載。 # 如果想要讓模型從主幹的預訓練權值開始訓練，則設置model_path = ''，下面的pretrain = True，此時僅加載主幹。 # 如果想要讓模型從0開始訓練，則設置model_path = ''，下面的pretrain = Fasle，Freeze_Train = Fasle，此時從0開始訓練，且沒有凍結主幹的過程。 # # 一般來講，網絡從0開始的訓練效果會很差，因爲權值太過随機，特征提取效果不明顯，因此非常、非常、非常不建議大家從0開始訓練！ # 如果一定要從0開始，可以了解imagenet數據集，首先訓練分類模型，獲得網絡的主幹部分權值，分類模型的 主幹部分 和該模型通用，基于此進行訓練。 #----------------------------------------------------------------------------------------------------------------------------# model_path = 'model_data/detr_resnet50_weights_coco.pth' #------------------------------------------------------# # input_shape 輸入的shape大小 #------------------------------------------------------# input_shape = [800, 800] #---------------------------------------------# # vgg # resnet50 #---------------------------------------------# backbone = "resnet50" #----------------------------------------------------------------------------------------------------------------------------# # pretrained 是否使用主幹網絡的預訓練權重，此處使用的是主幹的權重，因此是在模型構建的時候進行加載的。 # 如果設置了model_path，則主幹的權值無需加載，pretrained的值無意義。 # 如果不設置model_path，pretrained = True，此時僅加載主幹開始訓練。 # 如果不設置model_path，pretrained = False，Freeze_Train = Fasle，此時從0開始訓練，且沒有凍結主幹的過程。 #----------------------------------------------------------------------------------------------------------------------------# pretrained = False #----------------------------------------------------------------------------------------------------------------------------# # 訓練分爲兩個階段，分别是凍結階段和解凍階段。設置凍結階段是爲了滿足機器性能不足的同學的訓練需求。 # 凍結訓練需要的顯存較小，顯卡非常差的情況下，可設置Freeze_Epoch等于UnFreeze_Epoch，此時僅僅進行凍結訓練。 # # 在此提供若幹參數設置建議，各位訓練者根據自己的需求進行靈活調整： # （一）從整個模型的預訓練權重開始訓練： # AdamW： # Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 100，Freeze_Train = True，optimizer_type = 'adamw'，Init_lr = 1e-4，weight_decay = 1e-4。（凍結） # Init_Epoch = 0，UnFreeze_Epoch = 100，Freeze_Train = False，optimizer_type = 'adamw'，Init_lr = 1e-4，weight_decay = 1e-4。（不凍結） # 其中：UnFreeze_Epoch可以在100-300之間調整。 # （二）從主幹網絡的預訓練權重開始訓練： # AdamW： # Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 300，Freeze_Train = True，optimizer_type = 'adamw'，Init_lr = 1e-4，weight_decay = 1e-4。（凍結） # Init_Epoch = 0，UnFreeze_Epoch = 300，Freeze_Train = False，optimizer_type = 'adamw'，Init_lr = 1e-4，weight_decay = 1e-4。（不凍結） # 其中：由于從主幹網絡的預訓練權重開始訓練，主幹的權值不一定适合目标檢測，需要更多的訓練跳出局部最優解。 # UnFreeze_Epoch可以在150-300之間調整，YOLOV5和YOLOX均推薦使用300。 # Adam相較于SGD收斂的快一些。因此UnFreeze_Epoch理論上可以小一點，但依然推薦更多的Epoch。 # （三）batch_size的設置： # 在顯卡能夠接受的範圍内，以大爲好。顯存不足與數據集大小無關，提示顯存不足（OOM或者CUDA out of memory）請調小batch_size。 # 受到BatchNorm層影響，batch_size最小爲2，不能爲1。 # 正常情況下Freeze_batch_size建議爲Unfreeze_batch_size的1-2倍。不建議設置的差距過大，因爲關系到學習率的自動調整。 #----------------------------------------------------------------------------------------------------------------------------# #------------------------------------------------------------------# # 凍結階段訓練參數 # 此時模型的主幹被凍結了，特征提取網絡不發生改變 # 占用的顯存較小，僅對網絡進行微調 # Init_Epoch 模型當前開始的訓練世代，其值可以大于Freeze_Epoch，如設置： # Init_Epoch = 60、Freeze_Epoch = 50、UnFreeze_Epoch = 100 # 會跳過凍結階段，直接從60代開始，并調整對應的學習率。 # （斷點續練時使用） # Freeze_Epoch 模型凍結訓練的Freeze_Epoch # (當Freeze_Train=False時失效) # Freeze_batch_size 模型凍結訓練的batch_size # (當Freeze_Train=False時失效) #------------------------------------------------------------------# Init_Epoch = 0 Freeze_Epoch = 50 Freeze_batch_size = 8 #------------------------------------------------------------------# # 解凍階段訓練參數 # 此時模型的主幹不被凍結了，特征提取網絡會發生改變 # 占用的顯存較大，網絡所有的參數都會發生改變 # UnFreeze_Epoch 模型總共訓練的epoch # SGD需要更長的時間收斂，因此設置較大的UnFreeze_Epoch # Adam可以使用相對較小的UnFreeze_Epoch # Unfreeze_batch_size 模型在解凍後的batch_size #------------------------------------------------------------------# UnFreeze_Epoch = 300 Unfreeze_batch_size = 4 #------------------------------------------------------------------# # Freeze_Train 是否進行凍結訓練 # 默認先凍結主幹訓練後解凍訓練。 #------------------------------------------------------------------# Freeze_Train = True #------------------------------------------------------------------# # 其它訓練參數：學習率、優化器、學習率下降有關 #------------------------------------------------------------------# #------------------------------------------------------------------# # Init_lr 模型的最大學習率，在DETR中，Backbone的學習率爲Transformer模塊的0.1倍 # Min_lr 模型的最小學習率，默認爲最大學習率的0.01 #------------------------------------------------------------------# Init_lr = 1e-4 Min_lr = Init_lr * 0.01 #------------------------------------------------------------------# # optimizer_type 使用到的優化器種類，可選的有adam、sgd # 當使用Adam優化器時建議設置 Init_lr=1e-4 # 當使用AdamW優化器時建議設置 Init_lr=1e-4 # 當使用SGD優化器時建議設置 Init_lr=1e-2 # momentum 優化器内部使用到的momentum參數 # weight_decay 權值衰減，可防止過拟合 # adam會導緻weight_decay錯誤，使用adam時建議設置爲0。 #------------------------------------------------------------------# optimizer_type = "adamw" momentum = 0.9 weight_decay = 1e-4 #------------------------------------------------------------------# # lr_decay_type 使用到的學習率下降方式，可選的有step、cos #------------------------------------------------------------------# lr_decay_type = "cos" #------------------------------------------------------------------# # save_period 多少個epoch保存一次權值 #------------------------------------------------------------------# save_period = 10 #------------------------------------------------------------------# # save_dir 權值與日志文件保存的文件夾 #------------------------------------------------------------------# save_dir = 'logs' #------------------------------------------------------------------# # eval_flag 是否在訓練時進行評估，評估對象爲驗證集 # 安裝pycocotools庫後，評估體驗更佳。 # eval_period 代表多少個epoch評估一次，不建議頻繁的評估 # 評估需要消耗較多的時間，頻繁評估會導緻訓練非常慢 # 此處獲得的mAP會與get_map.py獲得的會有所不同，原因有二： # （一）此處獲得的mAP爲驗證集的mAP。 # （二）此處設置評估參數較爲保守，目的是加快評估速度。 #------------------------------------------------------------------# eval_flag = True eval_period = 10 #------------------------------------------------------------------# # 官方提示爲TODO this is a hack # 穩定性未知，默認爲不開啓 #------------------------------------------------------------------# aux_loss = False #------------------------------------------------------------------# # num_workers 用于設置是否使用多線程讀取數據 # 開啓後會加快數據讀取速度，但是會占用更多内存 # 内存較小的電腦可以設置爲2或者0 #------------------------------------------------------------------# num_workers = 4 #----------------------------------------------------# # 獲得圖片路徑和标簽 #----------------------------------------------------# train_annotation_path = '2007_train.txt' val_annotation_path = '2007_val.txt'四、訓練結果預測
訓練結果預測需要用到兩個文件，分别是detr.py和predict.py。
我們首先需要去detr.py裏面修改model_path以及classes_path，這兩個參數必須要修改。
model_path指向訓練好的權值文件，在logs文件夾裏。
classes_path指向檢測類别所對應的txt。
完成修改後就可以運行predict.py進行檢測了。運行後輸入圖片路徑即可檢測。
