torch-mlir/e2e_testing/torchscript/vision_models.py

128 lines
3.7 KiB
Python
Raw Normal View History

# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# Also available under a BSD-style license. See LICENSE.
import torch
import torchvision.models as models
from torch_mlir_e2e_test.torchscript.framework import TestUtils
from torch_mlir_e2e_test.torchscript.registry import register_test_case
from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
# ==============================================================================
class ResNet18Module(torch.nn.Module):
def __init__(self):
super().__init__()
# Reset seed to make model deterministic.
torch.manual_seed(0)
self.resnet = models.resnet18()
self.train(False)
@export
@annotate_args([
None,
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
([-1, 3, -1, -1], torch.float32, True),
])
def forward(self, img):
return self.resnet.forward(img)
@register_test_case(module_factory=lambda: ResNet18Module())
def ResNet18Module_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 3, 224, 224))
class ResNet18StaticModule(torch.nn.Module):
def __init__(self):
super().__init__()
# Reset seed to make model deterministic.
torch.manual_seed(0)
self.resnet = models.resnet18()
self.train(False)
@export
@annotate_args([
None,
([1, 3, 224, 224], torch.float32, True),
])
def forward(self, img):
return self.resnet.forward(img)
@register_test_case(module_factory=lambda: ResNet18StaticModule())
def ResNet18StaticModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 3, 224, 224))
class IouOfModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, bbox1, bbox2):
area1 = (bbox1[:, 2] - bbox1[:, 0]) * (bbox1[:, 3] - bbox1[:, 1])
area2 = (bbox2[:, 2] - bbox2[:, 0]) * (bbox2[:, 3] - bbox2[:, 1])
lt = torch.maximum(bbox1[:, :2], bbox2[:, :2])
rb = torch.minimum(bbox1[:, 2:], bbox2[:, 2:])
overlap_coord = (rb - lt).clip(0)
overlap = overlap_coord[:, 0] * overlap_coord[:, 1]
union = area1 + area2 - overlap
return overlap / union
@register_test_case(module_factory=lambda: IouOfModule())
def IouOfModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1024, 4), tu.rand(1024, 4))
class MobilenetV2Module(torch.nn.Module):
def __init__(self):
super().__init__()
# Reset seed to make model deterministic.
torch.manual_seed(0)
self.mobilenetv2 = models.mobilenet_v2()
self.train(False)
@export
@annotate_args([
None,
([-1, 3, -1, -1], torch.float32, True),
])
def forward(self, img):
return self.mobilenetv2.forward(img)
@register_test_case(module_factory=lambda: MobilenetV2Module())
def MobilenetV2Module_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 3, 224, 224))
class MobilenetV3Module(torch.nn.Module):
def __init__(self):
super().__init__()
# Reset seed to make model deterministic.
torch.manual_seed(0)
self.mobilenetv3 = models.mobilenet_v3_small()
self.train(False)
@export
@annotate_args([
None,
([-1, 3, -1, -1], torch.float32, True),
])
def forward(self, img):
return self.mobilenetv3.forward(img)
@register_test_case(module_factory=lambda: MobilenetV3Module())
def MobilenetV3Module_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 3, 224, 224))