xref: /aosp_15_r20/external/executorch/examples/cadence/models/resnet50.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1*523fa7a6SAndroid Build Coastguard Worker# Copyright (c) Meta Platforms, Inc. and affiliates.
2*523fa7a6SAndroid Build Coastguard Worker# All rights reserved.
3*523fa7a6SAndroid Build Coastguard Worker#
4*523fa7a6SAndroid Build Coastguard Worker# This source code is licensed under the BSD-style license found in the
5*523fa7a6SAndroid Build Coastguard Worker# LICENSE file in the root directory of this source tree.
6*523fa7a6SAndroid Build Coastguard Worker
7*523fa7a6SAndroid Build Coastguard Worker# Example script for exporting simple models to flatbuffer
8*523fa7a6SAndroid Build Coastguard Worker
9*523fa7a6SAndroid Build Coastguard Workerimport logging
10*523fa7a6SAndroid Build Coastguard Worker
11*523fa7a6SAndroid Build Coastguard Workerimport torch
12*523fa7a6SAndroid Build Coastguard Worker
13*523fa7a6SAndroid Build Coastguard Workerfrom executorch.backends.cadence.aot.ops_registrations import *  # noqa
14*523fa7a6SAndroid Build Coastguard Worker
15*523fa7a6SAndroid Build Coastguard Worker
16*523fa7a6SAndroid Build Coastguard Workerfrom executorch.backends.cadence.aot.export_example import export_model
17*523fa7a6SAndroid Build Coastguard Workerfrom torchvision.models import resnet50, ResNet50_Weights
18*523fa7a6SAndroid Build Coastguard Worker
19*523fa7a6SAndroid Build Coastguard Worker
20*523fa7a6SAndroid Build Coastguard WorkerFORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
21*523fa7a6SAndroid Build Coastguard Workerlogging.basicConfig(level=logging.INFO, format=FORMAT)
22*523fa7a6SAndroid Build Coastguard Worker
23*523fa7a6SAndroid Build Coastguard Worker
24*523fa7a6SAndroid Build Coastguard Workerif __name__ == "__main__":
25*523fa7a6SAndroid Build Coastguard Worker
26*523fa7a6SAndroid Build Coastguard Worker    model = resnet50(weights=ResNet50_Weights.DEFAULT)
27*523fa7a6SAndroid Build Coastguard Worker    model.eval()
28*523fa7a6SAndroid Build Coastguard Worker    example_inputs = (torch.randn(1, 3, 64, 64),)
29*523fa7a6SAndroid Build Coastguard Worker
30*523fa7a6SAndroid Build Coastguard Worker    export_model(model, example_inputs)
31