I am trying to train a model that detects whether someone is using sunglasses using tfx and a subset of the celebsA dataset (~ 26k images). I have written the images and labels to a TFrecord that is 232.9MB.
When I then go through the different components, I always run out of memory when running the Transform component. Is this normal? By the way, I am running this on a TPU with 32GB of RAM as I am using Google Colab Pro.
If so, what would be the best way to overcome the problem? Just create many smaller records and pass them through the components one by one?
Here is the code I have been using:
Code for writing to TFRecord:
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential
import tensorflow_datasets as tfds
import pandas as pd
import numpy as np
from PIL import Image
import shutil
import random
import os
import io
import matplotlib.pyplot as plt
%matplotlib inline
from google.colab import drive
drive.mount('/content/gdrive')
%cd gdrive/MyDrive/Machine_Learning_stuff/celebs/
RAW_SUNGLASSES_DIR='./sunglasses_classifier/sunglasses_imgs/raw/'
SUNGLASSES_TFRECORD_DIR= './sunglasses_classifier/data/rec_sunglasses/sunglasses_full.tfrecords'
def _bytes_feature(value):
"""Returns a bytes_list from a string / byte."""
if isinstance(value, type(tf.constant(0))):
value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def _float_feature(value):
"""Returns a float_list from a float / double."""
return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
def _int64_feature(value):
"""Returns an int64_list from a bool / enum / int / uint."""
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def image_resize_to_byte_array(image:Image):
imgByteArr = io.BytesIO()
image=image.resize((256,256))
image.save(imgByteArr, format="jpeg")
imgByteArr = imgByteArr.getvalue()
return imgByteArr
#Remove any corrupted files and non jpeg files
!find ${RAW_SUNGLASSES_DIR} -size 0 -exec rm {} +
!find ${RAW_SUNGLASSES_DIR} -type f ! -name "*.jpg" -exec rm {} +
image_labels={}
for filename in os.listdir(RAW_SUNGLASSES_DIR + '1-sunglasses'):
if '.jpg' in filename:
file_path=os.path.join(RAW_SUNGLASSES_DIR,'1-sunglasses',filename)
#print(file_path)
image_labels[file_path]=1
for filename in os.listdir(RAW_SUNGLASSES_DIR + 'no_sunglasses'):
if '.jpg' in filename:
file_path=os.path.join(RAW_SUNGLASSES_DIR,'no_sunglasses',filename)
#print(file_path)
image_labels[file_path]=0
# Create a dictionary with features that are relevant.
def image_example(image_string, label):
image_shape = tf.io.decode_jpeg(image_string).shape
feature = {
'label': _int64_feature(label),
'image_raw': _bytes_feature(image_string),
}
return tf.train.Example(features=tf.train.Features(feature=feature))
with tf.io.TFRecordWriter(SUNGLASSES_TFRECORD_DIR) as writer:
for filepath, label in image_labels.items():
image_bytes=image_resize_to_byte_array(Image.open(filepath,mode='r'))
#image_string = open(filepath, 'rb').read()
tf_example = image_example(image_bytes, label)
writer.write(tf_example.SerializeToString())
Code for TFX pipeline:
import tensorflow as tf
from tensorflow import keras
#import tensorflow_datasets as tfds
import os
import pprint
#import tfx
from tfx.components import ImportExampleGen
from tfx.components import ExampleValidator
from tfx.components import SchemaGen
from tfx.components import StatisticsGen
from tfx.components import Transform
from tfx.components import Tuner
from tfx.components import Trainer
from tfx.proto import example_gen_pb2
from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext
# Location of the pipeline metadata store
_pipeline_root = 'pipeline/'
# Directory of the raw data files
_data_root = './data/rec_sunglasses/'
from google.colab import drive
drive.mount('/content/gdrive')
%cd gdrive/MyDrive/Machine_Learning_stuff/celebs/sunglasses_classifier/
context = InteractiveContext(pipeline_root=_pipeline_root)
#ExampleGen
example_gen = ImportExampleGen(input_base=_data_root)
context.run(example_gen)
#StatisticsGen
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
context.run(statistics_gen)
#SchemaGen
schema_gen = SchemaGen(
statistics=statistics_gen.outputs['statistics'],infer_feature_shape=True)
context.run(schema_gen)
#ExampleValidator
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
context.run(example_validator)
#Transform
_transform_module_file = 'sunglasses_transform.py'
%%writefile {_transform_module_file}
import tensorflow as tf
import tensorflow_transform as tft
# Keys
_LABEL_KEY = "label"
_IMAGE_KEY = "image_raw"
def _transformed_name(key):
return key + '_xf'
def _image_parser(image_str):
'''converts the images to a float tensor'''
image = tf.image.decode_image(image_str,channels=3)
image = tf.reshape(image, (256, 256, 3))
image = tf.cast(image, tf.float32)
return image
def _label_parser(label_id):
'''converts the labels to a float tensor'''
label = tf.cast(label_id, tf.float32)
return label
def preprocessing_fn(inputs):
"""tf.transform's callback function for preprocessing inputs.
Args:
inputs: map from feature keys to raw not-yet-transformed features.
Returns:
Map from string feature key to transformed feature operations.
"""
# Convert the raw image and labels to a float array
#print(inputs)
outputs = {
_transformed_name(_IMAGE_KEY):
tf.map_fn(
_image_parser,
tf.squeeze(inputs[_IMAGE_KEY], axis=1),
dtype=tf.float32),
_transformed_name(_LABEL_KEY):
tf.map_fn(
_label_parser,
inputs[_LABEL_KEY],
dtype=tf.float32)
}
# scale the pixels from 0 to 1
outputs[_transformed_name(_IMAGE_KEY)] = tft.scale_to_0_1(outputs[_transformed_name(_IMAGE_KEY)])
return outputs
When I then run the code below I always get a message after about 23 mins stating that my runtime was restarted because I ran out of RAM.
# Ignore TF warning messages
tf.get_logger().setLevel('ERROR')
# Setup the Transform component
transform = Transform(
examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=os.path.abspath(_transform_module_file))
# Run the component
context.run(transform)