Nightysky
Nightysky

Reputation: 57

Out of memory when extracting training images features from VGG16 pretrained model

I've been trying some tensorflow tutorial.
Now I have 3539 images in size of 224*224, and then split them into train and test sets.
Then, I used VGG16 pre-trained model to extract features, but I got the following error.

enter image description here I had looked for lots of solutions saying that I have to reduce the batch size, but I don't know how to do it and where to do it.
Anyone can help? Here is my code:

import keras
from keras.models import Sequential
from keras.applications.vgg16 import VGG16
from keras.layers import Dense, InputLayer, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D, GlobalMaxPooling2D
from keras.preprocessing import image
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm
from sklearn.model_selection import train_test_split

train=pd.read_csv('output/train_new.csv')
train_image=[]

for i in tqdm(range(train.shape[0])):
    img=image.load_img('train_1/'+train['image'][i], target_size=(224,224,3))
    img=image.img_to_array(img)
    img=img/255
    train_image.append(img)

X=np.array(train_image)
X.shape

y=train['class']
#creating the training and validation set
X_train, X_test, y_train, y_test=train_test_split(X, y, random_state=42, test_size=0.2, stratify=y)

#creating dummies of target  variable for train and validation set
y_train=pd.get_dummies(y_train)
y_test=pd.get_dummies(y_test)

# creating the base model of pre-trained VGG16 model
base_model=VGG16(weights='imagenet', include_top=False)

# extracting features for training frames
X_train=base_model.predict(X_train)
X_train.shape

Upvotes: 0

Views: 327

Answers (1)

Gerry P
Gerry P

Reputation: 8102

In this part of your code you are loading all 3539 images of size 224 X 224 X 3 which takes up a massive amount of memory.

for i in tqdm(range(train.shape[0])):
    img=image.load_img('train_1/'+train['image'][i], target_size=(224,224,3))
    img=image.img_to_array(img)
    img=img/255
    train_image.append(img)

What you want to do is to load your data in batches, and the easiest way to do that is to use the Keras ImageDataGGenerator.flow_from_datafames. Documentation for that is here.. Now I am assuming you have all your images in a directory, I will call it the train_1. Your csv file I assume contains an identifier for the image file with a column name train. I assume this is a partial path to the image file. I assume the csv file has a column for the class labels of the file called class. Below is a function to read in the csv file, change the column names to filepaths, labels, make the filepaths column the full path to thee image and then split the dataframe into a train_df, a test_df and a valid_df

def preprocess (sdir, csvpath, trsplit, vsplit):
    df=pd.read_csv(csvpath)
    print (df.head())
    df.columns=['filepaths', 'labels'] # assumes there are only 2 xolumns in dataframe
    df['filepaths']=df['filepaths'].apply(lambda x: os.path.join(sdir,x))  # now filepath column is the full path to the image 
    # split into a train_df, a test_df and a valid_df
    dsplit=vsplit/(1-trsplit)
    strat=df['labels']
    train_df, dummy_df=train_test_split(df, train_size=trsplit, shuffle=True, random_state=123, stratify=strat)
    strat=dummy_df['labels']
    valid_df, test_df= train_test_split(dummy_df, train_size=dsplit, shuffle=True, random_state=123, stratify=strat)
    print('train_df length: ', len(train_df), '  test_df length: ',len(test_df), '  valid_df length: ', len(valid_df))
    classes=list(train_df['labels'].unique()) 
    class_count =len(classes)  
    print ('below is a list of image samples in each class to evaluate train_df.balance')
    print(list(train_df['labels'].value_counts()))
    return train_df, test_df, valid_df

Now call the function

sdir=r'c:\train_1'
csvpath=r''output/train_new.csv'
trsplit=.9  # percent of images to use for training
vsplit=.05 # percent of images to use for validation NOTE test split is then 1-.9=.05=.05
train_df, test_df, valid_df= preprocess(sdir,csvpath,trsplit, vsplit)

Now create three ImageDataGenerators a train_gen, a test_gen and a valid_gen. Code below creates these generators

channels=3
batch_size=20 # set batch size based on model complexity and sie of images
img_shape=(img_size[0], img_size[1], channels)
# calculate test_batch_size and test_steps so that test_batch_size X test_steps = number of test images
# this ensures you go through the test set exactly once when doing predictions on the test set
length=len(test_df)
test_batch_size=sorted([int(length/n) for n in range(1,length+1) if length % n ==0 and length/n<=80],reverse=True)[0]  
test_steps=int(length/test_batch_size)
print ( 'test batch size: ' ,test_batch_size, '  test steps: ', test_steps)
trgen=ImageDataGenerator(horizontal_flip=True)
tvgen=ImageDataGenerator()
msg='                                                              for the train generator'
print(msg, '\r', end='') 
train_gen=trgen.flow_from_dataframe( train_df, x_col='filepaths', y_col='labels', target_size=img_size, class_mode='categorical',
                                    color_mode='rgb', shuffle=True, batch_size=batch_size)
msg='                                                              for the test generator'
print(msg, '\r', end='') 
test_gen=tvgen.flow_from_dataframe( test_df, x_col='filepaths', y_col='labels', target_size=img_size, class_mode='categorical',
                                    color_mode='rgb', shuffle=False, batch_size=test_batch_size)
msg='                                                             for the validation generator'
print(msg, '\r', end='')
valid_gen=tvgen.flow_from_dataframe( valid_df, x_col='filepaths', y_col='labels', target_size=img_size, class_mode='categorical',
                                    color_mode='rgb', shuffle=True, batch_size=batch_size)
classes=list(train_gen.class_indices.keys())
class_count=len(classes)
train_steps=int(np.ceil(len(train_gen.labels)/batch_size))
labels=test_gen.labels

Now create your model and train on the train_gen

history=model.fit(x=train_gen,  epochs=epochs, verbose=0, validation_data=valid_gen,
               validation_steps=None,  shuffle=False,  initial_epoch=0)`

after the model is trained you can evaluate its performance on the test set.

loss, acc=model.evaluate(test_gen, steps=test_steps)

Upvotes: 1

Related Questions