yield生成器之keras模型训练代码实现

1 . 何为yield?

包含yield的函数,会生成一个生成器generator ,听过next函数可以不断生成一批数据.

示例:

>>> def test():
...     testData = ['a' , 'b' , 'c' , 'd' , 'e' ,'f']
...     for x in testData:
...         yield x
...         print ("test----"+x)
... 
>>> Generator = test()
>>> x = Generator.next()
>>> x
'a'
>>> x = Generator.next()
test----a
>>> x
'b'
>>> x = Generator.next()
test----b
>>> x
'c'

可以看出,当第一次调用next的时候,函数test返回第一个字母a,然后便停止运行,当第二次调用next的时候,函数从第一次中断的地方继续往下运行,于是便打印出test----a,接着继续循环,遇到yield返回b,以此循环下去,直到for循环结束.

 

2 . next方法

编写的生成器,需要使用next去调用程序,并获取返回的数据,要注意的是,python2使用next的方法是Generator.next().python3使用next的方法是next(Generator)

 

3 . 如何使用yield编写生成器

因为通过yield我们可以不断返回数据,并且可以保持上次调用next的记录点,那么我们可以通过一个while(1)循环,不断产生数据.生成的生成器可直接通过next循环调用,或者直接传入keras的model,使用model.fit_generator(generator=train_generator)方法训练模型.

代码格式如下:

###
#  imgLabelDict为词典dict , 里面存放数据格式如下{key(图片文件名) : value(图片对应的label)}
#  imgNameSet为所有图片的文件名集合
#  batch_size为每一个batch图片大小
#  img_num为图片总共的张数
###
def Generator(imgLabelDict , imgNameSet ,  batch_size , img_num):
    i = 0
    labelGenerator = []
    imgGenerator = []
    num = 0
    while 1:
        #当迭代完所有图片时,可对图片进行shuffle操作并新一轮迭代.
        if num == img_num or num == 0:
            #shuffle
            num = 0
            
        #get the img name
        img_name = imgNameSet[num]
        
        #labelProcess
        label = labels.get(img_name)
        labelGenerator.append(label)
        
        #imgProcess
        img = imread(img_name)
        img = preprocess_input(img)
        imgGenerator.append(img)

        if i>=batch_size-1:
            yield np.array(imgGenerator) ,np.array(labelGenerator)
            i = -1
            labelGenerator = []
            imgGenerator = []
        i+=1
        num += 1
        
        
trainGenerator = Generator(imgLabelDict , imgNameSet ,  batch_size , img_num)
trainBatchImg = trainGenerator.next()

 

存在问题:

因为使用到while(1)操作,所以会占用CPU资源比较高,具体原因可以参考这博客:https://blog.csdn.net/S1amDuncan/article/details/78840031