TensorBoard对YOLOv5可视化训练，稀疏化bn_weight 不显示

使用的是博主（早茶和猫）开源项目：YOLOv5旋转目标框模型剪枝

传送门：yolov5_obb旋转框检测（kld loss、probloss）,剪枝，跟踪_yolov5obb-CSDN博客

问题：

一、该问题解决的前提：

解决正常和稀疏化训练没有results.csv文件（注：YOLOv5s训练未生成result文件-CSDN博客）

前言：

       之前采用YOLOv5剪枝（地址：GitHub - midasklr/yolov5prune）稀疏化训练时，可以通过TensorBoard可视化BN权值层分布。但是，在使用旋转目标框进行训练时，对BN层进行可视化失败了。

敲黑板：在这之前请确保runs训练的文件下有results.csv文件，这个可以参考上一篇（YOLOv5s训练未生成result文件-CSDN博客）

二、解决办法：

train_sparity.py文件中修改416行，修改前：

# Update best mAP
fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
if fi > best_fitness:
    best_fitness = fi
log_vals = list(mloss) + list(results) + lr
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)

在callbacks中添加bn_weights，添加后

# Update best mAP
    fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    if fi > best_fitness:
        best_fitness = fi
    log_vals = list(mloss) + list(results) + lr
    callbacks.run('on_fit_epoch_end', log_vals, bn_weights.numpy(), epoch, best_fitness, fi)

出现报错

找到utils/loggers/_init_.py文件，在on_fit_epoch_end函数中，添加bn_weights，

修改前：

def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
    # Callback runs at the end of each fit (train+val) epoch
    x = {k: v for k, v in zip(self.keys, vals)}  # dict
    if self.csv:
        file = self.save_dir / 'results.csv'
        n = len(x) + 1  # number of cols
        #s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n')  # add header
        s = '' if file.exists() else (('{:<20}' * n).format(*['epoch'] + self.keys)).rstrip(',') + '\n'

        with open(file, 'a') as f:
            f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')

    if self.tb:
        for k, v in x.items():
            self.tb.add_scalar(k, v, epoch)

    if self.wandb:
        self.wandb.log(x)
        self.wandb.end_epoch(best_result=best_fitness == fi)

修改后：

def on_fit_epoch_end(self, vals, bn_weights, epoch, best_fitness, fi):
    # Callback runs at the end of each fit (train+val) epoch
    x = {k: v for k, v in zip(self.keys, vals)}  # dict
    if self.csv:
        file = self.save_dir / 'results.csv'
        n = len(x) + 1  # number of cols
        # s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n')  # add header
        s = '' if file.exists() else (('{:<20}' * n).format(*['epoch'] + self.keys)).rstrip(',') + '\n'

        with open(file, 'a') as f:
            f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')

    if self.tb:
        for k, v in x.items():
            self.tb.add_scalar(k, v, epoch)
        self.tb.add_histogram('bn_weights/hist', bn_weights, epoch, bins='doane')

    if self.wandb:
        self.wandb.log(x)
        self.wandb.end_epoch(best_result=best_fitness == fi)

再次训练，并输入tensorboard指令可视化训练，为了防止对tensorboard不是很熟悉小伙伴在这里出错，也多说两句首先使用指令ipconfig检查下电脑的ip地址(注：这边是192.168.0.105)。

之后回到项目，在terminal窗口，首先cd到runs目录下，然后输入指令：tensorboard --logdir ./ --host=192.168.0.105获取连接。

点击连接或者复制连接到浏览器中打开，惊喜！

三、结语

哈哈，多唠一句，希望大家科研和工作顺利，在各自领域发光发热吧，加油！！！