SSD源碼閱讀二

 個人博客:http://www.chenjianqu.com/

原文鏈接:http://www.chenjianqu.com/show-92.html

 

上篇博客<SSD源碼閱讀一>閱讀了ssd.pytorch項目的部分代碼,包括包括數據集、數據增強和構建網絡結構的部分。在模型類class SSD(nn.Module)中,構造函數定義這個變量self.priorbox = PriorBox(self.cfg) ,PriorBox是用於構造先驗框的類,代碼如下:

layers/functions/prior_box.py

class PriorBox(object):
    """Compute priorbox coordinates in center-offset form for each source feature map.
    """
    #cfg是定義在data/config.py裏面,是模型的相關配置
    def __init__(self, cfg):
        super(PriorBox, self).__init__()
        #輸入數據分辨率,一般爲300
        self.image_size = cfg['min_dim']
        #輸出特徵圖的數量
        self.num_priors = len(cfg['aspect_ratios'])
        #尺度
        self.variance = cfg['variance'] or [0.1]
        #各個輸出特徵圖的分辨率
        self.feature_maps = cfg['feature_maps']
        #計算先驗框用到的Smin和Smax
        self.min_sizes = cfg['min_sizes']
        self.max_sizes = cfg['max_sizes']
        #各個輸出特徵圖每個像素對應到原圖的大小
        self.steps = cfg['steps']
        #輸出特徵圖用到的比例
        self.aspect_ratios = cfg['aspect_ratios']
        self.clip = cfg['clip']
        #VOC或COCO
        self.version = cfg['name']
        for v in self.variance:
            if v <= 0:
                raise ValueError('Variances must be greater than 0')
                
    def forward(self):
        mean = []
        
        for k, f in enumerate(self.feature_maps):#f是特徵圖的邊長
            for i, j in product(range(f), repeat=2):#product(range(f), repeat=2)求range(f)和range(f)的笛卡爾積,生成各個像素的座標
'''
每個先驗框中心的計算公式是:(cx,cy)=((i+0.5)/|fk|, (j+0.5)/|fk|),其中|fk|是特徵圖的邊長
得到的座標範圍是[0,1]
'''
                f_k = self.image_size / self.steps[k]#steps是輸出特徵圖的下採樣率,則f_k是特徵圖的邊長
                # unit center x,y
                cx = (j + 0.5) / f_k
                cy = (i + 0.5) / f_k
’‘’
每個特徵圖使用的先驗框大小:
Sk=Smin+(Smax-Smin)(k-1)/(m-1),k的值[1,m],m是輸出特徵圖的數量,
Sk是先驗框相對於整張圖片的比例,Smin=0.2,Smax=0.95
先驗框的寬度w_k_a=Sk*ar^0.5,高度h_k_a=Sk/ar^0.5,
其中ar是特徵圖的比例,有{3, 2, 1, 1/2, 1/3}
‘’‘
                #下面計算各個比例的先驗框的的寬高,並把它們放到list()裏面
                #注意這裏的寬高是相對於原圖的比例
                # aspect_ratio: 1
                # rel size: min_size
                s_k = self.min_sizes[k]/self.image_size
                mean += [cx, cy, s_k, s_k]
                
                # aspect_ratio: 1
                # rel size: sqrt(s_k * s_(k+1))
                s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size))
                mean += [cx, cy, s_k_prime, s_k_prime]
                
                # rest of aspect ratios                
                for ar in self.aspect_ratios[k]:
                    mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)]
                    mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)]
        
        #reshape先驗框
        output = torch.Tensor(mean).view(-1, 4)
        
        #clamp_ 將output張量每個元素的夾緊到區間 [min,max],並返回結果到一個新張量。其實就是我們常用的截斷函數(分段閾值)。
        if self.clip:
            output.clamp_(max=1, min=0)
        
        return output

    上面代碼的代碼嚴格按照論文<SSD論文筆記>中先驗框的計算,forward()函數的中mean變量保存了一張圖片中各個先驗框的中心座標和長寬(均是相對值)。先驗框主要是計算損失函數的時候用到,前向傳播的時候,直接對每張圖片拷貝mean張量即可。

    回到SSD網絡結構,根據原論文的說法,與其他層相比,conv4_3具有不同的特徵尺度,因此使用ParseNet中介紹的L2 normalization技術將特徵圖中每個位置的feature norm縮放到20,並在反向傳播期間學習尺度。因此這裏也實現了L2 normalization,如下:

layers/modules/l2norm.py

class L2Norm(nn.Module):
    #參數:輸入特徵圖的通道數,縮放像素值到達的範圍
    def __init__(self,n_channels, scale):
        super(L2Norm,self).__init__()
        self.n_channels = n_channels
        self.gamma = scale or None
        self.eps = 1e-10
        self.weight = nn.Parameter(torch.Tensor(self.n_channels))
        self.reset_parameters()
        
    def reset_parameters(self):
        init.constant_(self.weight,self.gamma)
        
    def forward(self, x):
        norm = x.pow(2).sum(dim=1, keepdim=True).sqrt()+self.eps
        #x /= norm
        x = torch.div(x,norm)
        out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x

        return out

    到此,網絡結構部分基本完成。接下來,回到訓練的主程序train.py train()裏面,在構建網絡之後,因爲作者有多個GPU,因此它這裏設置數據並行:

train.py train()

    if args.cuda:
        net = torch.nn.DataParallel(ssd_net)#數據並行,多個GPU的時候使用
        cudnn.benchmark = True

    對於cudnn.benchmark,大部分情況下,設置這個 flag 可以讓內置的 cuDNN 的 auto-tuner 自動尋找最適合當前配置的高效算法,來達到優化運行效率的問題。如果網絡的輸入數據維度或類型上變化不大,設置  torch.backends.cudnn.benchmark = true  可以增加運行效率;如果網絡的輸入數據在每次 iteration 都變化的話,會導致 cnDNN 每次都會去尋找一遍最優配置,這樣反而會降低運行效率。

    接下來,程序判斷是從零開始訓練還是恢復訓練,如果是恢復訓練則直接加載整個模型的權重,否則加載vgg的權重,然後將模型轉換爲GPU模式:

train.py train()

#權重加載
    if args.resume:
        print('Resuming training, loading {}...'.format(args.resume))
        ssd_net.load_weights(args.resume)
    else:
        vgg_weights = torch.load(args.save_folder + args.basenet)
        print('Loading base network...')
        ssd_net.vgg.load_state_dict(vgg_weights)
#使用GPU模式
    if args.cuda:
        net = net.cuda()

    如果是從零開始訓練,那麼對新增層的權重進行xavier初始化:

train.py train()

#權重初始化
    if not args.resume:
        print('Initializing weights...')
        #使用xavier方法進行初始化,apply(fn):將fn函數遞歸地應用到網絡模型的每個子模型中進行參數的初始化。
        ssd_net.extras.apply(weights_init)
        ssd_net.loc.apply(weights_init)
        ssd_net.conf.apply(weights_init)

    然後設置優化器、設置損失函數:

train.py train()

#使用SGD優化器
    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
                          weight_decay=args.weight_decay)
#定義損失函數
    criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,False, args.cuda)

    根據論文原文<SSD論文筆記>,我們知道SSD的損失函數由定義損失和置信度損失組成,這裏的代碼實現如下,這裏回顧以下SSD損失函數的計算:

100.jpg

    其中定位損失各項的計算:

圖片1.jpg

    要計算上式中的損失函數,首先要確定xijk,即要將gt框與先驗框進行匹配,匹配的策略:策略1:找到與gt bbox的IOU最大的先驗包圍框,這樣保證每個gt bbox至少對應一個先驗框。但是一個圖片中gt bbox非常少,而先驗框卻很多,這樣正負樣本極其不平衡。策略2:對於其它的先驗包圍框,若與gt bbox的 IOU 大於某個閾值(一般是0.5),那麼該先驗框也與這個ground truth進行匹配。

    首先知道IoU的實現:

layers/box_utils.py  jaccard()

#注意這裏的包圍框格式是:[xmin,ymin,xmax,ymax],即包圍框左上角和右下角的座標

#計算兩組框中兩兩的交集,box_a.shape=[A,4],box_b.shape=[B,4]
def intersect(box_a, box_b):
        #box_a中框的數量
	A = box_a.size(0) 
	#box_b中框的數量
	B = box_b.size(0) 		
	
	#box_a[:,2:]原本的維度是(A,2),經過unsqueeze(1)函數變爲(A,1,2),再經過expand(A,B,2)變爲(A,B,2),該操作不會分配內存
	t1=box_a[:, 2:].unsqueeze(1).expand(A, B, 2)
	t2=box_b[:, 2:].unsqueeze(0).expand(A, B, 2)	
	max_xy = torch.min(t1,t2)#兩個[xmax,ymax]之間的比較,求他們兩點中小的點
	
	#兩個[xmin,ymin]之間的比較,求他們兩個的大的點
	n1=box_a[:, :2].unsqueeze(1).expand(A, B, 2)
	n2=box_b[:, :2].unsqueeze(0).expand(A, B, 2)
	min_xy = torch.max(n1,n2)
	
	#求出交集區域的寬高,inter.shape=[A,B,2]
	inter = torch.clamp((max_xy - min_xy), min=0)
	#求出交集的面積,返回.shape=[A,B]
	return inter[:, :, 0] * inter[:, :, 1] 


#計算兩組包圍框中兩兩包圍框的IoU,box_a.shape=[A,4],box_b.shape=[B,4]
def jaccard(box_a, box_b):
    #計算兩組包圍框中兩兩包圍框的交集面積,inter.shape=[A,B]
    inter = intersect(box_a, box_b)
    
    area_a = (box_a[:, 2]-box_a[:, 0]) * (box_a[:, 3]-box_a[:, 1]) #box_a中所有框的區域面積
    area_a = area_a.unsqueeze(1).expand_as(inter)  #將area_a的維度由[A,B]拓展爲[A,B]
    
    area_b = (box_b[:, 2]-box_b[:, 0]) * (box_b[:, 3]-box_b[:, 1])#box_b中所有框的區域面積
    area_b = area_b.unsqueeze(0).expand_as(inter)  # [A,B]
    
    #IoU=A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
    union = area_a + area_b - inter
    return inter / union  #維度:[A,B],
    #返回矩陣的第i行、第j列元素的值代表box_a中的第i個包圍框與box_b中的第j個包圍框的IoU

    然後就是匹配算法的實現:

layers/box_utils.py  match()

#將box的 [框中心座標的x,框中心座標的y,框的寬,框的高],
#轉換爲 [框左上角座標的x,框左上角座標的y,框右下角座標的x,框右下角座標的y]
def point_form(boxes):
    #[cx,cy]-[w/2,h/2]=[cx - w/2 , cy - h/2]=[xmin,ymin]
    return torch.cat(
                     (
                         boxes[:, :2] - boxes[:, 2:]/2, # xmin, ymin
                         boxes[:, :2] + boxes[:, 2:]/2  # xmax, ymax
                     ), 
                  1) #在列進行拼接



#GT框和先驗框匹配算法實現
def match(threshold,#挑選正樣本的閾值
		  truths, #當前圖片中的gt框,shape:[num_truths,4]
		  priors, #當前圖片中的先驗框,shape:[num_priors,4]
		  variances, #
		  labels, #當前圖片中每個gt框的標籤
		  loc_t,#匹配好的gt框
		  conf_t, #匹配好的置信度
		  idx#當前圖片在batch的索引
		 ):
	#計算IOU,overlaps的維度是[A,B],其中A是truths的數量,B是priors的數量
	overlaps = jaccard(
		truths,
		point_form(priors))
#point_form(priors)將priors的格式[框中心座標的x,框中心座標的y,框的寬,框的高],
#轉換爲 [框左上角座標的x,框左上角座標的y,框右下角座標的x,框右下角座標的y]

#計算每個gt框最匹配的先驗框。即overlaps中每行數據中的最大值,一行代表一個gt box與所有先驗框的IOU	
#shape爲[1,B],兩個返回值分別代表gt框最大的IoU值和對應的先驗框的索引
	best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)

#計算每個先驗框最匹配的gt框,即計算overlaps中每列數據中的最大值,一列代表一個先驗框與所有gt box的IOU,後面用於計算最大是否大於閾值
#shape爲[1,B],代表的是所有先驗框的最大IOU和對應的gt框的索引
	best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)

	best_truth_idx.squeeze_(0) #將shape從[1,B]變爲[B]
	best_truth_overlap.squeeze_(0)
	
	best_prior_idx.squeeze_(1) #將shape從[A,1]變爲[A]
	best_prior_overlap.squeeze_(1)

'''
index_fill_(dim, index, val)
按照index,將val的值填充self的dim維度。如:
>>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float)
>>> index = torch.tensor([0, 2])
>>> x.index_fill_(1, index, -1)
tensor([[-1.,  2., -1.],
        [-1.,  5., -1.],
        [-1.,  8., -1.]])

A是truths的數量,B是priors的數量
best_truth_overlap代表各個先驗框與所有gt box的最大的IoU,shape=[B]
best_prior_idx代表各個gt box的最大IoU匹配先驗框的索引,shape=[A]
因此best_truth_overlap.index_fill_(0, best_prior_idx, 2)的效果是
    設置每個gt框與最大IoU先驗框的的IoU爲2,這確保每個gt框匹配的prior框不會因爲IoU太低,而被過濾掉
'''
	best_truth_overlap.index_fill_(0, best_prior_idx, 2)  # ensure best prior
	
	#確保每個gt框都能匹配到一個先驗框。best_prior_idx.shape=[A],裏面的每個值代表一個gt框與最大IoU的先驗框的索引。
	for j in range(best_prior_idx.size(0)):
		idx=best_prior_idx[j] #第j個gt框 與各個先驗框IoU最大 的 先驗框索引
		#best_truth_idx表示所有先驗框與各個gt框IoU最大的gt框的索引
		best_truth_idx[idx]=j #第j個gt框匹配的先驗框設置其匹配的gt框索引爲j

	#獲取每個prior匹配的gt框
	matches = truths[best_truth_idx]          # Shape: [num_priors,4]
	#每個prior對應類別標籤,加1是因爲0是背景類,best_truth_idx.shape=[B]
	conf = labels[best_truth_idx] + 1         # Shape: [num_priors]

	#令IoU小於threshold的先驗框的類別標籤爲0,即爲背景類
	conf[best_truth_overlap < threshold] = 0  # label as background

	#對在裏面會進行預計算
	loc = encode(matches, priors, variances)

	
	loc_t[idx] = loc    # [num_priors,4] encoded offsets to learn
	conf_t[idx] = conf  # [num_priors] top class label for each prior
	

#對位置進行編碼,參數:每個先驗框匹配的gt框,先驗框,沒搞懂variances這個參數
def encode(matched, priors, variances):
    #下面計算的變量均來自 定位損失函數裏面的
    #匹配好的gt框的中心和先驗框中心的x、y的距離,即計算公式中的g_cxcy
    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]
    #計算公式中的g_cxcy_^
    g_cxcy /= (variances[0] * priors[:, 2:])

    #計算公式中的g_w和g_h
    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
    g_wh = torch.log(g_wh) / variances[1]

    return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]

    上面就是先驗框和gt框匹配的過程,經過這個步驟,就可以計算損失函數了,損失函數的計算放在layers/modules/multibox_loss.py裏面的MultiBoxLoss類裏面,來看損失函數計算的前半部分:

layers/modules/multibox_loss.py   forward()

def forward(self, predictions, targets):
'''
預測的數據
loc_data:預測框位置點矩陣 size: (batch_size,num_priors,4) -> (32,8732,4)
conf_data:預測框置信度矩陣 size:(batch_size,num_priors,num_classes) -> (32,8732,3)
priors:先驗框矩陣 size:(num_priors,4) -> (8732,4)
'''
    loc_data, conf_data, priors = predictions

    num = loc_data.size(0) #batch_size
    num_priors=loc_data.size(1) #每張圖片預測的先驗框數量
    priors = priors[:num_priors, :]
    
    num_priors = (priors.size(0))#先驗框的數量
    num_classes = self.num_classes #類別數

    #匹配先驗框和gt框,gt框的格式爲[:,xmin,ymin,xmax,ymax]
    loc_t = torch.Tensor(num, num_priors, 4)
    conf_t = torch.LongTensor(num, num_priors)
    for idx in range(num):
        truths = targets[idx][:, :-1].data #第idx張圖片的所有gt框
        labels = targets[idx][:, -1].data #第idx張圖片的所有gt框對應的標籤,爲整數
        defaults = priors.data #先驗框
        #先驗框匹配
        match(self.threshold, truths, defaults, self.variance, labels,loc_t, conf_t, idx)
    #經過匹配,loc_t是每個先驗框對應的gt框,維度是[32,8732,4],    
    #conf_t是先驗框的標籤,維度是[32,8732]
    
    if self.use_gpu:
        loc_t = loc_t.cuda()
        conf_t = conf_t.cuda()
        
    # wrap targets
    loc_t = Variable(loc_t, requires_grad=False)
    conf_t = Variable(conf_t, requires_grad=False)
    
    #如果某個先驗框的標籤>0,則設置pos標誌位爲true。pos的維度是[32,8732]
    pos = conf_t > 0
    
    #每張圖片中非背景類的先驗框數,num_pos的維度[32,1]
    num_pos = pos.sum(dim=1, keepdim=True)
    
    # Shape:[batch,num_priors,4],即將pos維度擴展爲[32,8732,4]
    pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
    
    #選出先驗框類別非0 對應的預測框和gt框
    loc_p = loc_data[pos_idx].view(-1, 4)
    loc_t = loc_t[pos_idx].view(-1, 4)
    
    #定位損失使用smooth_L1
    loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)

    上面是計算定位損失的過程,需要注意的是,網絡預測的是相對於包圍框中心和寬度和高度相對於先驗框的偏移值,因此也需要將先驗框對應的gt框轉換爲gt框相對於先驗框的偏移值,也就是box_utils.py encode()函數的作用。

    定位損失只用到了positive輸出框用於計算,而計算置信度損失的時候,不僅需要計算positive輸出框作爲正例,還需要其它輸出框作爲負例。但是剩餘的預測框的數量太多,會導致正負樣本極度不平衡,這會導致得到的模型偏向性十分嚴重,甚至不可用。因此我們需要從剩餘的輸出框中挑選出高置信度的作爲負樣本,這就是難例挖掘,SSD使用的正負樣本比例爲1:3。難例挖掘和置信度損失的代碼如下:

layers/modules/multibox_loss.py   forward()

#將[batch_size,8732,21] reshape爲[batch_size*8732,21]
batch_conf = conf_data.view(-1, self.num_classes)

'''
batch_conf是該批次的所有預測框的類別置信度,[batch_size*8732, 21]
conf_t是該批次的所有先驗框的匹配的標籤,[batch_size*8732, 1]
gather函數的作用是沿着定軸dim(1),按照Index(conf_t.view(-1, 1))取出元素
batch_conf.gather(1, conf_t.view(-1, 1))得到矩陣的維度是[batch_size*8732,1],代表每個預測框的最匹配類的置信度
計算所有預測框的置信度損失
'''
loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))


#難例挖掘
loss_c[pos] = 0  #過濾掉positive預測框
loss_c = loss_c.view(num, -1)  #reshape爲[batch_size,8732,1]
_, loss_idx = loss_c.sort(1, descending=True)  #將每張圖片中所有的先驗框 按照置信度從高到低排序
_, idx_rank = loss_idx.sort(1)
num_pos = pos.long().sum(1, keepdim=True) #pos是bool,先將其轉換爲long,在求和,結果是每張圖片中positive預測框的數量

'''
torch.clamp(input, min, max, out=None) → Tensor 將張量的每個元素夾緊到區間[min,max]內,
      | min, if x_i < min
y_i = | x_i, if min <= x_i <= max
      | max, if x_i > max
  
self.negpos_ratio*num_pos的shape是[batch_size,1],下面語句的意思是將每張圖片負樣本的數量限制在pos.size(1)-1
'''
num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1)

#neg是難例挖掘最終得到的負樣本索引,
neg = idx_rank < num_neg.expand_as(idx_rank)

#包含正樣本和負樣本的置信度損失
pos_idx = pos.unsqueeze(2).expand_as(conf_data) 
neg_idx = neg.unsqueeze(2).expand_as(conf_data)

'''
gt()函數是比較兩個同維度張量的對應元素,比如C=A.gt(B),若A[i,j]>B[i,j],則C[i,j]==0
(pos_idx+neg_idx).gt(0)表示該樣本是否是訓練樣本
'''
#挑出預測結果的正負樣本
conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes)

#挑出匹配的gt標籤
targets_weighted = conf_t[(pos+neg).gt(0)]

#置信度損失使用交叉熵
#交叉熵的計算公式爲loss=sum(label*log(predict))
loss_c = F.cross_entropy(conf_p, targets_weighted, size_average=False)

# Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
N = num_pos.data.sum() #batch中所有positive框的總數

loss_l /= N
loss_c /= N
return loss_l, loss_c

 

    回到訓練代碼,前面已經定義了數據集和損失,接着就差不多可以開始訓練了,剩餘的訓練代碼如下:

train.py   train()

net.train()#訓練模式

# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')

epoch_size = len(dataset) // args.batch_size
print('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args)

step_index = 0
if args.visdom:
    vis_title = 'SSD.PyTorch on ' + dataset.name
    vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
    iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
    epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)

#定義數據迭代器
data_loader = data.DataLoader(dataset, 
  args.batch_size,
  num_workers=args.num_workers,#處理數據集的線程數
  shuffle=True, 
  collate_fn=detection_collate,#將一個list的sample組成一個mini-batch的函數
  pin_memory=True #如果設置爲True,那麼data loader將會在返回它們之前,將tensors拷貝到CUDA中的固定內存(CUDA pinned memory)中.
)

# create batch iterator
batch_iterator = iter(data_loader)

#迭代訓練
for iteration in range(args.start_iter, cfg['max_iter']):
    #添加數據到visdom
    if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
        update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None,'append', epoch_size)
        # reset epoch loss counters
        loc_loss = 0
        conf_loss = 0
        epoch += 1
        
    #學習率調整
    if iteration in cfg['lr_steps']:
    step_index += 1
    adjust_learning_rate(optimizer, args.gamma, step_index)
    
    #加在訓練數據
    images, targets = next(batch_iterator)
    
    if args.cuda:
        images = Variable(images.cuda())
        targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
    else:
        images = Variable(images)
        targets = [Variable(ann, volatile=True) for ann in targets]
        
    t0 = time.time()
    #前向傳播
    out = net(images)
    #反向傳播
    optimizer.zero_grad()
    
    #計算損失函數
    loss_l, loss_c = criterion(out, targets)
    loss = loss_l + loss_c
    loss.backward()
    optimizer.step()
    
    t1 = time.time()
    loc_loss += loss_l.data[0]
    conf_loss += loss_c.data[0]
    if iteration % 10 == 0:
        print('timer: %.4f sec.' % (t1 - t0))
        print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
    if args.visdom:
        update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],iter_plot, epoch_plot, 'append')
    if iteration != 0 and iteration % 5000 == 0:
        print('Saving state, iter:', iteration)
        torch.save(ssd_net.state_dict(), 'weights/ssd300_COCO_' +repr(iteration) + '.pth')

torch.save(ssd_net.state_dict(),args.save_folder + '' + args.dataset + '.pth')


def adjust_learning_rate(optimizer, gamma, step):
    """Sets the learning rate to the initial LR decayed by 10 at every
        specified step
    # Adapted from PyTorch Imagenet example:
    # https://github.com/pytorch/examples/blob/master/imagenet/main.py
    """
    lr = args.lr * (gamma ** (step))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

      這部分代碼沒有什麼太多需要解釋的。

       本篇博客詳細解釋了先驗框計算損失函數計算模型的訓練

 

 

 

 

 

 

 

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章