Spaces:
Sleeping
Sleeping
File size: 2,675 Bytes
3e93e14 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import os
import tensorflow as tf
from pathlib import Path
from cnnClassifier.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, compile=False
)
self.model.compile(
optimizer=tf.keras.optimizers.SGD(learning_rate=self.config.params_learning_rate),
loss=tf.keras.losses.CategoricalCrossentropy(),
metrics=["accuracy"]
)
def train_valid_generator(self):
datagenerator_kwargs = dict(rescale=1.0 / 255, validation_split=0.20)
dataflow_kwargs = dict(
target_size=self.config.params_image_size[:-1],
batch_size=self.config.params_batch_size,
interpolation="bilinear"
)
valid_datagenerator = tf.keras.preprocessing.image.ImageDataGenerator(
**datagenerator_kwargs
)
self.valid_generator = valid_datagenerator.flow_from_directory(
directory=self.config.training_data,
subset="validation",
shuffle=False,
**dataflow_kwargs
)
if self.config.params_is_augmentation:
train_datagenerator = tf.keras.preprocessing.image.ImageDataGenerator(
rotation_range=40,
horizontal_flip=True,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
**datagenerator_kwargs
)
else:
train_datagenerator = valid_datagenerator
self.train_generator = train_datagenerator.flow_from_directory(
directory=self.config.training_data,
subset="training",
shuffle=True,
**dataflow_kwargs
)
@staticmethod
def save_model(path: Path, model: tf.keras.Model):
model.save(path)
def train(self):
self.steps_per_epoch = self.train_generator.samples // self.train_generator.batch_size
self.validation_steps = self.valid_generator.samples // self.valid_generator.batch_size
self.model.fit(
self.train_generator,
epochs=self.config.params_epochs,
steps_per_epoch=self.steps_per_epoch,
validation_steps=self.validation_steps,
validation_data=self.valid_generator
)
self.save_model(path=self.config.trained_model_path, model=self.model)
|