Skip to main content
 首页 » 编程设计

python之Tensorflow,大图像推理之内存不足

2024年08月15日31www_RR

我训练了一个超分辨率模型。推理有效:

def inference(self, filename): 
    img_in = misc.imread(filename) / 255.0 
    res = self.sess.run(self.out, feed_dict={ self.x: [img_in], self.is_training: False, self.lr_input: 0.0 })[0] * 255.0 
    image = misc.toimage(res, cmin=0, cmax=255) 
    fn_res = filename.replace(".png", "_result.png").replace(".jpg", "_result.jpg") 
    misc.imsave(fn_res, image) 

但是当图像大于 600x600px 时,它说:

terminate called after throwing an instance of 'std::bad_alloc' 
what():  std::bad_alloc 
Aborted (core dumped) 

我使用 16GB RAM 和 16GB 交换文件,使用 CPU 进行推理。但是在 Ubuntu 16.04 中设置的交换文件尽管足够了,但无济于事。

我知道我可以通过将图像拆分成多个部分并按顺序处理它们来解决这个问题。但是在输出中有明显的 split 痕迹。

为什么不起作用?有什么方法可以使用 Python 或 Tensorflow 实现内存交换文件吗?还有其他解决方法吗?

请您参考如下方法:

通过将图像分成几 block ,一张一张地处理并组合它们来解决它。为了避免输出图像中出现可见的“分割线”,我用不同的 block 尺寸重复上述过程并计算平均值。

我最终得到了这个:

def inference(self, filename, self_ensemble=False): 
    img_q = Image.open(filename) 
 
    scale = 4 
 
    if img_q.width > 400: 
        total = np.zeros((img_q.height * scale, img_q.width * scale, 3)) 
        #repeat 7 times with different piece size 
        for x in range(7): 
            total += self.inference_inter(filename, scale, 32 * (x + 8), ensemble=self_ensemble) 
 
        total = total / 7.0 
 
        image = misc.toimage(total, cmin=0, cmax=255) 
        fn_res = filename.replace(".png", "_result.png").replace(".jpg", "_result.jpg") 
        misc.imsave(fn_res, image) 
    else: 
        img_in = misc.imread(filename) / 255.0 
        res = self.sess.run(self.out, feed_dict={ self.x: [img_in], self.is_training: False, self.lr_input: 0.0 })[0] * 255.0 
        image = misc.toimage(res, cmin=0, cmax=255) 
        fn_res = filename.replace(".png", "_result.png").replace(".jpg", "_result.jpg") 
        misc.imsave(fn_res, image) 
 
 
def inference_inter(self, filename, scale, pat_size, ensemble=True): 
    img_in = misc.imread(filename) / 255.0 
    img_q = Image.open(filename) 
 
    res = np.zeros((img_q.height * scale, img_q.width * scale, 3)) 
 
    for qx in range(0, img_q.width, pat_size): 
        for qy in range(0, img_q.height, pat_size): 
            res[(qy * scale):((qy + pat_size) * scale)][(qx * scale):((qx + pat_size) * scale)] = self.sess.run(self.out, feed_dict={ self.x: [img_in[qy:(qy + pat_size)][qx:(qx + pat_size)]], self.is_training: False, self.lr_input: 0.0 })[0] * 255.0 
 
    return res