diff --git a/examples/training/classification/imagenet/training_history.json b/examples/training/classification/imagenet/training_history.json index ad28d79a82..621c044b9d 100644 --- a/examples/training/classification/imagenet/training_history.json +++ b/examples/training/classification/imagenet/training_history.json @@ -1,4 +1,25 @@ { + "cspdarknet": { + "v0": { + "accelerators": 8, + "args": { + "batch_size": "64", + "epochs": "300", + "initial_learning_rate": ".0125", + "learning_rate_schedule": "CosineDecayWithWarmup", + "warmup_hold_steps_percentage": ".45", + "warmup_steps_percentage": ".01" + }, + "contributor": "ianstenbit", + "epochs_trained": 299, + "script": { + "name": "basic_training.py", + "version": "dceea23c954e59c5884e98384140e0a8ad5bd320" + }, + "tensorboard_logs": "https://tensorboard.dev/experiment/3eUYhaFMQ3O3fvdg1lDRMw/", + "validation_accuracy": "0.7744" + } + }, "darknet53": { "v0": { "accelerators": 8, diff --git a/keras_cv/models/csp_darknet.py b/keras_cv/models/csp_darknet.py index a966b9147e..d20c1a042b 100644 --- a/keras_cv/models/csp_darknet.py +++ b/keras_cv/models/csp_darknet.py @@ -30,6 +30,7 @@ from keras_cv.models.__internal__.darknet_utils import DarknetConvBlockDepthwise from keras_cv.models.__internal__.darknet_utils import Focus from keras_cv.models.__internal__.darknet_utils import SpatialPyramidPoolingBottleneck +from keras_cv.models.weights import parse_weights def CSPDarkNet( @@ -74,8 +75,9 @@ def CSPDarkNet( should be used over a regular darknet block. Defaults to False classes: optional number of classes to classify images into, only to be specified if `include_top` is True. - weights: one of `None` (random initialization), or a pretrained weight - file path. + weights: one of `None` (random initialization), a pretrained weight file + path, or a reference to pre-trained weights (e.g. 'imagenet/classification') + (see available pre-trained weights in weights.py) input_tensor: optional Keras tensor (i.e. output of `layers.Input()`) to use as image input for the model. input_shape: optional shape tuple, defaults to (None, None, 3). @@ -95,6 +97,8 @@ def CSPDarkNet( Returns: A `keras.Model` instance. """ + weights = parse_weights(weights, include_top, "cspdarknet") + if weights and not tf.io.gfile.exists(weights): raise ValueError( "The `weights` argument should be either `None` or the path to the " diff --git a/keras_cv/models/weights.py b/keras_cv/models/weights.py index 5083a19c8b..987a1929fe 100644 --- a/keras_cv/models/weights.py +++ b/keras_cv/models/weights.py @@ -39,6 +39,10 @@ def parse_weights(weights, include_top, model_type): BASE_PATH = "https://storage.googleapis.com/keras-cv/models" ALIASES = { + "cspdarknet": { + "imagenet": "imagenet/classification-v0", + "imagenet/classification": "imagenet/classification-v0", + }, "darknet53": { "imagenet": "imagenet/classification-v0", "imagenet/classification": "imagenet/classification-v0", @@ -78,6 +82,10 @@ def parse_weights(weights, include_top, model_type): } WEIGHTS_CONFIG = { + "cspdarknet": { + "imagenet/classification-v0": "8bdc3359222f0d26f77aa42c4e97d67a05a1431fe6c448ceeab9a9c5a34ff804", + "imagenet/classification-v0-notop": "9303aabfadffbff8447171fce1e941f96d230d8f3cef30d3f05a9c85097f8f1e", + }, "darknet53": { "imagenet/classification-v0": "7bc5589f7f7f7ee3878e61ab9323a71682bfb617eb57f530ca8757c742f00c77", "imagenet/classification-v0-notop": "8dcce43163e4b4a63e74330ba1902e520211db72d895b0b090b6bfe103e7a8a5",