Spaces:
Sleeping
Sleeping
File size: 1,312 Bytes
8fe2e46 |
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 |
import keras
import tensorflow as tf
from make_dataset import train_dataset, valid_dataset
from src.components.model import get_cnn_model, TransformerEncoderBlock, TransformerDecoderBlock, ImageCaptioningModel, image_augmentation, LRSchedule
EMBED_DIM = 512
FF_DIM = 512
EPOCHS = 30
cnn_model = get_cnn_model()
encoder = TransformerEncoderBlock(
embed_dim=EMBED_DIM, dense_dim=FF_DIM, num_heads=1)
decoder = TransformerDecoderBlock(
embed_dim=EMBED_DIM, ff_dim=FF_DIM, num_heads=2)
caption_model = ImageCaptioningModel(
cnn_model=cnn_model,
encoder=encoder,
decoder=decoder,
image_aug=image_augmentation,
)
early_stopping = keras.callbacks.EarlyStopping(
patience=3, restore_best_weights=True)
num_train_steps = len(train_dataset) * EPOCHS
num_warmup_steps = num_train_steps // 15
lr_schedule = LRSchedule(post_warmup_learning_rate=1e-4,
warmup_steps=num_warmup_steps)
caption_model.compile(optimizer=keras.optimizers.Adam(lr_schedule), loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
caption_model.fit(
train_dataset,
epochs=EPOCHS,
validation_data=valid_dataset,
callbacks=[early_stopping],
)
caption_model.save("./artifacts/caption_model1.keras")
|