caffe小工具（分析log画输出折线）

背景

前段时间玩了一下tensorflow，发现tensroflow有一个十分实用的工具叫做tensorboard。在tensorboard的帮助下我们可以在训练的过程中监控程序的输出（通过添加summary），这样能够帮助我们在模型无法按理想的情况工作时更好的分析原因。
最近的实验又用回了caffe，在用caffe的过程中发现整个训练过程真是又长又臭，除了面对枯燥的输出以外几乎什么都做不了。如果想要将loss，accuracy等信息画出来的画只能自己控制训练过程，并在训练过程中记录loss和accuracy等输出。
简单查了下资料，发现caffe有提供现成的工具，通过分析训练输出的log来画折线图。感觉这个思路挺不错，但是有一点不太满意的就是，只能在训练完或者训练中的某个时刻，调用该工具来画出折线图。于是就想着，反正也是分析log，那自己写一个python脚本，隔1分钟就分析一次log重画一次折线图不就可以达到类似tensorboard的那种伪实时监控嘛。于是就在一个等待模型训练的下午写了下这个小工具。

思路

这个小工具主要有三部分组成：
1. 分析log，从整个日志文件中取出每次迭代的输出。这里将一次迭代的输出称为一个block；
2. 分析block，从block中提取出所有输出值得名称以及数值；
3. 将从每一个block中提取到的输出值以折线图的形式画出。
具体代码如下（只是根据自己的需要写的简单脚本，鲁棒性可能不太好）：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2017-04-04 15:48:14
# @Author  : Raymond Wong

'''
this script is meant to parse the caffe log and plot the output in order to track the training process
and find out problems as soon as possible
'''

import os
import sys

import matplotlib.pyplot as plt
import argparse
import re
import math
import numpy as np
import time
import datetime

ITERATION_START_PATTERN = re.compile(r'.*Iteration \d*, loss = \d+\.{0,1}\d*')
ITERATION_END_PATTERN = re.compile(r'.*Iteration \d*, lr = \d+\.{0,1}\d*')
ITERATION_PATTERN = re.compile(r'Iteration \d+,')
OUTPUT_PATTERN = re.compile(r'[a-zA-Z_]+[\w@]+ = \d+\.{0,1}\d*')

def next_line_index(lines, pattern, start=0):
  for i in xrange(start, len(lines)):
    line = lines[i].strip()
    match = pattern.search(line)
    if match:
      return i
  return -1

def parse_block(block):
  iter_line = block[0]
  iter_match = ITERATION_PATTERN.search(iter_line)
  iter_num = 0
  if iter_match:
    iter_match = iter_match.group()
    iter_num = int(iter_match[:-1].split(' ')[-1])
  else:
    return None
  ret = dict()
  for line in block[:-1]:
    match = OUTPUT_PATTERN.search(line)
    if not match:
      continue
    name, value = map(lambda x:x.strip(), match.group().split('='))
    ret[name] = [iter_num, float(value)]
  return ret

def plot(name, data):
  plt.grid(True)
  plt.title(name + ' v.s. iter')
  data = np.array(data, dtype="float32")
  x = data[:, 0].astype(np.int)
  y = data[:, 1]
  if name.find('loss') >= 0:
    y = np.log(y + 1)
  plt.plot(x, y, 'x-')

if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('log', help="log file path")
  parser.add_argument('refresh', help="refresh(>0) or just print(<=0)")
  args = parser.parse_args()
  if int(args.refresh) > 0:
    plt.ion()
  while True:
    print 'refresh at ' + datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    infile = open(args.log, 'r')
    lines = infile.readlines()
    infile.close()
    start = next_line_index(lines, ITERATION_START_PATTERN, 0)
    output = dict()
    while start >= 0:
      end = next_line_index(lines, ITERATION_END_PATTERN, start)
      if end < start:
        break
      block = lines[start:end + 1]
      c_output = parse_block(block)
      for key in c_output.keys():
        if not output.has_key(key):
          output[key] = list()
        output[key].append(c_output[key])
      start = next_line_index(lines, ITERATION_START_PATTERN, end)
    total = len(output)
    row = int(math.ceil(math.sqrt(total)))
    col = int(math.ceil(total / row))
    plt.clf()
    for idx, key in enumerate(output.keys()):
      plt.subplot(row, col, idx + 1)
      plot(key, output[key])
    if int(args.refresh) <= 0:
      plt.show()
      break
    plt.draw()
    plt.pause(60)
  print 'Quit!'

简单说明

简单解释一下。在main函数中我们遍历整个日志文件，调用next_line_index方法使用正则匹配匹配到下一个block。然后调用parse_block 方法从该block中分析出当前迭代的次数以及迭代的输出。将parse_block 的返回值存入记录所有迭代输出的字典中，直至分析完所有的block。逐一对输出字典的每个key调用plot函数，画出这个输出的折线图。然后等待60秒，再重新分析一次日志。
现在的实现是每次刷新的时候都从头开始分析log文件，但其实是很没有必要的，因为每次刷新时，前面迭代的输出都是相同的，所以我们只要分析并记录当前刷新中新增的block就可以了。另外就是，现在我只考虑了画出train net的输出（正则表达式过滤掉了test net的输出）。有时间的话可以把test net的分析加上。

使用方法

1. 训练的过程中将caffe的日志保存下来，具体保存方法有很多，caffe的命令行有参数，也可以借助linux的管道等实现（以train.log为例）
2. 调用脚本python parse_plot_log.py train.log 1 来对输出文件进行监听。如果是已经训练完想要画出一次日志文件，则将最后的1改成0，即python parse_plot_log.py train.log 0