Spaces:
Sleeping
Sleeping
File size: 4,389 Bytes
0054a0f 36b8bfd cb053fc 36b8bfd cb053fc 36b8bfd cb053fc 36b8bfd cb053fc 36b8bfd cb053fc 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd cb053fc 0054a0f cb053fc 0054a0f 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd 0054a0f 36b8bfd | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | import os
import urllib.request as request
from zipfile import ZipFile
import tensorflow as tf
from pathlib import Path
from cnnClassifier.entity.config_entity import PrepareBaseModelConfig
#model for different architectures
MODEL_CONFIGS = {
"VGG16": {
"fn": tf.keras.applications.VGG16,
"pooling": "flatten", # VGG16 spatial output suits Flatten
"dropout": 0.5,
"unfreeze": 0, # 138M params — too large to fine-tune safely
"lr_override": None, # respects params.yaml LEARNING_RATE
"optimizer": "sgd" # VGG16 benefits from momentum and lower LR
},
"EfficientNetV2B3": {
"fn": tf.keras.applications.EfficientNetV2B3,
"pooling": "gap", # GlobalAveragePooling2D — built for this
"dropout": 0.4,
"unfreeze": 20, # unfreeze top 20 layers for CT domain adaptation
"lr_override": 0.0001, # overrides params.yaml — EfficientNet needs lower LR
"optimizer": "adam"
},
"DenseNet121": {
"fn": tf.keras.applications.DenseNet121,
"pooling": "gap",
"dropout": 0.4,
"unfreeze": 20,
"lr_override": 0.0001,
"optimizer": "adam"
},
"ResNet50V2": {
"fn": tf.keras.applications.ResNet50V2,
"pooling": "gap",
"dropout": 0.4,
"unfreeze": 20,
"lr_override": 0.0001,
"optimizer": "adam"
},
}
DEFAULT_CONFIG = {
"fn": tf.keras.applications.VGG16,
"pooling": "gap",
"dropout": 0.5,
"unfreeze": 0,
"lr_override": None,
"optimizer": "sgd"
}
class PrepareBaseModel:
def __init__(self, config: PrepareBaseModelConfig):
self.config = config
self.model_name = getattr(config, "params_model_name", "VGG16")
self.model_cfg = MODEL_CONFIGS.get(self.model_name, DEFAULT_CONFIG)
def get_base_model(self):
self.model = self.model_cfg["fn"](
input_shape=self.config.params_image_size,
weights=self.config.params_weights,
include_top=self.config.params_include_top
)
self.save_model(path=self.config.base_model_path, model=self.model)
@staticmethod
def _prepare_full_model(model, classes, learning_rate, model_cfg):
# freeze all layers
for layer in model.layers:
layer.trainable = False
#unfreeze top N layers for domain adaptation
unfreeze = model_cfg["unfreeze"]
if unfreeze > 0:
for layer in model.layers[-unfreeze:]:
layer.trainable = True
#pooling Flatten for VGG16, GAP for all modern architectures
if model_cfg["pooling"] == "flatten":
x = tf.keras.layers.Flatten()(model.output)
else:
x = tf.keras.layers.GlobalAveragePooling2D()(model.output)
#classification head
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dropout(model_cfg["dropout"])(x)
prediction = tf.keras.layers.Dense(
units=classes,
activation="softmax"
)(x)
full_model = tf.keras.models.Model(
inputs=model.input,
outputs=prediction
)
#lr_override takes priority over params.yaml value
effective_lr = model_cfg["lr_override"] or learning_rate
if model_cfg["optimizer"] == "sgd":
optimizer = tf.keras.optimizers.SGD(learning_rate=effective_lr)
else:
optimizer = tf.keras.optimizers.Adam(learning_rate=effective_lr)
full_model.compile(
optimizer=optimizer,
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=["accuracy"]
)
full_model.summary()
return full_model
def update_base_model(self):
self.full_model = self._prepare_full_model(
model=self.model,
classes=self.config.params_classes,
learning_rate=self.config.params_learning_rate,
model_cfg=self.model_cfg
)
self.save_model(
path=self.config.updated_base_model_path,
model=self.full_model
)
@staticmethod
def save_model(path: Path, model: tf.keras.Model):
model.save(path) |