def __init__(self, dataset_dir, base_model, head_model, optimizer):
    'Creates a wrapper for a set of models and a data set.'
    self.dataset_dir = dataset_dir
    datagen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=(1.0 / 255), validation_split=VALIDATION_SPLIT)
    self.train_img_generator = datagen.flow_from_directory(self.dataset_dir, target_size=(IMAGE_SIZE, IMAGE_SIZE), batch_size=BATCH_SIZE, subset='training')
    self.val_img_generator = datagen.flow_from_directory(self.dataset_dir, target_size=(IMAGE_SIZE, IMAGE_SIZE), batch_size=BATCH_SIZE, subset='validation')
    converter = tflite_transfer_converter.TFLiteTransferConverter(NUM_CLASSES, base_model, head_model, optimizer, BATCH_SIZE)
    models = converter._convert()
    self.initialize_model = models['initialize']
    self.bottleneck_model = models['bottleneck']
    self.train_head_model = models['train_head']
    self.inference_model = models['inference']
    self.optimizer_model = models['optimizer']
    self.variables = self._generate_initial_variables()
    optim_state_shapes = self._optimizer_state_shapes()
    self.optim_state = [np.zeros(shape, dtype=np.float32) for shape in optim_state_shapes]
Creates a wrapper for a set of models and a data set.
def __init__(self, dataset_dir, base_model, head_model, optimizer):
    
    self.dataset_dir = dataset_dir
    datagen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=(1.0 / 255), validation_split=VALIDATION_SPLIT)
    self.train_img_generator = datagen.flow_from_directory(self.dataset_dir, target_size=(IMAGE_SIZE, IMAGE_SIZE), batch_size=BATCH_SIZE, subset='training')
    self.val_img_generator = datagen.flow_from_directory(self.dataset_dir, target_size=(IMAGE_SIZE, IMAGE_SIZE), batch_size=BATCH_SIZE, subset='validation')
    converter = tflite_transfer_converter.TFLiteTransferConverter(NUM_CLASSES, base_model, head_model, optimizer, BATCH_SIZE)
    models = converter._convert()
    self.initialize_model = models['initialize']
    self.bottleneck_model = models['bottleneck']
    self.train_head_model = models['train_head']
    self.inference_model = models['inference']
    self.optimizer_model = models['optimizer']
    self.variables = self._generate_initial_variables()
    optim_state_shapes = self._optimizer_state_shapes()
    self.optim_state = [np.zeros(shape, dtype=np.float32) for shape in optim_state_shapes]