cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Remove unnecessary architectures from image classifier ModelSpec

Browse Source 
PiperOrigin-RevId: 481974529
This commit is contained in:
MediaPipe Team 2022-10-18 11:27:44 -07:00 committed by Copybara-Service
parent bc47589c9b
commit 51879ae81a
3 changed files with 0 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
mediapipe/model_maker/python/vision/image_classifier/image_classifier_test.py
View File

@ -60,11 +60,6 @@ class ImageClassifierTest(tf.test.TestCase, parameterized.TestCase):
model_spec=image_classifier.SupportedModels.MOBILENET_V2,
hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)),
dict(
testcase_name='resnet_50',
model_spec=image_classifier.SupportedModels.RESNET_50,
hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)),
dict(
testcase_name='efficientnet_lite0',
model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE0,
@ -75,21 +70,6 @@ class ImageClassifierTest(tf.test.TestCase, parameterized.TestCase):
model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE1,
hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)),
dict(
testcase_name='efficientnet_lite2',
model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE2,
hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)),
dict(
testcase_name='efficientnet_lite3',
model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE3,
hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)),
dict(
testcase_name='efficientnet_lite4',
model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE4,
hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)),
)
def test_create_and_train_model(self,
model_spec: image_classifier.SupportedModels,

27
mediapipe/model_maker/python/vision/image_classifier/model_spec.py
View File

@ -48,11 +48,6 @@ mobilenet_v2_spec = functools.partial(
uri='https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4',
name='mobilenet_v2')
resnet_50_spec = functools.partial(
ModelSpec,
uri='https://tfhub.dev/google/imagenet/resnet_v2_50/feature_vector/4',
name='resnet_50')
efficientnet_lite0_spec = functools.partial(
ModelSpec,
uri='https://tfhub.dev/tensorflow/efficientnet/lite0/feature-vector/2',
@ -64,36 +59,14 @@ efficientnet_lite1_spec = functools.partial(
input_image_shape=[240, 240],
name='efficientnet_lite1')
efficientnet_lite2_spec = functools.partial(
ModelSpec,
uri='https://tfhub.dev/tensorflow/efficientnet/lite2/feature-vector/2',
input_image_shape=[260, 260],
name='efficientnet_lite2')
efficientnet_lite3_spec = functools.partial(
ModelSpec,
uri='https://tfhub.dev/tensorflow/efficientnet/lite3/feature-vector/2',
input_image_shape=[280, 280],
name='efficientnet_lite3')
efficientnet_lite4_spec = functools.partial(
ModelSpec,
uri='https://tfhub.dev/tensorflow/efficientnet/lite4/feature-vector/2',
input_image_shape=[300, 300],
name='efficientnet_lite4')
# TODO: Document the exposed models.
@enum.unique
class SupportedModels(enum.Enum):
"""Image classifier model supported by model maker."""
MOBILENET_V2 = mobilenet_v2_spec
RESNET_50 = resnet_50_spec
EFFICIENTNET_LITE0 = efficientnet_lite0_spec
EFFICIENTNET_LITE1 = efficientnet_lite1_spec
EFFICIENTNET_LITE2 = efficientnet_lite2_spec
EFFICIENTNET_LITE3 = efficientnet_lite3_spec
EFFICIENTNET_LITE4 = efficientnet_lite4_spec
@classmethod
def get(cls, spec: 'SupportedModels') -> 'ModelSpec':

24
mediapipe/model_maker/python/vision/image_classifier/model_spec_test.py
View File

@ -30,12 +30,6 @@ class ModelSpecTest(tf.test.TestCase, parameterized.TestCase):
expected_uri='https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4',
expected_name='mobilenet_v2',
expected_input_image_shape=[224, 224]),
dict(
testcase_name='resnet_50_spec_test',
model_spec=ms.resnet_50_spec,
expected_uri='https://tfhub.dev/google/imagenet/resnet_v2_50/feature_vector/4',
expected_name='resnet_50',
expected_input_image_shape=[224, 224]),
dict(
testcase_name='efficientnet_lite0_spec_test',
model_spec=ms.efficientnet_lite0_spec,
@ -48,24 +42,6 @@ class ModelSpecTest(tf.test.TestCase, parameterized.TestCase):
expected_uri='https://tfhub.dev/tensorflow/efficientnet/lite1/feature-vector/2',
expected_name='efficientnet_lite1',
expected_input_image_shape=[240, 240]),
dict(
testcase_name='efficientnet_lite2_spec_test',
model_spec=ms.efficientnet_lite2_spec,
expected_uri='https://tfhub.dev/tensorflow/efficientnet/lite2/feature-vector/2',
expected_name='efficientnet_lite2',
expected_input_image_shape=[260, 260]),
dict(
testcase_name='efficientnet_lite3_spec_test',
model_spec=ms.efficientnet_lite3_spec,
expected_uri='https://tfhub.dev/tensorflow/efficientnet/lite3/feature-vector/2',
expected_name='efficientnet_lite3',
expected_input_image_shape=[280, 280]),
dict(
testcase_name='efficientnet_lite4_spec_test',
model_spec=ms.efficientnet_lite4_spec,
expected_uri='https://tfhub.dev/tensorflow/efficientnet/lite4/feature-vector/2',
expected_name='efficientnet_lite4',
expected_input_image_shape=[300, 300]),
)
def test_predefiend_spec(self, model_spec: Callable[..., ms.ModelSpec],
expected_uri: str, expected_name: str,