|
|
import tensorflow as tf |
|
|
from pathlib import Path |
|
|
from kidney_classification.entity.config_entity import TrainingConfig |
|
|
|
|
|
|
|
|
class Training: |
|
|
def __init__(self, config: TrainingConfig): |
|
|
self.config = config |
|
|
|
|
|
def get_base_model(self): |
|
|
self.model = tf.keras.models.load_model(self.config.updated_base_model_path) |
|
|
|
|
|
def train_valid_generator(self): |
|
|
img_height, img_width = self.config.params_image_size[:-1] |
|
|
|
|
|
train = tf.keras.utils.image_dataset_from_directory( |
|
|
self.config.training_data, |
|
|
image_size=(img_height, img_width), |
|
|
validation_split=0.1, |
|
|
subset="training", |
|
|
seed=123, |
|
|
) |
|
|
|
|
|
val = tf.keras.utils.image_dataset_from_directory( |
|
|
self.config.training_data, |
|
|
image_size=(img_height, img_width), |
|
|
validation_split=0.2, |
|
|
subset="validation", |
|
|
seed=123, |
|
|
) |
|
|
train = train.map(lambda x, y: (x / 255, y)) |
|
|
val = val.map(lambda x, y: (x / 255, y)) |
|
|
AUTOTUNE = tf.data.AUTOTUNE |
|
|
|
|
|
self.train_dataset = train.cache().prefetch(buffer_size=AUTOTUNE) |
|
|
self.val_dataset = val.cache().prefetch(buffer_size=AUTOTUNE) |
|
|
|
|
|
@staticmethod |
|
|
def save_model(path: Path, model: tf.keras.Model): |
|
|
model.save(path) |
|
|
|
|
|
def define_and_train_model(self): |
|
|
self.model.fit( |
|
|
self.train_dataset, |
|
|
validation_data=self.val_dataset, |
|
|
epochs=self.config.params_epochs, |
|
|
) |
|
|
|
|
|
self.save_model(path=self.config.trained_model_path, model=self.model) |
|
|
|
