Blog de Amazon Web Services (AWS)
Clasificación de imágenes médicas con SageMaker y Tensorflow2
Por María Gaska, AI/ML Specialist SA de AWS
Uno de los grandes desafíos para la transformación digital en instituciones médicas es el desarrollo de herramientas de soporte al diagnóstico a través de imágenes. Con este tipo de aplicaciones, los médicos proveen una imagen (resonancias magnéticas, rayos X, etc) a través de una aplicación y reciben como resultado una probabilidad asociada a un determinado diagnóstico. El desafío a la hora de construirlas es que se necesita no sólo un modelo de machine learning con una alta precisión si no también de una implementación robusta que permita realizar varias consultas de manera simultánea y en tiempo real.
En esta publicación veremos cómo construir un modelo de clasificación de imágenes médicas, implementarlo y consumirlo en tiempo real utilizando el SDK de SageMaker. El código completo que acompaña a este artículo se encuentra aquí.
Clasificación de distintos tipos de Alzheimer.
El data set que vamos a utilizar para esta demostración proviene de esta competencia de kaggle. Las imágenes consisten en un conjunto de resonancias magnéticas clasificadas por la gravedad de los síntomas de Alzheimer que presentan los pacientes: Ninguno, Suave, Muy Suave o Moderado.
Vamos a utilizar este conjunto de imágenes para armar un modelo de clasificación que posteriormente podamos implementar como un microservicio.
Para este proyecto vamos a utilizar containers manejados de SageMaker para Tensorflow tanto para la etapa de training como la de serving que se desplegarán de forma completamente transparente para nosotros a través del SDK. Estos containers implementan un patrón de diseño llamado «script mode» donde SageMaker provee un menú de containers que pueden utilizarse y el usuario escribe su propio script de entrenamiento personalizado para correr sobre uno de esos entornos.
Paso 1: Obtener los datos de S3 e importar las librerías.
Primero descargamos y descomprimimos los datos
%%bash wget s3://tensorflow-images-classification/images.zip unzip images.zip
Paso 2: Entrenar en «modo local»
Una de las ventajas de SageMaker es su capacidad de aprovisionar y desaprovisionar infraestructura de cualquier tamaño, tanto por unidades como en un clúster. Si bien esto es una gran ventaja a la hora de correr procesos de entrenamientos sobre grandes volúmenes de datos, puede no ser lo más cómodo a la hora de desarrollar y probar un script. El modo local de SageMaker permite levantar el container localmente y entrenar el algoritmo con datos que se encuentran en el storage local sin necesidad de esperar a que se aprovisione nueva infraestructura. Noten que en este paso es importante utilizar un sólo epoch ya que únicamente estamos probando que el script funciona.
Para utilizar el estimador de Tensorflow necesitamos crear un script que entrene un modelo a partir de los datos de «train» y «validation» que fueron cargados en S3 y luego lo guarde en una ubicación determinada. En este caso, el archivo classification.py implementa una VGG16 donde sólo las 3 últimas capas son entrenables.
Noten que los directorios de los cuales se sacan los datos de entrenamiento son ubicaciones locales.
train_dir = os.path.join(os.getcwd(), 'images/train') validation_dir = os.path.join(os.getcwd(), 'images/validation')
A continuación, escribimos el script que entrena el modelo.
%%writefile
import argparse
import numpy as np
import os
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Model
from tensorflow.keras.layers import *
from tensorflow.keras import optimizers
#para prediccion
from tensorflow.keras import backend as K
from tensorflow.keras.models import *
from tensorflow.keras.layers import *
from tensorflow.keras.optimizers import *
from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras.callbacks import ModelCheckpoint, LearningRateScheduler, EarlyStopping, ReduceLROnPlateau, TensorBoard
train_data_gen_args = dict(rescale=1./255,
shear_range=0.01,
rotation_range = 20,
zoom_range=0.2,
height_shift_range = 0.2,
width_shift_range = 0.2,
brightness_range=[0.1, 1.9],
horizontal_flip=True)
data_gen_args = dict(target_size=(224, 224),
batch_size=16,
shuffle=True,
#color_mode='grayscale',
class_mode='categorical')
num_classes = 2
def parse_args():
parser = argparse.ArgumentParser()
# hyperparameters sent by the client are passed as command-line arguments to the script
parser.add_argument('--epochs', type=int, default=1)
parser.add_argument('--batch_size', type=int, default=16)
# data directories
parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAIN'))
parser.add_argument('--validation', type=str, default=os.environ.get('SM_CHANNEL_VALIDATION'))
# model directory: we will use the default set by SageMaker, /opt/ml/model
parser.add_argument('--model_dir', type=str, default=os.environ.get('SM_MODEL_DIR'))
return parser.parse_known_args()
def get_train_val_data(train_dir,validation_dir):
train_datagen = ImageDataGenerator(**train_data_gen_args)
train_generator = train_datagen.flow_from_directory(train_dir, **data_gen_args)
val_generator = train_datagen.flow_from_directory(validation_dir, **data_gen_args)
print('train shape:', train_generator[0][0].shape,'val shape:', val_generator[0][0].shape)
return train_generator, val_generator
def get_model():
base_model = VGG16(input_shape=(224,224,3), weights='imagenet', include_top=False)
#x = GlobalAveragePooling2D()(base_model.output)
x = Flatten()(base_model.output)
x = Dense(1024, activation='relu')(x)
x = Dropout(0.3)(x)
x = Dense(num_classes, activation='softmax')(x)
model = Model(inputs=base_model.input, outputs=x)
for layer in base_model.layers:
layer.trainable = False
return model
if __name__ == "__main__":
args, _ = parse_args()
train_generator, val_generator = get_train_val_data(args.train,args.validation)
model = get_model()
opt = Adam(learning_rate=0.0025)
model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
callbacks_list = []
STEP_SIZE_TRAIN=train_generator.n//train_generator.batch_size
STEP_SIZE_VALID=val_generator.n//val_generator.batch_size
model.fit(train_generator, steps_per_epoch=STEP_SIZE_TRAIN, epochs=args.epochs, validation_data=val_generator, validation_steps=STEP_SIZE_VALID, callbacks=callbacks_list)
# create a TensorFlow SavedModel for deployment to a SageMaker endpoint with TensorFlow Serving
model.save(args.model_dir + '/1')
#tf.keras.models.save_model(model, args.model_dir)
Con el SDK de SageMaker entrenamos el modelo. Noten que el parámetro train_instance_type se debe configurar como «local».
import sagemaker
from sagemaker.tensorflow import TensorFlow
model_dir = '/opt/ml/model'
train_instance_type = 'local'
hyperparameters = {'epochs': 1}
local_estimator = TensorFlow(
entry_point='classification.py',
model_dir=model_dir,
train_instance_type=train_instance_type,
train_instance_count=1,
hyperparameters=hyperparameters,
role=sagemaker.get_execution_role(),
base_job_name='tf-keras-clasif',
framework_version='2.0.0',
py_version='py3',
script_mode=True)
Paso 3: Subir los datos a S3
Una vez que estamos listos para el proceso de entrenamiento definitivo, deberíamos subir los datos a S3 y guardar la ubicación para informársela al proceso de entrenamiento.
El SDK de SageMaker ofrece un bucket por default cuyo nombre está asociado a nuestra cuenta y se puede acceder a través del método upload_data.
s3_prefix = 'medicalimgs'
traindata_s3_prefix = '{}/data/train'.format(s3_prefix)
validation_s3_prefix = '{}/data/validation'.format(s3_prefix)
train_s3 = sagemaker.Session().upload_data(path='./images/train/', key_prefix=traindata_s3_prefix)
validation_s3 = sagemaker.Session().upload_data(path='./images/validation/', key_prefix=validation_s3_prefix)
inputs = {'train':train_s3,'validation':validation_s3}
print(inputs)
Paso 4: Crear el Estimator
Ahora que tenemos definida la arquitectura de entrenamiento, podemos usar este script en la clase Tensorflow del SDK de SageMaker para entrenar el modelo durante 100 epochs sobre una instancia de tipo ml.p3.2xlarge.
train_instance_type = 'ml.p3.2xlarge'
hyperparameters = {'epochs': 100}
estimator = TensorFlow(
entry_point='classification.py',
model_dir=model_dir,
train_instance_type=train_instance_type,
train_instance_count=1,
hyperparameters=hyperparameters,
role=sagemaker.get_execution_role(),
base_job_name='tf-keras-clasif',
framework_version='2.0.0',
py_version='py3',
script_mode=True)
estimator.fit(inputs)
Paso 5: Crear el endpoint de predicción
Una vez terminado el entrenamiento definitivo, podemos utilizar ese modelo para crear un endpoint que entregue predicciones en tiempo real. Vamos a desplegar el modelo en una sola instancia de tipo ml.m5.xlarge.
predictor = estimator.deploy(initial_instance_count=1, instance_type='ml.m5.xlarge')
Paso 6: Evaluación del modelo
Por último, vamos evaluar el modelo sobre un conjunto de datos de test. Recuerden utilizar los mismos argumentos para los generadores que usaron en classification.py. A la hora de inferir, no necesitamos los parámetros relacionados a data augmentation, así que podemos quedarnos sólo con el “rescale”.
test_data_gen_args = dict(rescale=1./255)
test_datagen = ImageDataGenerator(**test_data_gen_args)
test_generator = test_datagen.flow_from_directory('images/test/', **data_gen_args)
Luego de configurar los generadores para transformar las imágenes de test, podemos hacer la inferencia y calcular la matriz de confusión. Noten que el payload que recibe el endpoint es un json con una clave instances donde pasamos el numpy array sobre el cual queremos predecir, serializado utilizando el método tolist().
number_of_examples = 1279
number_of_generator_calls = math.ceil(number_of_examples / (1.0 * data_gen_args['batch_size']))
test_labels = []
predictions = []
for i in range(0,int(number_of_generator_calls)):
instances = test_generator[i][0]
print(instances.shape)
for instance in instances:
array = instance.reshape((1,) + instance.shape)
payload = {
'instances': array.tolist()
}
resp = predictor.predict(payload)['predictions']
predictions.append(np.array(resp))
test_labels.extend(np.array(test_generator[i][1]))
Con estas predicciones podemos construir la matriz de confusion.
predictions = np.array(predictions).reshape(1279,4)
predictions = np.argmax(predictions,axis=1)
labels = np.argmax(np.array(test_labels),axis=1)
classes = list(test_generator.class_indices.keys())
df_cm = confusion_matrix(labels,predictions,labels=np.unique(labels))
heatmap = sns.heatmap(df_cm, annot=True, fmt="d")
heatmap.yaxis.set_ticklabels(classes, rotation=0, ha='right')
heatmap.xaxis.set_ticklabels(classes, rotation=45, ha='right')
plt.ylabel('Valor Verdadero')
plt.xlabel('Valor Predicho');
Esta es la matriz de confusión resultante. Los resultados se pueden refinar con un proceso de optimización de hiperparámetros o probando otras opciones de arquitectura como por ejemplo Resnet.
Conclusión
Utilizando estas herramientas podemos construir un microservicio liviano, enteramente «serverless» y por lo tanto muy efectivo desde el punto de vista de los costos y del mantenimiento de infraestructura.
Limpieza
Para evitar incurrir en cargos pueden borrar las imágenes del bucket S3 y la instancia de Jupyter notebook donde se haya ejecutado el código.
Sobre el autor
María Gaska es arquitecta de soluciones en AWS desde hace casi dos años. En su rol, ayuda a los clientes tanto a determinar la mejor arquitectura para sus distintas aplicaciones como a encontrar los mejores algoritmos para resolver problemas de Machine Learning e IA. Antes de AWS, trabajó como desarrolladora de modelos de deep learning en un startup enfocado en NLP y chatbots y también como profesora full time en una coding school a cargo de un curso de data science.
Sobre los revisores
Andres Palacios es arquitecto de soluciones especialista en Analytics en AWS. En su rol apoya a los clientes a encontrar la mejor solucion y arquitectura para sus necesidades al igual que aprovechar los servicios de AI/ML para generar innovación y mejorar la productividad. Antes de AWS, trabajó para consultoras Big Four en las áreas de Data y Analytics, tanto como en consultoría estratégica, arquitectura e implementación de soluciones de procesamiento y consumo distribuidas.
Sergio Beltran es arquitecto de soluciones especialista en AI/ML en AWS. En su rol apoya a los clientes a encontrar la mejor solución y arquitectura para sus necesidades al igual que aprovechar los servicios de AI/ML para generar innovación y mejorar la productividad. Antes de AWS, trabajó data scientist y gerente de business development en la industria Telco.