xref: /aosp_15_r20/external/executorch/examples/models/mobilenet_v2/model.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import logging
8
9import torch
10
11from torchvision.models import mobilenet_v2  # @manual
12from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
13
14from ..model_base import EagerModelBase
15
16
17class MV2Model(EagerModelBase):
18    def __init__(self):
19        pass
20
21    def get_eager_model(self) -> torch.nn.Module:
22        logging.info("Loading mobilenet_v2 model")
23        mv2 = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
24        logging.info("Loaded mobilenet_v2 model")
25        return mv2
26
27    def get_example_inputs(self):
28        tensor_size = (1, 3, 224, 224)
29        return (torch.randn(tensor_size),)
30
31
32class MV2UntrainedModel(EagerModelBase):
33    def __init__(self):
34        pass
35
36    def get_eager_model(self) -> torch.nn.Module:
37        mv2 = mobilenet_v2()
38        return mv2
39
40    def get_example_inputs(self):
41        tensor_size = (1, 3, 224, 224)
42        return (torch.randn(tensor_size),)
43