Spaces:
Sleeping
Sleeping
File size: 1,586 Bytes
47488ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
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)
|