Yolo2实现笔记（pytorch）

基于mileiston的github仓库，感谢m同学的无私分享。

 

 

label.py的制作逻辑

1.根据人工预设的几组year（如果2007,2012等）和预设的trainset（如 'train','val'），和预设的格式，组装出{trainset}.txt所在路径。

2.根据获得的路径，读取该pascalvoc数据集的{trainset}.txt文件，里面存着list of xml_paths。

3.根据获得的list of xml_paths，传入parser解析器（brambox.parser）。每个xml里面有img_id，annos等信息。每个anno里面又包含class和bbox等信息。

4.返回dict{'img_id': list_of_annos}

5.保存该dict为pkl文件

【brambox是lightnet的作者搞的另一个工具，解析几个常用数据集的文件。】

 

train.py里

1.根据传入的argument里的model_name，调用initEnv()方法

2.在initEnv()中，

->2.1先调用getConfig()，根据model_name进行yaml.load({对应modelname.yaml})，返回cur_cfg，是个yaml字典。

->2.2 向cur_cfg字典内加入 model_name 。

cur_cfg['model_name'] = model_name

->2.3根据yaml里的参数组合出三个dir路径。

work_dir = os.path.join(root_dir, model_name, version)
backup_dir = os.path.join(work_dir, backup_name)
log_dir = os.path.join(work_dir, log_name)

>建立对应的dir文件夹。

>再把backup_dir添加到cur_cfg字典中。(trainflag==1)

->2.4不明意义的gpu环境设置

gpus = cur_cfg['train']['gpus'] 

os.environ['CUDA_VISIBLE_DEVICES'] = gpus

->2.5完成上述工作后，调用combineConfig(cur_cfg, train_flag)。

这个函数的作用是根据trainflag，例如=1时，把test和speed都去掉，把train下面的参数的层级拔高一层。

举例：原来是cfg['train']['trainset']=xxx，现在是cfg['trainset']=xxx。

->2.6 initEnv() end，得到config字典。

 

3.根据config字典，实例化一个HyperParams类的对象。

->3.1 基本把config里面的所有{key:value}结构都转变成了self.key = value的结构。

->3.2 新增了几个attr，这是用户config的yaml中所不包含的。

self.jitter = 0.3
self.flip = 0.5
self.hue = 0.1
self.sat = 1.5
self.val = 1.5

self.rs_steps = []
self.rs_rates = []

->3.3 自适应self.cuda = True or False

->3.4根据cfg['labels']这个list，获得类别数量

self.classes = len(self.labels)

->3.5 HyperParams end。

主要是把key变成了attr，方便调用。又隐藏了几个参数，对用户不可见。

经过此步骤，我们获得hyper_params。

 

4.根据hyper_params，实例化一个VOCTrainingEngine类对象。

->4.1 部分attr赋值为self.attr

->4.2 搭建网络

根据model_name，在本地包models.__dict__内找到对应的model类。

net = models.__dict__[model_name](hyper_params.classes,
                hyper_params.weights,     # pre-trained文件路径
                train_flag=1,             # 1表示train mode ,2 表示test mode
                clear=hyper_params.clear) # clear表示是否重置self.seen

#本例中，是找到了class Yolov2(YoloABC):这个类，定义了backbone和head两部分。
#self.backbone = backbone.Darknet19()
#self.head = head.Yolov2(num_anchors=len(anchors_mask[0]), num_classes=num_classes)

if self.cuda:
    net.cuda()

 

->4.3 定义optimizer，使用SGD

log.debug('Creating optimizer')
learning_rate = hyper_params.learning_rate
momentum = hyper_params.momentum
decay = hyper_params.decay
batch = hyper_params.batch

optim = torch.optim.SGD(net.parameters(), 
                lr=learning_rate/batch,
                momentum=momentum, 
                dampening=0, 
                weight_decay=decay*batch)

 

->4.4 定义dataloader

#VOCDataset是torch.utils.data.dataset的子类
dataset = VOCDataset(hyper_params)

#引入了4个不同的transform操作
#
#rf  = data.transform.RandomFlip(flip)
#rc  = data.transform.RandomCropLetterbox(self, jitter)
#hsv = data.transform.HSVShift(hue, sat, val)
#it  = tf.ToTensor()     # torchvision.transforms as tf
#
#img_tf = data.transform.Compose([rc, rf, hsv, it])
#anno_tf = data.transform.Compose([rc, rf])
#
#这样，在每次loader调用__getitem__时，
#都会根据概率0.5，来决定返回(raw_img，raw_anno)
#还是返回 img_tf和anno_tf之后的(img,anno)
#
#@Dataset.resize_getitem
#def __getitem__(self, index):
#


#data.Dataloder是torch.utils.data.dataloader的子类
#
#使用自定义的子类，主要是为了实现的对不同尺寸图片的训练。
#def change_input_dim(self, multiple=32, random_range=(10, 19), finish=False):
#显然，input_dim可以在320~608之间以32的步长来切换。

dataloader = data.DataLoader(
            dataset,
            batch_size = self.mini_batch_size,
            shuffle = True,
            drop_last = True,
            num_workers = hyper_params.nworkers if self.cuda else 0,
            pin_memory = hyper_params.pin_mem if self.cuda else False,
            collate_fn = data.list_collate,
        )

 

->4.5 用父类engine.Engine的初始化方法。

super(VOCTrainingEngine, self).__init__(net, optim, dataloader)

其实就是把4.2的net，4.3的optim，4.4的dataloader传进去了，串成了self的成员对象。

 

->4.6 为之后设置loss function进行前置工作。

self.nloss = self.network.nloss  #after super,self.network=net ,so this equals net.nloss, since we have 5 anchors, so got 5 nloss

#预定义一个list of (dict of train_loss)
self.train_loss = [{'tot': [], 'coord': [], 'conf': [], 'cls': []} for _ in range(self.nloss)]

 

->4.7 VOCTrainingEngine 实例化完成。

结构还是很简单的，追根究底也只有三大块， 4.1，(4.2+4.3+4.4)->4.5，4.6。

结束后，我们得到实例对象eng。

 

5.获得eng.batch

因为我们可能会从中断处重新开始训练，所以获取这个还是有用的，可以在结束后告诉我们，本次训练从第n batch 运行到了第m batch。

#self.batch定义在父类engine.Engine中

    @property
    def batch(self):
        """ Get current batch number.

        Return:
            int: Computed as self.network.seen // self.batch_size
        """
        return self.network.seen // self.batch_size

self.batch是一个@property修饰的属性方法，计算方式是self.network.seen // self.batch_size。含义是当前处于第几批次。

解释一下。

self.network就是上文的net，本质为Yolov2类，继承自YoloABC->继承自Darknet->继承自lightnet类。

network里面含有backbone和head两层，还定义了_forward等相关的方法，不必细究。

network.seen的含义是，目前为止network处理了多少张图片。

底层的lightnet类的forward(self,x)方法里，每次会更新 self.seen += x.size(0) 。

我们再用处理了n张图片//batch_size，来计算当前处于第几批次。

 

例如，原代码中，取batch_size=64，mini_batch_size=16。每64张图片为1个batch。

这样每次loader给出16张图片和相应annos。每次network.forward之后，seen += 16。

当seen=16时，计算self.batch = 16//64 = 0 ，处于第0批次。

当seen=64时，计算self.batch = 64//64 = 1，处于第1批次。

以此类推。

 

6. 启动eng()，开始！

->6.1 运行eng()，会调用父类engine.Engine中定义的 __call__()方法, 在__call__中处理所有标准的训练步骤。

->6.2 先进行self.start()，这个函数由子类VOCTrainingEngine实现。

主要效果是取出一些self.hyper_params里的参数，然后用self.add_rate方法加入到自身的self._rates字典中。

做这步是因为，我们希望有阶梯式的learning_rate， 比如在400 batch的时候用中规中矩的0.0001来初训练，在1000 batch的时候用较大的0.001防止掉入局部最优解，在大后期60000 batch的时候用极小的0.00001来微调。

#start()由子类VOCTrainingEngine负责实现。
#lr_steps: [400,700,900,1000, 40000,60000] 
#lr_rates: [0.0001,0.0005,0.0005,0.001, 0.0001,0.00001] 

    def start(self):
        log.debug('Creating additional logging objects')
        hyper_params = self.hyper_params

        lr_steps = hyper_params.lr_steps
        lr_rates = hyper_params.lr_rates

        bp_steps = hyper_params.bp_steps
        bp_rates = hyper_params.bp_rates
        backup = hyper_params.backup    #cfg['backup_interval']

        rs_steps = hyper_params.rs_steps
        rs_rates = hyper_params.rs_rates
        resize = hyper_params.resize    #cfg['resize_interval']

        self.add_rate('learning_rate', lr_steps, [lr/self.batch_size for lr in lr_rates])
        self.add_rate('backup_rate', bp_steps, bp_rates, backup)
        self.add_rate('resize_rate', rs_steps, rs_rates, resize)

        self.dataloader.change_input_dim()



#add_rate()定义在父类engine.Engine中
#效果为，把预设的阶梯learning_rate等加入到self.__rates字典中。
self.__rates[name] = (steps, values)

->6.3 self._update_rates()

因为可能是从中断处开始的，self.batch可能是一个比较大的值，所以保险起见update一下，恢复到断点处的rates。

#定义在父类engine.Engine中
#核心代码就这么两句，如果到了我们需要的batch，就切换下一个阶段的rate

if self.batch >= steps[i]:
    new_rate = values[i]
else:
    break

 ->6.4 开始一个循环，while True:

不用担心结束的问题。

yaml中有max_batches，每个batch训练完之后都会检查一下self.batch是否超出这个值。

超出就会保存weights之后quit。

或者人为地ctrl+c发送中断signal也会quit。

->6.5 在循环内，不断遍历loader里的全部mini_batch。本例一次mini_batch=16张图片。

loader = self.dataloader
for idx, data in enumerate(loader):

    ...

    if (len(loader) - idx) <= self.batch_subdivisions:
        break

#且我们会对尾数进行抛弃
#batch_subdivisions = self.batch_szie / self.mini_batch_size
#本例 = 64/16 = 4。 即4个mini_batch等于一个batch。
#当loader内的mini_batch不足4个，即不足以构成一个batch时，予以drop

->6.6 遍历过程中，对每个mini_batch，先调用self.process_batch(data)，定义在子类中。

先把data移到cuda内存区，然后用已经在cuda上的network计算loss。

最后把loss细分成 total_loss,coord_loss（坐标损失），confidence_loss（确信有物体），class_loss（分类）等，存入self.train_loss这个list of dict中。

def process_batch(self, data):
    data, target = data
    # to(device)
    if self.cuda:
        data = data.cuda()

    # network.forward function defined in _lightnet.py ,
    # firstly we run self._forward(x),
    # then we calculate loss for each self.loss, which is a list of loss_functions from loss.RegionLoss
    loss = self.network(data, target)
    loss.backward()

    for ii in range(self.nloss):
        self.train_loss[ii]['tot'].append(self.network.loss[ii].loss_tot.item() / self.mini_batch_size)
        self.train_loss[ii]['coord'].append(self.network.loss[ii].loss_coord.item() / self.mini_batch_size)
        self.train_loss[ii]['conf'].append(self.network.loss[ii].loss_conf.item() / self.mini_batch_size)
        if self.network.loss[ii].loss_cls is not None:
            self.train_loss[ii]['cls'].append(self.network.loss[ii].loss_cls.item() / self.mini_batch_size)

关于network的具体forward()方法，另寻一处说明。

总而言之，无论是什么网络，究其根底都不过是传入data和target计算loss然后backward罢了。

 

->6.7遍历过程中，对每个mini_batch，再调用self.train_batch()，定义在子类中。

主要是进行了optimizer.step()。这步的作用是根据6.6中loss.backward()算出的梯度，更新网络参数network.parameters()。

然后顺便的进行了收尾工作。

计算loss，print一下当前第几batch了，loss是多少。

每backup_rate，save_weights()一下weights_{batch}.pt。

固定每100batch，save_weights()一下backup.pt。

同时检查是否超出max_batches，若是则设置finish_flag=True。

 

->6.8 遍历过程中，对每个mini_batch，第三步是检查是否要停止，若不需要停止就self._update_rates()。

在前面的process_batch过程中，经过network算出了loss，如上文所述，network.seen会增加。

所以结束了process_batch()和train_batch()之后，self.batch数值可能会变动。此时update一下，就可以更新我们的阶梯速率。

 

->6.9 结束了 eng()，即 __call__() 过程。

 

最后附上__call__的代码，结合上面的文字来看。

    def __call__(self):
        """ Start the training cycle. """
        self.start()            #to add rates, implemented in subclass
        self._update_rates()    #then update rates
        if self.test_rate is not None:
            last_test = self.batch - (self.batch % self.test_rate)

        log.info('Start training')
        self.network.train()
        while True:
            loader = self.dataloader
            for idx, data in enumerate(loader):
                print('new loaded batch is ', self.batch)
                # Forward and backward on (mini-)batches
                self.process_batch(data)    # implemented in subclass, got self.train_loss as a list of dict of tot,coord,conf,cls
                if (idx + 1) % self.batch_subdivisions != 0:
                    continue

                # Optimizer step
                self.train_batch()          # implemented in subclass

                # Check if we need to stop training
                if self.quit() or self.sigint:
                    log.info('Reached quitting criteria')
                    return

                # Check if we need to perform testing
                if self.test_rate is not None and self.batch - last_test >= self.test_rate:
                    log.info('Start testing')
                    last_test += self.test_rate
                    self.network.eval()     #set to evaluation mode
                    self.test()
                    log.debug('Done testing')
                    self.network.train()    #set to train mode

                # Check if we need to stop training
                if self.quit() or self.sigint:
                    log.info('Reached quitting criteria')
                    return

                # Automatically update registered rates
                self._update_rates()

                # Not enough mini-batches left to have an entire batch
                if (len(loader) - idx) <= self.batch_subdivisions:
                    break

 

 

7. 结束了整个train过程。记录一下耗时和batch变动。

    # run eng
    b1 = eng.batch
    t1 = time.time()
    eng()
    t2 = time.time()
    b2 = eng.batch

    log.info(f'\nDuration of {b2-b1} batches: {t2-t1} seconds [{round((t2-t1)/(b2-b1), 3)} sec/batch]')