Tf-pose-estimation: Training speed improvements

[ildoonet]

• Data Augmentation is Slow
• Memory Consumption
• Cpu Load too high

[ildoonet]

Numba

• 1.6x-1.7x speed up
• Use less subprocesses
  
  
  
  • less memory consumption
  
  
  • less cpu load
  
  

TODO

• some hangups before writing logs
• some hangups before calculating validation logs

[skyler1253]

hey @ildoonet - I'm able to step through training quite a bit quicker in the dev/architecture-mobilenet2 branch (also training on mobilenet_v2). I'm about to normalize the inputs as described in #369, but i seem to get hung up on the sess.run() calls happening every 500 steps:

sess.run([total_loss, total_loss_ll, total_loss_ll_paf, total_loss_ll_heat, learning_rate, merged_summary_op, enqueuer.size()])

These calls seem to be executing on a single CPU core and it takes a good 20 minutes. Any ideas on how this might be sped up?

[filipetrocadoferreira]

@staticmethod
    def put_vectormap(vectormap, countmap, plane_idx, center_from, center_to, threshold=8):
        _, height, width = vectormap.shape[:3]

        vec_x = center_to[0] - center_from[0]
        vec_y = center_to[1] - center_from[1]

        min_x = max(0, int(min(center_from[0], center_to[0]) - threshold))
        min_y = max(0, int(min(center_from[1], center_to[1]) - threshold))

        max_x = min(width, int(max(center_from[0], center_to[0]) + threshold))
        max_y = min(height, int(max(center_from[1], center_to[1]) + threshold))

        norm = math.sqrt(vec_x ** 2 + vec_y ** 2)
        if norm == 0:
            return

        vec_x /= norm
        vec_y /= norm

        x = np.arange(min_x,max_x)
        y = np.arange(min_y,max_y)

        bec_x = x - center_from[0]
        bec_y = y - center_from[1]

        dist = np.abs(bec_x * vec_y - bec_y[:,np.newaxis] * vec_x)

        countmap[plane_idx][min_y:max_y,min_x:max_x][dist<=threshold] += 1
        vectormap[plane_idx * 2 + 0][min_y:max_y,min_x:max_x][dist<=threshold] = vec_x
        vectormap[plane_idx * 2 + 1][min_y:max_y,min_x:max_x][dist<=threshold] = vec_y `

@staticmethod
    def put_heatmap(heatmap, plane_idx, center, sigma):
       center_x, center_y = center
        _, height, width = heatmap.shape[:3]
        kp_heat_proj = np.zeros_like(heatmap[plane_idx])


        th = 4.6052

        kp_height, kp_width = kp_heat.shape[:2]

        x0 = int(np.clip(center_x - kp_width / 2, 0, width))
        y0 = int(np.clip(center_y - kp_height / 2, 0, height))

        x1 = int(np.clip(center_x + kp_width / 2, 0, width))
        y1 = int(np.clip(center_y + kp_height / 2, 0, height))

        #print([y0,y1,x0,x1])
        #print([y0-center_y+kp_height//2,y1-center_y+kp_height//2,x0-center_x+kp_height//2,x1-center_x+kp_height//2])

        kp_heat_proj[y0:y1,x0:x1] = kp_heat[y0-center_y+kp_height//2:y1-center_y+kp_height//2,x0-center_x+kp_height//2:x1-center_x+kp_height//2]

        kp_heat_proj[y0:y1,x0:x1] = np.max(np.dstack([kp_heat_proj[y0:y1,x0:x1],heatmap[plane_idx][y0:y1,x0:x1]]),axis=-1)
        kp_heat_proj[kp_heat_proj>1.0] = 1.0


        heatmap[plane_idx][y0:y1, x0:x1] = kp_heat_proj[y0:y1,x0:x1]

This numpy-based versions of put_vectormap() and put_vectormap() really speedup pre-processing.

[ildoonet]

@filipetrocadoferreira Thanks. I will test it. With numba, it will be faster.