深度学习中动态加载样本

发表于 | 分类于
字数统计: | 阅读时长 ≈

在用CNN处理图像的时候,经常我们不是用一整张完整的图像去训练,而是经常会采用取一个图像块也就是patch来输入网络进行训练,构成我们的训练集样本和测试集样本.

一种常用的方法就是先把他们保存到本地硬盘中,然后每次训练时一张张去读取(或者将他们转换为tfrecord格式使用TensorFlow中的队列读取这些文件),最近在做3D图像的操作,直接对3D图像进行卷积,卷积核也是3维的,因此图像尺寸不可能取到很大,也是voxel patch.

但是3D数据是较大的,硬盘空间有时候没那么大,尤其是固态,因为固态下的文件读取速度要比机械盘快很多,所以能放到固态下的话是个不错的选择,但是对于这个问题来说,完全不够,只能放到机械盘下面.

还有一种想法就是: “把鸡蛋放到一个篮子里.🥚🥚🥚”

就是在训练时把所有数据一次性读取加载到内存中,然后训练的时候从这些图片中随机取样本,要求就是这些图片可以被同时读取到内存中,需要一定大小的内存,尤其是在一些用来跑程序的主机上,内存占用很低.

Example

例如可以把要提取的patch的点以pickle文件保存为list对象,用pytorch实现为:

将所有的读取到图像以及要提取样本的关键点传入MyDataset类, 在__getitem__方法中对已有的图像进行切片就能得到需要的样本.

这样的好处就是节省硬盘空间,从内存获取数据速度较快.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from torch.utils.data import Dataset
import numpy as np
import pickle

class MyDataset(Dataset):
	def __init__(self, src_img_list, std_img_list, points_pickle):	
    	self.src_img_list = src_img_list
        self.std_img_list = std_img_list
        with open(points_pickle, 'rb') as f:
              self.points = pickle.load(f)
  
    def __len__(self):
        return len(self.points)
  
    def __getitem__(self, idx):
        i, c, h, w = self.points[idx]
        data = self.src_img_list[i][c-64:c+64, h-64:h+64, w-64:w+64].reshape([1, 128, , 128]).
        target = self.std_img_list[i][c-64:c+64, h-64:h+64, w-64:w+64].reshape([1, , 128, 128]
        return data, target