xref: /aosp_15_r20/external/executorch/examples/cadence/models/resnet50.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
7# Example script for exporting simple models to flatbuffer
8
9import logging
10
11import torch
12
13from executorch.backends.cadence.aot.ops_registrations import *  # noqa
14
15
16from executorch.backends.cadence.aot.export_example import export_model
17from torchvision.models import resnet50, ResNet50_Weights
18
19
20FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
21logging.basicConfig(level=logging.INFO, format=FORMAT)
22
23
24if __name__ == "__main__":
25
26    model = resnet50(weights=ResNet50_Weights.DEFAULT)
27    model.eval()
28    example_inputs = (torch.randn(1, 3, 64, 64),)
29
30    export_model(model, example_inputs)
31