pax_global_header00006660000000000000000000000064143671355410014523gustar00rootroot0000000000000052 comment=1c23fcdde49b4b286f76abaab78bfe5174b09cc5 msadapter/000077500000000000000000000000001436713554100130475ustar00rootroot00000000000000msadapter/.drone.yml000066400000000000000000000004351436713554100147610ustar00rootroot00000000000000kind: pipeline type: kubernetes name: linter platform: os: linux arch: amd64 trigger: event: - push - pull_request steps: - name: Code Inspection image: swr.cn-north-4.myhuaweicloud.com/hanjr/msadapter:2.0.0.dev20230103_torch1.12.1 commands: - sh run.shmsadapter/.gitignore000066400000000000000000000005471436713554100150450ustar00rootroot00000000000000# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class *.egg-info/ # Jupyter Notebook .ipynb_checkpoints # Environments .env .venv env/ # IDE Specific directories .DS_Store .idea .vscode/ # Data Files and ByteCode files *.gz *.npz *.ckpt *.pytest_cache *.dat # packaging build/ dist/ eggs/ .eggs/ sdist/ var/ wheels/ #datasets/ #mnist/msadapter/CONTRIBUTING_CN.md000066400000000000000000000163771436713554100156760ustar00rootroot00000000000000# MSAdapter贡献指南 ## 1.环境安装 测试环境选择Python3.7.5、MindSpore2.0.0Nightly、PyTorch1.12.0 ### 1.1安装MindSpore 我们选择使用最新版本的MindSpore开发，请使用2.0.0Nightly版本。安装指令可以参考[MindSpore安装。](https://mindspore.cn/install) 例如需要安装CPU版本的MindSpore2.0.0Nightly可以用如下指令： ``` pip install mindspore-dev -i https://pypi.tuna.tsinghua.edu.cn/simple ``` ### 1.2安装PyTorch 安装PyTorch是用于和开发的MSAdapter接口进行结果和功能比对，选择PyTorch1.12.0，安装可以参考[PyTorch安装。](https://pytorch.org/get-started/previous-versions/) 例如需要安装CPU版本PyTorch1.12.0可以用如下指令： ``` pip install torch==1.12.0+cpu torchvision==0.13.0+cpu torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cpu ``` ## 2.开发指南 ### 2.1目录结构说明 接口代码开发主要涉及到nn的modules里面文件及nn.functional.py和functional.py；分别对应开发模型构建组件接口代码（对应torch.nn.Conv2d这类），nn函数接口代码（对应torch.nn.functional.conv2d这类）以及常用函数接口代码（对应torch.abs这些）。 * ├── nn * │ ├── modules 模型构建组件层 * │ │ ├── __init__.py * │ │ ├── activation.py 激活函数层接口 * │ │ └── batchnorm.py BN层接口 * │ │ └── ... * │ └── functionanl.py nn.functional接口 * ├── optim * │ └── init * ├── utils * └── functional.py 函数接口 代码测试用例编写目录结构如下，在开发完成接口后需要在对应的文件夹内创建测试代码 * |testing * ├── ut * │ ├── pytorch * │ │ ├── data * │ │ ├── function 函数接口测试用例位置 * │ │ └── nn 模型构建组件类测试用例位置 * │ │ └── ... ### 2.2深度学习框架前端层次结构说明 深度学习框架前端层次结构可以描述为如下表示。低阶接口包含Tensor、nn.functional.xx等；中阶接口是对低阶接口的封装，调用低阶接口实现用户友好的API；高阶接口又可以是中低阶接口的封装。因此我们开发中阶接口Layers（nn.Conv2d、nn.MaxPool2d）相关的API时,需要先开发对应的低级接口如nn.functional.xx(nn.functional.conv2d、nn.functional.max_pool2d)。 ![深度学习框架层次结构](https://openi.pcl.ac.cn/laich/pose_data/raw/branch/master/%e5%b1%82%e6%ac%a1%e7%bb%93%e6%9e%84.png) ### 2.3 模型构建组件开发示例 假设我们需要构建一个和PyTorch线性层（Linear）对应的接口。首先应该开发Linear所调用的低阶接口nn.functional.linear 其中低阶接口linear可以使用MindSpore的MatMul和bias_add算子，示例如下： ``` def linear(input, weight, bias=None): input = cast_to_ms_tensor(input) output = ms.ops.MatMul(transpose_b=True)(input, weight) if bias is not None: output = ms.ops.bias_add(output, bias) output = cast_to_adapter_tensor(output) return output ``` 需要注意的是开发函数接口调用MindSpore算子时需要将tensor调用cast_to_ms_tensor转换成MindSpore所需tensor，输出调用cast_to_adapter_tensor转换成MSAdapter所需tensor。 接下来就可以使用nn.functional.linear提供的接口来开发更高阶的神经网络组件接口Linear了，示例如下： ``` class Linear(Module): def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): super(Linear, self).__init__() self.in_features = in_features self.out_features = out_features self.has_bias = False self.bias = None self.weight = Parameter(empty((self.out_features, self.in_features)), requires_grad=True) if bias: self.bias = Parameter(empty(self.out_features), requires_grad=True) self.has_bias = True self.reset_parameters() unsupported_attr(device) unsupported_attr(dtype) def reset_parameters(self): init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.has_bias: fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 init.uniform_(self.bias, -bound, bound) def forward(self, input): input = cast_to_ms_tensor(input) x = linear(input, self.weight, self.bias) return cast_to_adapter_tensor(x) def extra_repr(self): return 'in_features={}, out_features={}, bias={}'.format( self.in_features, self.out_features, self.has_bias is not None ) ``` 在__init__方法里进行属性的定义和参数初始化，在reset_parameters里进行参数的初始化方法定义，在forward里实现前向计算过程。extra_repr里返回属性信息。 此时Linear Layer和其nn.functioal接口都已经开发完成，需要编写测试用例，在testing/ut/pytorch/nn下新建测试文件test_linear.py。我们先简单测试一下训练参数属性能否和torch一样改变初始化方法。也需要测试相同输入和初始化方法输出是否和torch一样，具体可以参考[test_conv.py](https://openi.pcl.ac.cn/OpenI/MSAdapter/src/branch/master/testing/ut/pytorch/nn/test_conv.py)。 ``` #!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch.nn import Module, Linear, Identity, Bilinear from ms_adapter.pytorch import tensor from mindspore import context import numpy as np import mindspore as ms context.set_context(mode=ms.PYNATIVE_MODE) def test_linear_model(): class LinearModel(Module): def __init__(self): super(LinearModel, self).__init__() self.line1 = Linear(in_features=32, out_features=64) self.line2 = Linear(in_features=64, out_features=128, bias=False) self.line3 = Linear(in_features=128, out_features=10) def forward(self, inputs): x = self.line1(inputs) x = self.line2(x) x = self.line3(x) return x model = LinearModel() model.train() def weight_init(m): if isinstance(m, Linear): m.weight.data.normal_(0, 0.01) if m.has_bias: m.bias.data.zero_() model.apply(weight_init) inputs = tensor(np.ones(shape=(5, 32))) output = model(inputs) assert output.shape == (5, 10) test_linear_model() ``` ### 2.4代码提交到仓库 测试 用例编写完成，此时就开发完成了。我们可以提交代码到代码仓库。 首先从代码仓库clone代码到本地机器： ``` git clone https://openi.pcl.ac.cn/OpenI/MSAdapter.git ``` clone代码后，本地开发，我们新建一个分支进行开发： ``` git checkout -b {新分支名称} ``` 一顿开发完成后，我们正式推送更新分支： ``` git add . git status # 查看更新了什么代码，避免提交错了 git commit -m "你的commit标题，例如开发了Linear起名为：Add Linear Layer" git push origin {新分支名称} ``` 接下来就在启智平台进行操作 点击： 合并请求-->创建合并请求 我们继续检查一下开发内容是否正确，没问题后确认创建合并请求，等待审核入库。msadapter/LICENSE000066400000000000000000000262561436713554100140670ustar00rootroot00000000000000Copyright (c) 2022 TLX2ONNX contributors. All Rights Reserved Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. msadapter/README.md000066400000000000000000000075701436713554100143370ustar00rootroot00000000000000# MSAdapter 简体中文 | [English](README_en.md) ## 简介 MSAdapter是MindSpore适配PyTorch接口的工具,其目的是在不改变原有PyTorch用户的使用习惯情况下,使得PyTorch代码能在昇腾上获得高效性能.

- PyTorch接口支持： MSAdapter的目的是支持PyTorch语法的原生态表达，用户只需要将PyTorch源代码中```import torch```替换为```import ms_adapter.pytorch```即可实现模型能支持昇腾上训练。模型中所使用的高阶APIs支持状态可以从这里找到 [Supported List](SupportedList.md) - PyTorch接口支持范围： MSAdapter目前主要适配PyTorch的数据处理和模型结构部分代码，目前完全支持MindSpore的PYNATIVE模式下训练，部分网络结构支持GRAPH模式训练。训练过程部分代码需要用户自定义编写具体使用可以参考[使用指南](USER_GUIDE.md) ## 安装 ### 安装MindSpore 请根据MindSpore官网[安装指南](https://www.mindspore.cn/install) ,安装2.0.0Nightly版本的MindSpore。 ### 安装MSAdapter #### 通过pip安装 ```bash pip install ms_adapter ``` #### 通过源码安装 ```bash git clone https://git.openi.org.cn/OpenI/MSAdapter.git cd MSAdapter python setup.py install ``` 如果出现权限不足的提示，请按照如下方式安装： ```bash python setup.py install --user || exit 1 ``` ## 使用 在数据处理和模型构建上,MSAdapter可以和PyTorch一样使用,模型训练部分代码需要自定义,示例如下： ### 1.数据处理(仅修改导入包) ```python from ms_adapter.pytorch.utils.data import DataLoader from ms_adapter.torchvision import datasets, transforms transform = transforms.Compose([transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) ]) train_images = datasets.CIFAR10('./', train=True, download=True, transform=transform) train_data = DataLoader(train_images, batch_size=128, shuffle=True, num_workers=2, drop_last=True) ``` ### 2.模型构建(仅修改导入包) ```python from ms_adapter.pytorch.nn import Module, Linear, Flatten class MLP(Module): def __init__(self): super(MLP, self).__init__() self.flatten = Flatten() self.line1 = Linear(in_features=1024, out_features=64) self.line2 = Linear(in_features=64, out_features=128, bias=False) self.line3 = Linear(in_features=128, out_features=10) def forward(self, inputs): x = self.flatten(inputs) x = self.line1(x) x = self.line2(x) x = self.line3(x) return x ``` ### 3.模型训练(自定义训练) ```python import ms_adapter.pytorch as torch import ms_adapter.pytorch.nn as nn import mindspore as ms net = MLP() net.train() epochs = 500 criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0005) # 定义训练过程 loss_net = ms.nn.WithLossCell(net, criterion) train_net = ms.nn.TrainOneStepCell(loss_net, optimizer) for i in range(epochs): for X, y in train_data: res = train_net(X, y) print("epoch:{}, loss:{:.6f}".format(i, res.asnumpy())) # 模型保存 ms.save_checkpoint(net, "save_path.ckpt") ``` ## 资源 - 模型库：MSAdapter支持丰富的深度学习应用，这里给出了从PyTorch官方代码迁移到MSAdapter模型。[已验证模型资源](https://git.openi.org.cn/OpenI/MSAdapterModelZoo) ## 贡献 欢迎开发者参与贡献。更多详情，请参阅我们的[贡献指南](https://openi.pcl.ac.cn/OpenI/MSAdapter/src/branch/master/CONTRIBUTING_CN.md). ## 许可证 [Apache License 2.0](https://openi.pcl.ac.cn/OpenI/MSAdapter/src/branch/master/LICENSE) msadapter/README.rst000066400000000000000000000061301436713554100145360ustar00rootroot00000000000000Introduction ============= MSAdapter is MindSpore tool for adapting the PyTorch interface, which is designed to make PyTorch code perform efficiently on Ascend without changing the habits of the original PyTorch users. |MSAdapter-architecture| Install ======= MSAdapter has some prerequisites that need to be installed first, including MindSpore, PIL, NumPy. .. code:: bash # for last stable version pip install ms_adapter # for latest release candidate pip install --upgrade --pre ms_adapter Alternatively, you can install the latest or development version by directly pulling from OpenI: .. code:: bash pip3 install git+https://openi.pcl.ac.cn/OpenI/MSAdapter.git User guide =========== For data processing and model building, MSAdapter can be used in the same way as PyTorch, while the model training part of the code needs to be customized, as shown in the following example. 1. Data processing (only modify the import package) .. code:: python from ms_adapter.pytorch.utils.data import DataLoader from ms_adapter.torchvision import datasets, transforms transform = transforms.Compose([transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) ]) train_images = datasets.CIFAR10('./', train=True, download=True, transform=transform) train_data = DataLoader(train_images, batch_size=128, shuffle=True, num_workers=2, drop_last=True) 2. Model construction (modify import package only) .. code:: python from ms_adapter.pytorch.nn import Module, Linear, Flatten class MLP(Module): def __init__(self): super(MLP, self).__init__() self.flatten = Flatten() self.line1 = Linear(in_features=1024, out_features=64) self.line2 = Linear(in_features=64, out_features=128, bias=False) self.line3 = Linear(in_features=128, out_features=10) def forward(self, inputs): x = self.flatten(inputs) x = self.line1(x) x = self.line2(x) x = self.line3(x) return x 3.Model training (custom training) .. code:: python import ms_adapter.pytorch as torch import ms_adapter.pytorch.nn as nn import mindspore as ms net = MLP() net.train() epochs = 500 criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0005) # Define the training process loss_net = ms.nn.WithLossCell(net, criterion) train_net = ms.nn.TrainOneStepCell(loss_net, optimizer) for i in range(epochs): for X, y in train_data: res = train_net(X, y) print("epoch:{}, loss:{:.6f}".format(i, res.asnumpy())) # Save model ms.save_checkpoint(net, "save_path.ckpt") License ======= MSAdapter is released under the Apache 2.0 license. .. |MSAdapter-architecture| image:: https://openi.pcl.ac.cn/laich/pose_data/raw/branch/master/MSA_F.pngmsadapter/README_en.md000066400000000000000000000103131436713554100150060ustar00rootroot00000000000000# MSAdapter [简体中文](README.md) | [English] ## Introduction MSAdapter is MindSpore tool for adapting the PyTorch interface, which is designed to make PyTorch code perform efficiently on Ascend without changing the habits of the original PyTorch users.

- PyTorch interface support: MSAdapter aims to support the original expression of PyTorch syntax, users just need to replace ``import torch`` in PyTorch source code with ``import ms_adapter.pytorch`` to realize that the model can support training on ascending. The support status of the higher-order APIs used in the model can be found here [Supported List](SupportedList.md) - PyTorch interface support scope: MSAdapter is currently mainly adapted to PyTorch data processing and model structure part of the code, currently fully supports MindSpore's PYNATIVE mode training, part of the network structure support GRAPH mode training. The training process part of the code needs to be written by the user to customize the specific use can refer to the [User Guide](USER_GUIDE.md) ## Install ### Install MindSpore Please install MindSpore version 2.0.0Nightly according to the [Installation Guide](https://www.mindspore.cn/install) on MindSpore official website. ### Install MSAdapter #### via pip ```bash pip install ms_adapter ``` #### via source code ```bash git clone https://git.openi.org.cn/OpenI/MSAdapter.git cd MSAdapter python setup.py install ``` If there is an insufficient permissions message, install as follows ```bash python setup.py install --user || exit 1 ``` ## User guide For data processing and model building, MSAdapter can be used in the same way as PyTorch, while the model training part of the code needs to be customized, as shown in the following example. ### 1.Data processing (only modify the import package) ```python from ms_adapter.pytorch.utils.data import DataLoader from ms_adapter.torchvision import datasets, transforms transform = transforms.Compose([transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) ]) train_images = datasets.CIFAR10('./', train=True, download=True, transform=transform) train_data = DataLoader(train_images, batch_size=128, shuffle=True, num_workers=2, drop_last=True) ``` ### 2.Model construction (modify import package only) ```python from ms_adapter.pytorch.nn import Module, Linear, Flatten class MLP(Module): def __init__(self): super(MLP, self).__init__() self.flatten = Flatten() self.line1 = Linear(in_features=1024, out_features=64) self.line2 = Linear(in_features=64, out_features=128, bias=False) self.line3 = Linear(in_features=128, out_features=10) def forward(self, inputs): x = self.flatten(inputs) x = self.line1(x) x = self.line2(x) x = self.line3(x) return x ``` ### 3.Model training (custom training) ```python import ms_adapter.pytorch as torch import ms_adapter.pytorch.nn as nn import mindspore as ms net = MLP() net.train() epochs = 500 criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0005) # Define the training process loss_net = ms.nn.WithLossCell(net, criterion) train_net = ms.nn.TrainOneStepCell(loss_net, optimizer) for i in range(epochs): for X, y in train_data: res = train_net(X, y) print("epoch:{}, loss:{:.6f}".format(i, res.asnumpy())) # Save model ms.save_checkpoint(net, "save_path.ckpt") ``` ## Resources - Model library: MSAdapter supports rich deep learning applications, migration to MSAdapter models from the official PyTorch code is given here. [Model Resources](https://git.openi.org.cn/OpenI/MSAdapterModelZoo). ## Contributing Developers are welcome to contribute. For more details, please see our [Contribution Guidelines](https://openi.pcl.ac.cn/OpenI/MSAdapter/src/branch/master/CONTRIBUTING_CN.md). ## License [Apache License 2.0](https://openi.pcl.ac.cn/OpenI/MSAdapter/src/branch/master/LICENSE) msadapter/SupportedList.md000066400000000000000000000056051436713554100162200ustar00rootroot00000000000000## List of PyTorch APIs supported by MSAdapter | MSAdapter APIs | Status | Notes | | --------------- | -------------------- | -------------- | | Conv1d | Supported| Pad支持不完善,权重不对齐,需要给出扩展为二维权重| | Conv2d | Supported| /| | Conv3d | Supported|Pad支持不完善 | | ConvTranspose1d |Supported |output_padding参数不支持、pad类型支持不完备 | | ConvTranspose2d |Supported |output_padding参数不支持、pad类型支持不完备 | | ConvTranspose3d |Supported |output_padding参数不支持、pad类型支持不完备 | | Linear | Supported | /| | MaxPool1d | Supported|/| | AvgPool1d | Supported|/| | MaxPool2d | Supported|/| | AvgPool2d | Supported|/| | MaxPool3d | Supported|/| | AvgPool3d | Supported|/| | AdaptiveAvgPool1d | Supported| /| | AdaptiveAvgPool2d | Supported| /| | AdaptiveAvgPool3d | Supported| /| | AdaptiveMaxPool1d | Supported| /| | AdaptiveMaxPool2d | Supported|/| | AdaptiveMaxPool3d | Supported| /| | Embedding |Supported | scale_grad_by_freq、sparse参数不支持| | Flatten | Supported| /| | Unflatten| Supported| /| | Dropout | Supported| /| |Dropout2D|Supported|/| |Dropout3D|Supported|/| | BatchNorm1d | Supported| /| | BatchNorm2d | Supported| /| | BatchNorm3d |Supported | /| | PRelu | Pending| /| | ReLU |Supported| /| | Tanh |Supported| /| | Sigmoid |Supported| /| | LeakyRelu Supported|| /| | Softplus |Supported| /| | ReLU6 | Supported| /| | LeakyReLU6 |Supported| /| |Hardtanh|Supported|/| |Hardswish|Supported|/| | Mish |Supported| /| | Softmax |Supported| /| | Elu |Supported | /| | RNN | Pending| /| | RNNCell | Pending| /| | LSTM | Pending| /| | LSTMCell | Pending| /| | GRU | Pending| /| | GRUCell | Pending| /| | FractionalMaxPool2d| Supported| /| | FractionalMaxPool3d| Supported| /| | LPPool1d| Supported| /| | LPPool2d| Supported| /| | ReflectionPad1d| Supported| /| | ReflectionPad2d| Supported| /| | ReflectionPad3d| Supported| /| | ReplicationPad2d| Supported| /| | ReplicationPad3d| Supported| /| | ConstantPad1d| Supported| /| | ConstantPad2d| Supported| /| | ConstantPad3d| Supported| /| | Tanhshrink| Supported| /| | Threshold| Supported| /| | GLU| Supported| /| | Softmin| Supported| /| | LogSoftmax| Supported| /| | SyncBatchNorm| Supported| /| | GroupNorm| Supported| 只支持2D| | LayerNorm| Supported| /| | AlphaDropout| Supported| /| | FeatureAlphaDropout| Supported| /| | CosineSimilarity| Supported| /| | PairwiseDistance| Supported| /| | L1Loss| Supported| /| | MSELoss| Supported| /| | CrossEntropyLoss| Supported| /| | NLLLoss| Supported| /| | BCELoss| Supported| /| | BCEWithLogitsLoss| Supported| /| | HuberLoss| Supported| /| | SmoothL1Loss| Supported| /| | SoftMarginLoss| Supported| /| | CosineEmbeddingLoss| Supported| /| | MultiMarginLoss| Supported| /| | TripletMarginLoss| Supported| /| | Upsample| Supported| /| | UpsamplingNearest2d| Supported| /| | UpsamplingBilinear2d| Supported| /| | | | | | | | | msadapter/USER_GUIDE.md000066400000000000000000000143151436713554100150700ustar00rootroot00000000000000# MSAdapter用户手册 ## 1.简介 MSAdapter是MindSpore适配PyTorch接口的工具,其目的是在不改变原有PyTorch用户的使用习惯情况下,使得PyTorch代码能在昇腾上获得高效性能，用户只需要将PyTorch源代码中import torch替换为import ms_adapter.pytorch即可实现模型能支持昇腾上训练。 ## 2.使用指南 目前MSAdapte支持MindSpore的PYNATIVE模式训练，GRAPH模式训练需要调整代码，训练过程部分代码需要用户自定义。我们做了很多接口的对齐工作，但是由于框架设计原因仍然有部分接口在MindSpore里不需要或者对不齐需要删除或修改. ### 2.1 冗余代码删除部分 在昇腾平台使用MindSpore不需要将模型或者数据给到GPU，遇到torch.cuda或者to gpu、cpu的操作都可以删除。如： * torch.cuda * torch.Tensor.to * torch.nn.Module.to 对于no_grad、grad等梯度操作，在MindSpore中不需要 * torch.no_grad ### 2.2 数据处理部分转换 MSAdapter的数据处理接口除了DataLoader中pin_memory=True不生效外，均PyTorch几乎完全对齐。仅需将数据处理相关导入包修改为从ms_adapter导入,示例如下 ```python from ms_adapter.pytorch.utils.data import DataLoader from ms_adapter.torchvision import datasets, transforms transform = transforms.Compose([transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) ]) train_images = datasets.CIFAR10('./', train=True, download=True, transform=transform) train_data = DataLoader(train_images, batch_size=128, shuffle=True, num_workers=2, drop_last=True) ``` ### 2.3 模型构建部分转换 我们开发了大量和PyTorch功能和输出一致的模型构建接口，仅需将模型构建组件导入包改为从ms_adapter导入,需要注意的是有些操作仍然不支持，如自定义反向传播。如果有一些必要的接口和功能可以通过[ISSUE](https://openi.pcl.ac.cn/OpenI/MSAdapter/issues) 向我们反馈. ```python from ms_adapter.pytorch.nn import Module, Linear, Flatten class MLP(Module): def __init__(self): super(MLP, self).__init__() self.flatten = Flatten() self.line1 = Linear(in_features=1024, out_features=64) self.line2 = Linear(in_features=64, out_features=128, bias=False) self.line3 = Linear(in_features=128, out_features=10) def forward(self, inputs): x = self.flatten(inputs) x = self.line1(x) x = self.line2(x) x = self.line3(x) return x ``` ### 2.4 模型训练部分定义 #### 2.4.1 简单模型训练 对于简单的模型训练可以调用MindSpore的原生接口，如mindspore.Model.train或者ms.nn.WithLossCell和ms.nn.TrainOneStepCell单步训练方式。具体使用如下： 方式一：函数式迭代训练 ```python from mindspore import ops from ms_adapter.pytorch import nn import ms_adapter.pytorch as torch model = LeNet() criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), learning_rate=0.1, momentum=0.9, weight_decay=1e-4) # 定义前向过程 def forward_fn(data, label): logits = model(data) loss = criterion(logits, label) return loss, logits # 计算梯度 grad_fn = ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True) # 单步训练 def train_step(data, label): (loss, _), grads = grad_fn(data, label) loss = ops.depend(loss, optimizer(grads)) return loss # 数据迭代训练 for i, (input, target) in enumerate(train_loader): loss = train_step(input, target) ``` 方式二：使用MindSpore的Model.train训练 ```python import mindspore as ms from mindspore.dataset import GeneratorDataset from mindspore.train.callback import LossMonitor, TimeMonitor model = LeNet() criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), learning_rate=0.1, momentum=0.9, weight_decay=1e-4) model = ms.Model(model, criterion, optimizer, metrics={'accuracy'}) dataset = GeneratorDataset(source=train_data, column_names=["data", "label"]) model.train(epochs, dataset, callbacks=[TimeMonitor(), LossMonitor()]) ``` 方式三：使用WithLossCell和TrainOneStepCell迭代训练 ```python import mindspore as ms from ms_adapter.pytorch import nn import ms_adapter.pytorch as torch model = LeNet() criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), learning_rate=0.1, momentum=0.9, weight_decay=1e-4) loss_net = ms.nn.WithLossCell(model, criterion) train_net = ms.nn.TrainOneStepCell(loss_net, optimizer) for i in range(epochs): for X, y in train_data: loss = train_net(X, y) ``` #### 2.4.2 自定义loss模型训练 #### 2.4.3 多个model交替训练 ## 3.性能提升方法 ### 3.1 使用Graph模式 ### 3.2 使用混合精度 ### 3.3 使用分布式训练 ## 4.常见问题 ### 4.1 训练问题 - 刚训练或者训练到一半时突然报错RuntimeError: std::bad_alloc可能是设置的batch size过大导致，修改小一些再次训练。 ## FAQ Q：设置context.set_context(mode=context.GRAPH_MODE)后运行出现类似问题：`Tensor.add_` is an in-place operation and "x.add_()" is not encouraged to use in MindSpore static graph mode. Please use "x = x.add()" or other API instead。 A：目前在设置GRAPH模式下不支持原地操作相关的接口，需要按照提示信息进行修改。需要注意的是，即使在PYNATIVE模式下,原地操作相关接口也是不鼓励使用的，因为目前在MSAdapter不会带来内存收益，而且会给反向梯度计算带来不确定性。 Q：运行代码出现类似报错信息：AttributeError: module 'ms_adapter.pytorch' has no attribute 'xxx'。 A：首先确定'xxx'是否为torch 1.12版本支持的接口，PyTorch官网明确已废弃或者即将废弃的接口和参数，MSAdapter不会兼容支持，请使用其他同等功能的接口代替。如果是PyTorch对应版本支持，而MSAdapter中暂时没有，欢迎参与[MSAdapter项目](https://openi.pcl.ac.cn/OpenI/MSAdapter)贡献你的代码，也可以通过[创建任务(New issue)](https://openi.pcl.ac.cn/OpenI/MSAdapter/issues/new)反馈需求。 msadapter/ms_adapter/000077500000000000000000000000001436713554100151665ustar00rootroot00000000000000msadapter/ms_adapter/__init__.py000066400000000000000000000002251436713554100172760ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter import pytorch from ms_adapter.utils import unsupported_attr, pynative_mode_condition msadapter/ms_adapter/pytorch/000077500000000000000000000000001436713554100166565ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/__init__.py000066400000000000000000000036551436713554100210000ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- # Register MSAdapter Tensor/Parameter to MindSpore, it should be executed at the top of all. from ms_adapter.pytorch._register import * from ms_adapter.pytorch.common import * from ms_adapter.pytorch.tensor import * from ms_adapter.pytorch import nn from ms_adapter.pytorch import optim from ms_adapter.pytorch.functional import * from ms_adapter.pytorch.utils import data from ms_adapter.pytorch._ref import * from ms_adapter.pytorch import cuda from ms_adapter.pytorch.conflict_functional import * import ms_adapter.pytorch.fft as fft from ms_adapter.pytorch import autograd from ms_adapter.pytorch.package_info import __version__, VERSION, version from ms_adapter.pytorch.random import * from ms_adapter.pytorch.storage import * # Variables with simple values, from math.py e = 2.718281828459045 pi = 3.141592653589793 tau = 6.283185307179586 def _assert(condition, message): assert condition, message def is_tensor(obj): r"""Returns True if `obj` is a ms_adapter.pytorch tensor. Note that this function is simply doing ``isinstance(obj, Tensor)``. Using that ``isinstance`` check is better for typechecking with mypy, and more explicit - so it's recommended to use that instead of ``is_tensor``. """ return isinstance(obj, Tensor) def is_floating_point(obj): # TODO: return mindspore.ops.is_floating_point(obj) if not is_tensor(obj): raise TypeError("is_floating_point(): argument 'input' (position 1) must be Tensor, not {}.".format(type(obj))) return obj._dtype in (mstype.float16, mstype.float32, mstype.float64) class Size(tuple): def __new__(cls, shape): if isinstance(shape, Tensor): _shape = shape.tolist() else: _shape = shape if not isinstance(_shape, (tuple, list)): raise TypeError("{} object is not supportted.".format(type(shape))) return tuple.__new__(Size, _shape) msadapter/ms_adapter/pytorch/_ref/000077500000000000000000000000001436713554100175715ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/_ref/__init__.py000066400000000000000000000024531436713554100217060ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.tensor import Tensor, cast_to_adapter_tensor def randn(*size, out=None, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) if isinstance(size[0], (tuple, list)): _size = size[0] elif isinstance(size[0], int): _size = size else: raise TypeError("`size` type in `randn` only support int, tuple and list") if dtype is None: dtype = ms.float32 out_value = ms.numpy.randn(_size, dtype=dtype) if out is not None: ms.ops.assign(out, out_value) return out return cast_to_adapter_tensor(out_value) def typename(o): if isinstance(o, Tensor): return o.type() module = '' class_name = '' if hasattr(o, '__module__') and o.__module__ != 'builtins' \ and o.__module__ != '__builtin__' and o.__module__ is not None: module = o.__module__ + '.' if hasattr(o, '__qualname__'): class_name = o.__qualname__ elif hasattr(o, '__name__'): class_name = o.__name__ else: class_name = o.__class__.__name__ return module + class_name msadapter/ms_adapter/pytorch/_register/000077500000000000000000000000001436713554100206415ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/_register/__init__.py000066400000000000000000000054011436713554100227520ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from mindspore.common.api import set_adapter_config from mindspore._extends.parse import trope as T from mindspore._extends.parse.resources import convert_object_map from ms_adapter.pytorch.tensor import Tensor from ms_adapter.pytorch.nn import Parameter from ms_adapter.pytorch._register import register_multitype_ops from ms_adapter.pytorch._register import register_standard_method as S convert_object_map[T.add] = register_multitype_ops.add # x+y convert_object_map[T.sub] = register_multitype_ops.sub # x-y convert_object_map[T.mul] = register_multitype_ops.mul # x*y convert_object_map[T.truediv] = register_multitype_ops.div # x/y convert_object_map[T.getitem] = register_multitype_ops.getitem # x[0] convert_object_map[T.setitem] = register_multitype_ops.setitem # x[0]=y convert_object_map[T.floordiv] = register_multitype_ops.floordiv # x//y convert_object_map[T.mod] = register_multitype_ops.mod # x%y convert_object_map[T.pow] = register_multitype_ops.pow_ # x**y convert_object_map[T.and_] = register_multitype_ops.bitwise_and # x&y convert_object_map[T.or_] = register_multitype_ops.bitwise_or # x|y convert_object_map[T.xor] = register_multitype_ops.bitwise_xor # x^y convert_object_map[T.neg] = register_multitype_ops.negative # -x convert_object_map[T.not_] = register_multitype_ops.logical_not # not x convert_object_map[T.eq] = register_multitype_ops.equal # x==y convert_object_map[T.ne] = register_multitype_ops.not_equal # x!=y convert_object_map[T.lt] = register_multitype_ops.less # x < y convert_object_map[T.gt] = register_multitype_ops.greater # x > y convert_object_map[T.le] = register_multitype_ops.less_equal # x <= y convert_object_map[T.ge] = register_multitype_ops.greater_equal # x >= y convert_object_map[T.contains] = register_multitype_ops.in_ # x in y convert_object_map[T.not_contains] = register_multitype_ops.not_in_ # x not in y convert_object_map[T.matmul] = S.adapter_matmul # x @ y convert_object_map[T.invert] = S.adapter_invert # ~x convert_object_map[T.abs] = S.adapter_abs # abs(x) convert_object_map[T.round] = S.adapter_round # round(x) convert_object_map[T.max] = S.adapter_max # max(x) convert_object_map[T.min] = S.adapter_min # min(x) convert_object_map[T.sum] = S.adapter_sum # sum(x) def register_msadapter_tensor(): adapter_config = {"Tensor": Tensor, "convert_object_map": convert_object_map} set_adapter_config(adapter_config) register_msadapter_tensor() msadapter/ms_adapter/pytorch/_register/getitem_impl.py000066400000000000000000000027621436713554100237010ustar00rootroot00000000000000 #!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms from mindspore import dtype as mstype from mindspore.ops.composite.multitype_ops import _compile_utils as compile_utils def _tensor_getitem_by_tensor(data, tensor_index): if tensor_index.dtype == mstype.bool_: ms_shape_len = len(data.shape) index_shape_len = len(tensor_index.shape) out_shape = [-1] while index_shape_len < ms_shape_len: out_shape.append(data.shape[index_shape_len]) tensor_index = tensor_index.expand_dims(-1) index_shape_len += 1 out = ms.ops.masked_select(data, tensor_index) if len(out_shape) > 1: out = out.reshape(out_shape) else: out = compile_utils.tensor_index_by_tensor(data, tensor_index) return out def _tensor_getitem_by_number(data, number_index): if isinstance(number_index, bool): if number_index: return data.expand_dims(0) else: index = ms.Tensor(False) out = ms.ops.masked_select(data, index) return out return compile_utils.tensor_index_by_number(data, number_index) def _tensor_getitem_by_tuple(data, tuple_index): if isinstance(tuple_index[0], bool): if False in tuple_index: index = ms.Tensor(False) out = ms.ops.masked_select(data, index) return out else: return data.expand_dims(0) return compile_utils.tensor_index_by_tuple(data, tuple_index) msadapter/ms_adapter/pytorch/_register/register_multitype_ops.py000066400000000000000000000151571436713554100260450ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from mindspore.ops.composite.multitype_ops.add_impl import add from mindspore.ops.composite.multitype_ops.sub_impl import sub from mindspore.ops.composite.multitype_ops.mul_impl import mul from mindspore.ops.composite.multitype_ops.div_impl import div from mindspore.ops.composite.multitype_ops.floordiv_impl import floordiv from mindspore.ops.composite.multitype_ops.mod_impl import mod from mindspore.ops.composite.multitype_ops.pow_impl import pow_ from mindspore.ops.composite.multitype_ops.bitwise_and_impl import bitwise_and from mindspore.ops.composite.multitype_ops.bitwise_or_impl import bitwise_or from mindspore.ops.composite.multitype_ops.bitwise_xor_impl import bitwise_xor from mindspore.ops.composite.multitype_ops.negative_impl import negative from mindspore.ops.composite.multitype_ops.logic_not_impl import logical_not from mindspore.ops.composite.multitype_ops.equal_impl import equal from mindspore.ops.composite.multitype_ops.not_equal_impl import not_equal from mindspore.ops.composite.multitype_ops.less_impl import less from mindspore.ops.composite.multitype_ops.greater_impl import greater from mindspore.ops.composite.multitype_ops.less_equal_impl import less_equal from mindspore.ops.composite.multitype_ops.greater_equal_impl import greater_equal from mindspore.ops.composite.multitype_ops.in_impl import in_ from mindspore.ops.composite.multitype_ops.not_in_impl import not_in_ from mindspore.ops.composite.multitype_ops.getitem_impl import getitem from mindspore.ops.composite.multitype_ops.setitem_impl import setitem from ms_adapter.pytorch._register import register_utils as utils from ms_adapter.pytorch._register.getitem_impl import _tensor_getitem_by_tensor, _tensor_getitem_by_number, \ _tensor_getitem_by_tuple # multitype_ops.add utils.update_multitype_ops_tensor_tensor(add) utils.update_multitype_ops_number_tensor(add) utils.update_multitype_ops_tensor_number(add) utils.update_multitype_ops_tuple_tensor(add) utils.update_multitype_ops_tensor_tuple(add) utils.update_multitype_ops_list_tensor(add) utils.update_multitype_ops_tensor_list(add) # multitype_ops.sub utils.update_multitype_ops_tensor_tensor(sub) utils.update_multitype_ops_number_tensor(sub) utils.update_multitype_ops_tensor_number(sub) utils.update_multitype_ops_tuple_tensor(sub) utils.update_multitype_ops_tensor_tuple(sub) utils.update_multitype_ops_list_tensor(sub) utils.update_multitype_ops_tensor_list(sub) # multitype_ops.mul utils.update_multitype_ops_tensor_tensor(mul) utils.update_multitype_ops_number_tensor(mul) utils.update_multitype_ops_tensor_number(mul) utils.update_multitype_ops_tuple_tensor(mul) utils.update_multitype_ops_tensor_tuple(mul) utils.update_multitype_ops_list_tensor(mul) utils.update_multitype_ops_tensor_list(mul) # multitype_ops.div utils.update_multitype_ops_tensor_tensor(div) utils.update_multitype_ops_number_tensor(div) utils.update_multitype_ops_tensor_number(div) utils.update_multitype_ops_tuple_tensor(div) utils.update_multitype_ops_tensor_tuple(div) utils.update_multitype_ops_list_tensor(div) utils.update_multitype_ops_tensor_list(div) # multitype_ops.floordiv utils.update_multitype_ops_tensor_tensor(floordiv) utils.update_multitype_ops_number_tensor(floordiv) utils.update_multitype_ops_tensor_number(floordiv) utils.update_multitype_ops_tuple_tensor(floordiv) utils.update_multitype_ops_tensor_tuple(floordiv) utils.update_multitype_ops_list_tensor(floordiv) utils.update_multitype_ops_tensor_list(floordiv) # multitype_ops.mod utils.update_multitype_ops_tensor_tensor(mod) utils.update_multitype_ops_number_tensor(mod) utils.update_multitype_ops_tensor_number(mod) utils.update_multitype_ops_tuple_tensor(mod) utils.update_multitype_ops_tensor_tuple(mod) utils.update_multitype_ops_list_tensor(mod) utils.update_multitype_ops_tensor_list(mod) # multitype_ops.pow_ utils.update_multitype_ops_tensor_tensor(pow_) utils.update_multitype_ops_number_tensor(pow_) utils.update_multitype_ops_tensor_number(pow_) utils.update_multitype_ops_tuple_tensor(pow_) utils.update_multitype_ops_tensor_tuple(pow_) utils.update_multitype_ops_list_tensor(pow_) utils.update_multitype_ops_tensor_list(pow_) # multitype_ops.bitwise_and utils.update_multitype_ops_tensor_tensor(bitwise_and) utils.update_multitype_ops_number_tensor(bitwise_and) utils.update_multitype_ops_tensor_number(bitwise_and) # multitype_ops.bitwise_or utils.update_multitype_ops_tensor_tensor(bitwise_or) utils.update_multitype_ops_number_tensor(bitwise_or) utils.update_multitype_ops_tensor_number(bitwise_or) # multitype_ops.bitwise_xor utils.update_multitype_ops_tensor_tensor(bitwise_xor) utils.update_multitype_ops_number_tensor(bitwise_xor) utils.update_multitype_ops_tensor_number(bitwise_xor) # multitype_ops.negative utils.update_multitype_ops_tensor(negative) # multitype_ops.logical_not # LogicalNot only support Tensor[Bool]. utils.update_multitype_ops_tensor(logical_not) # multitype_ops.equal utils.update_multitype_ops_tensor_tensor(equal) utils.update_multitype_ops_number_tensor(equal) utils.update_multitype_ops_tensor_number(equal) # multitype_ops.not_equal utils.update_multitype_ops_tensor_tensor(not_equal) utils.update_multitype_ops_number_tensor(not_equal) utils.update_multitype_ops_tensor_number(not_equal) # multitype_ops.less utils.update_multitype_ops_tensor_tensor(less) utils.update_multitype_ops_number_tensor(less) utils.update_multitype_ops_tensor_number(less) # multitype_ops.greater utils.update_multitype_ops_tensor_tensor(greater) utils.update_multitype_ops_number_tensor(greater) utils.update_multitype_ops_tensor_number(greater) # multitype_ops.less_equal utils.update_multitype_ops_tensor_tensor(less_equal) utils.update_multitype_ops_number_tensor(less_equal) utils.update_multitype_ops_tensor_number(less_equal) # multitype_ops.greater_equal utils.update_multitype_ops_tensor_tensor(greater_equal) utils.update_multitype_ops_number_tensor(greater_equal) utils.update_multitype_ops_tensor_number(greater_equal) # multitype_ops.in_ utils.update_multitype_ops_tensor_tuple(in_) utils.update_multitype_ops_tensor_list(in_) # multitype_ops.not_in_ utils.update_multitype_ops_tensor_tuple(not_in_) utils.update_multitype_ops_tensor_list(not_in_) # multitype_ops.getitem utils.update_multitype_ops_tensor_list(getitem) utils.update_multitype_ops_tensor_none(getitem) utils.update_multitype_ops_tensor_slice(getitem) utils.update_multitype_ops_tensor_tensor_with_fn(getitem, _tensor_getitem_by_tensor) utils.update_multitype_ops_tensor_number_with_fn(getitem, _tensor_getitem_by_number) utils.update_multitype_ops_tensor_tuple_with_fn(getitem, _tensor_getitem_by_tuple) # multitype_ops.setitem utils.update_multitype_ops_setitem_tensor(setitem) msadapter/ms_adapter/pytorch/_register/register_standard_method.py000066400000000000000000000051001436713554100262530ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from mindspore import dtype as mstype from mindspore._extends.parse import trope as T from mindspore._extends.parse.resources import convert_object_map from ms_adapter.pytorch.tensor import Tensor as adapter_Tensor from ms_adapter.pytorch._register.register_utils import convert_to_ms_tensor, convert_to_adapter_tensor matmul_fn = convert_object_map.get(T.matmul) invert_fn = convert_object_map.get(T.invert) abs_fn = convert_object_map.get(T.abs) round_fn = convert_object_map.get(T.round) max_fn = convert_object_map.get(T.max) min_fn = convert_object_map.get(T.min) sum_fn = convert_object_map.get(T.sum) def adapter_matmul(x, y): if isinstance(x, adapter_Tensor) and isinstance(y, adapter_Tensor): x = convert_to_ms_tensor(x) y = convert_to_ms_tensor(y) out = matmul_fn(x, y) out = convert_to_adapter_tensor(out) else: out = matmul_fn(x, y) return out def adapter_invert(x): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) if x.dtype != mstype.bool_: out = - 1 - x else: out = invert_fn(x) out = convert_to_adapter_tensor(out) else: out = invert_fn(x) return out def adapter_abs(x): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) out = abs_fn(x) out = convert_to_adapter_tensor(out) else: out = abs_fn(x) return out def adapter_round(*data): if (len(data) == 1 and isinstance(data[0], adapter_Tensor)) or \ (len(data) == 2 and isinstance(data[0], adapter_Tensor) and data[1] is None): x = data[0] x = convert_to_ms_tensor(x) out = round_fn(x) out = convert_to_adapter_tensor(out) else: out = round_fn(*data) return out def _has_adapter_tensor(*data): if len(data) == 1 and isinstance(data[0], adapter_Tensor): return True for elem in data: if isinstance(elem, adapter_Tensor): return True return False def adapter_max(*data): if _has_adapter_tensor(*data): out = max_fn(*data) out = convert_to_adapter_tensor(out) else: out = max_fn(*data) return out def adapter_min(*data): if _has_adapter_tensor(*data): out = min_fn(*data) out = convert_to_adapter_tensor(out) else: out = min_fn(*data) return out def adapter_sum(*data): if _has_adapter_tensor(*data): out = sum_fn(*data) out = convert_to_adapter_tensor(out) else: out = sum_fn(*data) return out msadapter/ms_adapter/pytorch/_register/register_utils.py000066400000000000000000000177351436713554100242740ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms from mindspore import dtype as mstype from mindspore.ops.operations import _inner_ops as inner from ms_adapter.pytorch.tensor import Tensor as adapter_Tensor def convert_to_ms_tensor(x): return inner.convert_to_ms_tensor(x) def convert_to_adapter_tensor(x): return inner.convert_to_adapter_tensor(x) def convert_output(out): if isinstance(out, ms.Tensor): out = convert_to_adapter_tensor(out) return out def get_registed_fn(ops, *type_names): types = tuple(map(mstype.typing.str_to_type, type_names)) for sigs, fn in ops.entries: if len(sigs) != len(types): continue if any(not mstype._issubclass_(type_, sig) for sig, type_ in zip(sigs, types)): continue return fn raise ValueError(f"For 'MultitypeFuncGraph', cannot find fn match given types: {types}.") def _multitype_ops_tensor_calcu(ops, func1, func2): @ops.register("Tensor") def _tensor(x): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) out = func1(x) out = convert_output(out) else: out = func2(x) return out def update_multitype_ops_tensor_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tensor") _multitype_ops_tensor_calcu(ops, func1, func2) def update_multitype_ops_tensor(ops): func = get_registed_fn(ops, "Tensor") _multitype_ops_tensor_calcu(ops, func, func) def _multitype_ops_tensor_tensor_calcu(ops, func1, func2): @ops.register("Tensor", "Tensor") def _tensor_and_tensor(x, y): if isinstance(x, adapter_Tensor) and isinstance(y, adapter_Tensor): x = convert_to_ms_tensor(x) y = convert_to_ms_tensor(y) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_tensor_tensor_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tensor", "Tensor") _multitype_ops_tensor_tensor_calcu(ops, func1, func2) def update_multitype_ops_tensor_tensor(ops): func = get_registed_fn(ops, "Tensor", "Tensor") _multitype_ops_tensor_tensor_calcu(ops, func, func) def _multitype_ops_number_tensor_calcu(ops, func1, func2): @ops.register("Number", "Tensor") def _number_and_tensor(x, y): if isinstance(y, adapter_Tensor): y = convert_to_ms_tensor(y) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_number_tensor_with_fn(ops, func1): func2 = get_registed_fn(ops, "Number", "Tensor") _multitype_ops_number_tensor_calcu(ops, func1, func2) def update_multitype_ops_number_tensor(ops): func = get_registed_fn(ops, "Number", "Tensor") _multitype_ops_number_tensor_calcu(ops, func, func) def _multitype_ops_tensor_number_calcu(ops, func1, func2): @ops.register("Tensor", "Number") def _tensor_and_number(x, y): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_tensor_number_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tensor", "Number") _multitype_ops_tensor_number_calcu(ops, func1, func2) def update_multitype_ops_tensor_number(ops): func = get_registed_fn(ops, "Tensor", "Number") _multitype_ops_tensor_number_calcu(ops, func, func) def _multitype_ops_tuple_tensor_calcu(ops, func1, func2): @ops.register("Tuple", "Tensor") def _tuple_and_tensor(x, y): if isinstance(y, adapter_Tensor): y = convert_to_ms_tensor(y) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_tuple_tensor_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tuple", "Tensor") _multitype_ops_tuple_tensor_calcu(ops, func1, func2) def update_multitype_ops_tuple_tensor(ops): func = get_registed_fn(ops, "Tuple", "Tensor") _multitype_ops_tuple_tensor_calcu(ops, func, func) def _multitype_ops_tensor_tuple_calcu(ops, func1, func2): @ops.register("Tensor", "Tuple") def _tensor_and_tuple(x, y): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_tensor_tuple_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tensor", "Tuple") _multitype_ops_tensor_tuple_calcu(ops, func1, func2) def update_multitype_ops_tensor_tuple(ops): func = get_registed_fn(ops, "Tensor", "Tuple") _multitype_ops_tensor_tuple_calcu(ops, func, func) def _multitype_ops_list_tensor_calcu(ops, func1, func2): @ops.register("List", "Tensor") def _list_and_tensor(x, y): if isinstance(y, adapter_Tensor): y = convert_to_ms_tensor(y) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_list_tensor_with_fn(ops, func1): func2 = get_registed_fn(ops, "List", "Tensor") _multitype_ops_list_tensor_calcu(ops, func1, func2) def update_multitype_ops_list_tensor(ops): func = get_registed_fn(ops, "List", "Tensor") _multitype_ops_list_tensor_calcu(ops, func, func) def _multitype_ops_tensor_list_calcu(ops, func1, func2): @ops.register("Tensor", "List") def _tensor_and_list(x, y): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_tensor_list_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tensor", "List") _multitype_ops_tensor_list_calcu(ops, func1, func2) def update_multitype_ops_tensor_list(ops): func = get_registed_fn(ops, "Tensor", "List") _multitype_ops_tensor_list_calcu(ops, func, func) def _multitype_ops_tensor_none_calcu(ops, func1, func2): @ops.register("Tensor", "None") def _tensor_and_none(x, y): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_tensor_none_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tensor", "None") _multitype_ops_tensor_none_calcu(ops, func1, func2) def update_multitype_ops_tensor_none(ops): func = get_registed_fn(ops, "Tensor", "None") _multitype_ops_tensor_none_calcu(ops, func, func) def _multitype_ops_tensor_slice_calcu(ops, func1, func2): @ops.register("Tensor", "Slice") def _tensor_and_slice(x, y): if isinstance(x, adapter_Tensor): x = convert_to_ms_tensor(x) out = func1(x, y) out = convert_output(out) else: out = func2(x, y) return out def update_multitype_ops_tensor_slice_with_fn(ops, func1): func2 = get_registed_fn(ops, "Tensor", "Slice") _multitype_ops_tensor_slice_calcu(ops, func1, func2) def update_multitype_ops_tensor_slice(ops): func = get_registed_fn(ops, "Tensor", "Slice") _multitype_ops_tensor_slice_calcu(ops, func, func) def update_multitype_ops_setitem_tensor(ops): def register_for_setitem(sigs, fn): @ops.register(*sigs) def _tensor_setitem(data, index, value): if isinstance(data, adapter_Tensor): data = convert_to_ms_tensor(data) out = fn(data, index, value) out = convert_to_adapter_tensor(out) else: out = fn(data, index, value) return out entries = ops.entries.copy() for sigs, fn in entries: if mstype._issubclass_(sigs[0], mstype.tensor): register_for_setitem(sigs, fn) msadapter/ms_adapter/pytorch/_six.py000066400000000000000000000000361436713554100201710ustar00rootroot00000000000000string_classes = (str, bytes) msadapter/ms_adapter/pytorch/_utils.py000066400000000000000000000020771436713554100205350ustar00rootroot00000000000000import sys import traceback # class KeyErrorMessage(str): # r"""str subclass that returns itself in repr""" # def __repr__(self): # return self class ExceptionWrapper: r"""Wraps an exception plus traceback to communicate across threads""" def __init__(self, exc_info=None, where="in background"): if exc_info is None: exc_info = sys.exc_info() self.exc_type = exc_info[0] self.exc_msg = "".join(traceback.format_exception(*exc_info)) self.where = where def reraise(self): r"""Reraises the wrapped exception in the current thread""" msg = "Caught {} {}.\nOriginal {}".format( self.exc_type.__name__, self.where, self.exc_msg) if self.exc_type == KeyError: # msg = KeyErrorMessage(msg) msg = str(msg) elif getattr(self.exc_type, "message", None): raise self.exc_type(message=msg) try: exception = self.exc_type(msg) except TypeError: raise RuntimeError(msg) from None raise exception msadapter/ms_adapter/pytorch/autograd/000077500000000000000000000000001436713554100204645ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/autograd/__init__.py000066400000000000000000000004651436713554100226020ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from .variable import Variable from .function import Function # MindSpore's autodiff mechanism is different from PyTorch' autograd， so it cannot be fully benchmarked. # Users can directly use the autograd API of MindSpore. __all__ = ["Variable", "Function"] msadapter/ms_adapter/pytorch/autograd/function.py000066400000000000000000000031621436713554100226650ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import warnings from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.nn import Module class Function(Module): """Base class to create custom `autograd.Function` Examples:: >>> class Exp(Function): >>> def __init__(self): >>> super(Exp, self).__init__() >>> >>> def forward(self, i): >>> result = i.exp() >>> return result >>> >>> def bprop(self, i, out, grad_output): >>> return grad_output * out >>> >>> # Use non-static forward method: >>> output = Exp()(input) """ def __init__(self, *args, **kwargs): unsupported_attr(args) unsupported_attr(kwargs) super(Function, self).__init__() warnings.warn("'autograd.Function' is the same as 'nn.Module', and do not support other Function methods") def apply(self, *args, **kwargs): """ # Don not use it by calling the apply method. """ unsupported_attr(args) unsupported_attr(kwargs) raise RuntimeError("To create a custom autograd.Function, use 'def forward(self, ...)' and " "'def bprop(self, ..., out, dout)' instead of 'forward()' and 'backward()' static methods. " "Then, use it as normal module class, do not call the class method 'apply'." " Please refer to the following example: " " https://openi.pcl.ac.cn/OpenI/MSAdapter/src/branch/master/" "testing/ut/pytorch/autograd/test_autograd_function.py") msadapter/ms_adapter/pytorch/autograd/variable.py000066400000000000000000000010221436713554100226160ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.tensor import Tensor class Variable(Tensor): def __init__(self, data, grad=None, requires_grad=None, volatile=None, creator=None): unsupported_attr(data) unsupported_attr(grad) unsupported_attr(requires_grad) unsupported_attr(volatile) unsupported_attr(creator) msg = "The Variable API has been deprecated, use Tensor instead." raise RuntimeError(msg) msadapter/ms_adapter/pytorch/common/000077500000000000000000000000001436713554100201465ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/common/__init__.py000066400000000000000000000010571436713554100222620ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch.common.dtype import * from ms_adapter.pytorch.common.device import * __all__ = ["float", "double", "float16", "float32", "float64", "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64", "bool_", "complex64", "complex128", "long", "bfloat16", "cfloat", "cdouble", "half", "short", "int", "bool", "iinfo", "finfo", "Device" ] msadapter/ms_adapter/pytorch/common/_inner.py000066400000000000000000000041611436713554100217740ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from mindspore.ops import constexpr from ms_adapter.pytorch.tensor import cast_to_adapter_tensor, Tensor from ms_adapter.utils import pynative_mode_condition, graph_mode_condition def _out_limit_pynative(out, op_name): if out is not None and graph_mode_condition(): # TODO: ms_function raise ValueError('In MindSpore static graph mode, `out` in `{}` shoud be None, ' 'please set out=None and use return value instead of `out`.'.format(op_name)) def _out_inplace_assign(out, output, op_name): if out is None: return cast_to_adapter_tensor(output) if pynative_mode_condition(): # TODO: ms_function def _assign(out, output): if isinstance(out, Tensor): out.assign_value(output) elif isinstance(out, (tuple, list)): for item in zip(out, output): _assign(item[0], item[1]) _assign(out, output) return out raise ValueError('In MindSpore static graph mode, `out` in `{}` shoud be None, ' 'please set out=None and use return value instead of `out`.'.format(op_name)) def _inplace_assign_pynative(input, inplace, output, op_name): if inplace is True: if pynative_mode_condition(): # TODO: ms_function input.assign_value(output) return input raise ValueError('In MindSpore static graph mode, `inplace` in `{}` shoud not be Ture, ' 'please set inplace=False and use return value instead of `input`.'.format(op_name)) return cast_to_adapter_tensor(output) @constexpr def _inplace_limit_pynative(inplace, op_name): if inplace is True and graph_mode_condition(): # TODO: ms_function raise ValueError('In MindSpore static graph mode, `inplace` in `{}` shoud not be Ture, ' 'please set inplace=False and use return value instead of `input`.'.format(op_name)) def _inplace_assign(input, inplace, output): if inplace is True: input.assign_value(output) return input return cast_to_adapter_tensor(output) msadapter/ms_adapter/pytorch/common/device.py000066400000000000000000000004271436713554100217620ustar00rootroot00000000000000 class Device(): def __init__(self, target, index): self.type = target self.index = index def __repr__(self): if self.index is None: return f"device(type='{self.type}')" return f"device(type='{self.type}', index={self.index})" msadapter/ms_adapter/pytorch/common/dtype.py000066400000000000000000000043721436713554100216530ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np from mindspore import dtype as mstype float = mstype.float32 double = mstype.float64 float16 = mstype.float16 # TODO: mindspore to support mstype.bfloat16 bfloat16 = mstype.float32 float32 = mstype.float32 float64 = mstype.float64 int8 = mstype.int8 int16 = mstype.int16 int32 = mstype.int32 int64 = mstype.int64 uint8 = mstype.uint8 uint16 = mstype.uint16 uint32 = mstype.uint32 uint64 = mstype.uint64 bool_ = mstype.bool_ complex64 = mstype.complex64 complex128 = mstype.complex128 long = mstype.int64 cfloat = mstype.complex64 cdouble = mstype.complex128 half = mstype.half short = mstype.short int = mstype.int32 bool = mstype.bool_ char = mstype.uint8 all_int_type = (mstype.int8, mstype.int16, mstype.int32, mstype.int64, mstype.uint8, ) _TypeDict = {mstype.float16: np.float16, mstype.float32: np.float32, mstype.float64: np.float64, mstype.int8: np.int8, mstype.int16: np.int16, mstype.int32: np.int32, mstype.int64: np.int64, mstype.uint8: np.uint8} class iinfo: def __init__(self, dtype): if dtype in (mstype.uint8, mstype.int8, mstype.int16, mstype.int32, mstype.int64): np_iinfo = np.iinfo(_TypeDict[dtype]) self.bits = np_iinfo.bits self.max = np_iinfo.max self.min = np_iinfo.min else: raise ValueError("iinfo currently only supports torch.uint8/torch.int8/torch.int16/torch.int32/" "torch.int64 as the input, but get a", dtype) class finfo: def __init__(self, dtype): if dtype in (mstype.float16, mstype.float32, mstype.float64): np_finfo = np.finfo(_TypeDict[dtype]) self.bits = np_finfo.bits self.eps = np_finfo.eps self.max = np_finfo.max self.min = np_finfo.min self.tiny = np_finfo.tiny # TODO: numpy vision >= 1.23 # self.smallest_normal = np_finfo.smallest_normal self.resolution = np_finfo.resolution else: raise ValueError("finfo currently only supports torch.float16/torch.float32/" "torch.float64 as the input, but get a", dtype) msadapter/ms_adapter/pytorch/conflict_functional.py000066400000000000000000000034741436713554100232630ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms from mindspore.common import dtype as mstype from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.common._inner import _out_inplace_assign def range(start, end, step=1, out=None, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) if dtype is None: dtype = ms.float32 start = ms.Tensor(start, dtype=dtype) end = ms.Tensor(end+0.001, dtype=dtype) # TODO This function is deprecated and will be removed in a future release # because its behavior is inconsistent with Python’s range builtin. Instead, use torch.arange(), # which produces values in [start, end). step = ms.Tensor(step, dtype=dtype) output = ms.ops.range(start, end, step) return _out_inplace_assign(out, output, "range") def arange(start, end, step=1, *, out=None, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) # TODO: use code below in future version # output = ms.ops.arange(start, end, step) # return _out_inplace_assign(out, output, "arange") if dtype is None: if isinstance(start, float) or isinstance(end, float) or isinstance(step, float): dtype = mstype.float32 else: #TODO # For now, `range` do not support `mstype.int64`, it should be changed # to 'dtype = mstype.int64' in mindspore 2.0 dtype = mstype.int32 start = ms.Tensor(start, dtype) end = ms.Tensor(end, dtype) step = ms.Tensor(step, dtype) output = ms.ops.range(start=start, limit=end, delta=step) return _out_inplace_assign(out, output, "arange") msadapter/ms_adapter/pytorch/cuda/000077500000000000000000000000001436713554100175725ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/cuda/__init__.py000066400000000000000000000015731436713554100217110ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms from mindspore.communication.management import init, get_group_size from ms_adapter.utils import get_backend from ms_adapter.pytorch.tensor import FloatTensor, LongTensor def is_available(): backend = get_backend() if backend == 'GPU': return True return False def current_device(): return 0 def device_count(): # TODO Use this method when supported # init() # return get_group_size() return 1 def set_device(device): if device in ("gpu", 'GPU'): ms.context.set_context(device_target="GPU") elif device in ("cpu", "CPU"): ms.context.set_context(device_target="CPU") elif device in ("ascend", "Ascend"): ms.context.set_context(device_target="Ascend") else: raise ValueError("device must be cpu, gpu, ascend or CPU, GPU, Ascend.") msadapter/ms_adapter/pytorch/fft/000077500000000000000000000000001436713554100174355ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/fft/__init__.py000066400000000000000000000001251436713554100215440ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from .fft import * __all__ = ['fft'] msadapter/ms_adapter/pytorch/fft/fft.py000066400000000000000000000007321436713554100205700ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import mindspore as ms from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor def fft(input, n=None, dim=-1, norm=None, out=None): input = cast_to_ms_tensor(input) input = input.asnumpy() output = np.fft.fft(input, n, axis=dim, norm=norm) output = cast_to_adapter_tensor(ms.Tensor(output)) if out is not None: out.assign_value(output) return output msadapter/ms_adapter/pytorch/functional.py000066400000000000000000002353161436713554100214040ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import collections import warnings import numpy as np from scipy import signal import mindspore as ms from mindspore import ops from mindspore.common import dtype as mstype from mindspore.ops import constexpr from mindspore.common.seed import _get_graph_seed from mindspore.ops._primitive_cache import _get_cache_prim from ms_adapter.pytorch.tensor import tensor, cast_to_ms_tensor, cast_to_adapter_tensor, _div_calcu from ms_adapter.utils import unsupported_attr, get_backend, pynative_mode_condition, is_under_gpu_context, \ is_under_ascend_context, _ascend_tensor_general_cast, _infer_size from ms_adapter.pytorch.tensor import Tensor as adapter_tensor from ms_adapter.pytorch.common._inner import _out_inplace_assign, _out_limit_pynative from ms_adapter.pytorch.common.dtype import _TypeDict, all_int_type from ms_adapter.pytorch.common.device import Device def empty(*size, out=None, dtype=None, layout=None, \ device=None, requires_grad=False, pin_memory=False, \ memory_format=None): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(pin_memory) unsupported_attr(memory_format) if dtype is None: dtype = ms.float32 _size = size if isinstance(size[0], (tuple, list)): _size = size[0] output = ms.numpy.empty(_size, dtype) return _out_inplace_assign(out, output, "empty") def eye(n, m=None, *, out=None, dtype=None, layout=None, \ device=None, requires_grad=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) if m is None: m = n if dtype is None: dtype = ms.float32 output = ms.ops.eye(n, m, dtype) return _out_inplace_assign(out, output, "eye") def cat(tensors, dim=0, *, out=None): if tensors is None: raise ValueError('`tensors` in `{}` should not be None'.format(cat.__name__)) if not isinstance(tensors, (tuple, list)): raise TypeError('`tensors` in `{}` should be tuple or list'.format(cat.__name__)) if is_under_ascend_context(): _rank = len(tensors[0].shape) dim = dim if dim >= 0 else dim + _rank inputs = cast_to_ms_tensor(tensors) output = ops.concat(inputs, dim) return _out_inplace_assign(out, output, "cat") def concat(tensors, dim=0, *, out=None): if tensors is None: raise ValueError('`tensors` in `{}` should not be None'.format(concat.__name__)) if not isinstance(tensors, (tuple, list)): raise TypeError('`tensors` in `{}` should be tuple or list'.format(concat.__name__)) if is_under_ascend_context(): _rank = len(tensors[0].shape) dim = dim if dim >= 0 else dim + _rank inputs = cast_to_ms_tensor(tensors) output = ops.concat(inputs, dim) return _out_inplace_assign(out, output, "concat") def ones(*size, out=None, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) if dtype is None: dtype = ms.float32 if isinstance(size[0], (tuple, list)): output = ms.ops.ones(*size, dtype=dtype) else: output = ms.ops.ones(size, dtype=dtype) return _out_inplace_assign(out, output, "ones") def stack(tensors, dim = 0, *, out=None): tensors = cast_to_ms_tensor(tensors) output = ops.stack(tensors, dim) return _out_inplace_assign(out, output, "stack") def meshgrid(*tensors, indexing='ij'): if isinstance(tensors[0], (list, tuple)): input_tensor = tuple(*tensors) else: input_tensor = tensors input_tensor = cast_to_ms_tensor(input_tensor) output = ops.meshgrid(input_tensor, indexing=indexing) return cast_to_adapter_tensor(output) def log(input, *, out=None): input = cast_to_ms_tensor(input) output = ops.log(input) return _out_inplace_assign(out, output, "log") def chunk(input, chunks, dim=0): input = cast_to_ms_tensor(input) output = ops.chunk(input, chunks, dim) return cast_to_adapter_tensor(output) def diag(input, diagonal=0, *, out=None): # TODO # May be use mindspore.ops.diag instead. Nowadays, this operator do not support CPU. input = cast_to_ms_tensor(input) output = ms.numpy.diag(input, diagonal) return _out_inplace_assign(out, output, "diag") def sqrt(input, *, out=None): if input.dtype == mstype.int32 or input.dtype == mstype.int64: input = input.astype(mstype.float32) input = cast_to_ms_tensor(input) output = ops.sqrt(input) return _out_inplace_assign(out, output, "sqrt") def mm(input, mat2, *, out=None): output_type = input.dtype if input.dtype == mstype.int32 or input.dtype == mstype.int64: input = input.astype(mstype.float32) input1 = cast_to_ms_tensor(input) input2 = cast_to_ms_tensor(mat2) output = ops.matmul(input1, input2) output = ops.cast(output, output_type) return _out_inplace_assign(out, output, "mm") def zeros(*size, out=None, dtype=None, device=None, requires_grad=False): unsupported_attr(device) unsupported_attr(requires_grad) if isinstance(size[0], (tuple, list)): size = tuple(size[0]) if dtype is None: dtype = mstype.float32 output = ms.ops.zeros(size, dtype) return _out_inplace_assign(out, output, "zeros") def div(input, other, *, rounding_mode=None, out=None): output = _div_calcu(input, other, rounding_mode) return _out_inplace_assign(out, output, "div") def divide(input, other, *, rounding_mode=None, out=None): output = _div_calcu(input, other, rounding_mode) return _out_inplace_assign(out, output, "divide") def flatten(input, start_dim=0, end_dim=-1): @constexpr def get_dst_shape(): input_shape = input.shape rank = len(input_shape) start = start_dim end = end_dim if start < 0: start += rank if end < 0: end += rank dst_shape = [] i = 0 while i != start: dst_shape.append(input_shape[i]) i = i + 1 flatten_shape = 1 while i <= end: flatten_shape = flatten_shape * input_shape[i] i = i + 1 dst_shape.append(flatten_shape) while i < rank: dst_shape.append(input_shape[i]) i = i + 1 return tuple(dst_shape) shape = get_dst_shape() input = cast_to_ms_tensor(input) output = ops.reshape(input, shape) return cast_to_adapter_tensor(output) def unflatten(input, dim, sizes): @constexpr def get_unflatten_size(): input_shape = input.shape input_rank = len(input_shape) if not isinstance(sizes, (tuple, list)): raise TypeError(f"Type of `sizes` should be `Tuple` or `List`, but got {type(sizes)}") if len(sizes) == 0: raise ValueError("`sizes` must be non-empty") if isinstance(dim, str): raise TypeError("Until Now, `dim` not support type of str in `unflatten`") _dim = dim if _dim < 0: _dim += input_rank if _dim < 0 or _dim >= input_rank: raise ValueError("`dim` should be in range [{}, {}), but got {}".format( -input_rank, input_rank, dim)) input_shape_list = list(input_shape) sizes_list = list(sizes) _sizes_mul = 1 for s in sizes_list: _sizes_mul *= s if _sizes_mul != input_shape_list[_dim]: raise ValueError(f"unflatten: Provided `sizes` {sizes_list} don't multiply up to the" f"size of dim {dim} ({input_shape_list[_dim]}) in the input tensor") out_shape = input_shape[:_dim] + tuple(sizes) + input_shape[_dim + 1:] return out_shape out_shape = get_unflatten_size() input = cast_to_ms_tensor(input) out = ms.ops.reshape(input, out_shape) return cast_to_adapter_tensor(out) def transpose(input, dim0, dim1): @constexpr def _check_dim(dim, rank): if dim >= rank or dim < -rank: raise ValueError("dim is out of bound, should be in range [{}, {})" .format(-rank, rank)) @constexpr def _get_perm(): rank = len(input.shape) _check_dim(dim0, rank) _check_dim(dim1, rank) _perm = list(range(rank)) _perm[dim0] = dim1 _perm[dim1] = dim0 return tuple(_perm) @constexpr def _get_perm_ascend(): rank = len(input.shape) _check_dim(dim0, rank) _check_dim(dim1, rank) _perm = list(range(rank)) _dim0 = dim0 if dim0 >=0 else dim0 + rank _dim1 = dim1 if dim1 >=0 else dim1 + rank _perm[_dim0] = _dim1 _perm[_dim1] = _dim0 return tuple(_perm) if is_under_ascend_context(): input_ms = cast_to_ms_tensor(input) out = ms.ops.transpose(input_ms, _get_perm_ascend()) return cast_to_adapter_tensor(out) input_ms = cast_to_ms_tensor(input) out = ops.transpose(input_ms, _get_perm()) return cast_to_adapter_tensor(out) def multinomial(input, num_samples, replacement=False, *, generator=None, out=None): unsupported_attr(generator) if generator is not None: warnings.warn("torch.multinomal don't support generator now.") input_tensor = cast_to_ms_tensor(input).astype(mstype.float32) # TODO: Ascend to support 1-dimention input if is_under_ascend_context(): _input_rank = len(input.shape) if _input_rank == 1: input_tensor = input_tensor.expand_dims(axis=0) if not replacement: output = ms.ops.multinomial(input_tensor, num_samples, replacement) else: # TODO: Multinomial not support on Ascend seed0, seed1 = _get_graph_seed(None, "multinomial") _op = ms.ops.Multinomial(seed0, seed1) _op.set_device("CPU") output = _op(input_tensor, num_samples) if _input_rank == 1: output = output[0] return _out_inplace_assign(out, output, "multinomial") output = ms.ops.multinomial(input_tensor, num_samples, replacement) return _out_inplace_assign(out, output, "multinomial") def randperm(n, *, generator=None, out=None, dtype=mstype.int64, layout=None, device=None, requires_grad=False, pin_memory=False): unsupported_attr(generator) unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(pin_memory) if generator is not None: warnings.warn("torch.randperm don't support generator now.") if layout is not None: warnings.warn("torch.randperm don't support layout now.") output = np.random.permutation(n) output = tensor(output, dtype=dtype) return _out_inplace_assign(out, output, "randperm") def randint(low, high, size, *, generator=None, out=None, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(generator) unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) if generator is not None: warnings.warn("torch.randperm don't support generator now.") if layout is not None: warnings.warn("torch.randperm don't support layout now.") output = np.random.randint(low, high, size) output = tensor(output, dtype=dtype) return _out_inplace_assign(out, output, "randint") def as_tensor(data, dtype=None, device=None): unsupported_attr(device) if isinstance(data, (tuple, list)): data = [i.data.item() if isinstance(i, adapter_tensor) else i for i in data ] output = ms.Tensor(data, dtype=dtype) return cast_to_adapter_tensor(output) def zeros_like(input, dtype=None, layout=None, device=None, requires_grad=False, memory_format=None): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(memory_format) input_x = cast_to_ms_tensor(input) # TODO: output = ms.ops.zeros_like(input_x, dtype=dtype) output = ms.ops.zeros_like(input_x) if dtype: output = output.astype(dtype) return cast_to_adapter_tensor(output) def ones_like(input, dtype=None, layout=None, device=None, requires_grad=False, memory_format=None): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(memory_format) input_x = ms.Tensor(input, dtype=dtype) output = ms.ops.ones_like(input_x) return cast_to_adapter_tensor(output) def empty_like(input, dtype=None, layout=None, device=None, requires_grad=False, memory_format=None): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(memory_format) if dtype is None: dtype = ms.float32 #TODO:replace whith mindspore.ops.empty_like ouput = ms.numpy.empty_like(input, dtype=dtype) return cast_to_adapter_tensor(ouput) def full(size, fill_value, out=None, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) #TODO:replace whith mindspore.ops.full output = ms.numpy.full(size, fill_value, dtype) return _out_inplace_assign(out, output, "full") def full_like(input, fill_value, dtype=None, layout=None, device=None, requires_grad=False, memory_format=None): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(memory_format) #TODO:replace whith mindspore.ops.full_like output = ms.numpy.full_like(input, fill_value=fill_value, dtype=dtype) return cast_to_adapter_tensor(output) def where(condition, x=None, y=None): if x is None and y is None: return nonzero(condition, as_tuple=True) # TODO:replace whith mindspore.ops.where output = ms.numpy.where(condition, x, y) return cast_to_adapter_tensor(output) def rand(*size, out=None, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(pin_memory) if dtype is None: dtype = ms.float32 output = ms.numpy.rand(*size, dtype=dtype) return _out_inplace_assign(out, output, "rand") def linspace(start, end, steps, out=None, dtype=None, device=None, requires_grad=False): unsupported_attr(device) unsupported_attr(requires_grad) if dtype is None: dtype = ms.float32 start = ms.Tensor(start, dtype) end = ms.Tensor(end, dtype) output = ms.ops.linspace(start, end, steps) return _out_inplace_assign(out, output, "linspace") def take(input, index): input = cast_to_ms_tensor(input) input = ms.numpy.array(input) index = ms.numpy.array(index) # TODO:replace whith mindspore.ops.take output = ms.numpy.take(input, index) return cast_to_adapter_tensor(output) def abs(input, out=None): input = cast_to_ms_tensor(input) output = ms.ops.abs(input) return _out_inplace_assign(out, output, "abs") def atan2(input, other, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.atan2(input, other) return _out_inplace_assign(out, output, "atan2") def clamp(input, min=None, max=None, out=None): # TODO:replace whith mindspore.ops.clamp input_ms = cast_to_ms_tensor(input) type = input_ms.dtype if min is not None and max is not None and min > max: output = ms.ops.ones_like(input_ms).astype(type)*max else: if min is not None: min = ms.Tensor(min, type) if max is not None: max = ms.Tensor(max, type) output = ms.ops.clip_by_value(input_ms, min, max) return _out_inplace_assign(out, output, "clamp") def cos(input, out=None): input = cast_to_ms_tensor(input) output = ms.ops.cos(input) return _out_inplace_assign(out, output, "cos") @constexpr def device(type=None, index=None): if type is not None: if isinstance(type, str): if ':' in type: if index is not None: raise ValueError(f"`type` must not include an index because index was passed explicitly: {type}") _target, _id = type.split(':') _id = int(_id) else: _target = type _id = index return Device(_target, _id) if isinstance(type, int): return Device(get_backend(), type) if isinstance(type, Device): if index is not None: raise ValueError("torch.device(): When input is torch.device, `index` can not be set.") return Device(type.type, type.index) raise TypeError("torch.device(): `type` must be type of 'str' or 'torch.device'.") raise ValueError("torch.device(): `type` can not be None") def fmod(input, other, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) #TODO: repalce with ms.ops.fmod output = input - ms.ops.div(input, other, rounding_mode="trunc") * other return _out_inplace_assign(out, output, "fmod") def frac(input, out=None): #TODO outout = input - floor(abs(input)) * sing(input) input = cast_to_ms_tensor(input) input = ms.numpy.array(input) output = input - ms.numpy.floor(ms.numpy.abs(input)) * ms.numpy.sign(input) return _out_inplace_assign(out, output, "frac") def log10(input, out=None): input = cast_to_ms_tensor(input) input = ms.numpy.array(input) output = ms.numpy.log10(input) return _out_inplace_assign(out, output, "log10") def log1p(input, out=None): input = cast_to_ms_tensor(input) output = ms.ops.log1p(input) return _out_inplace_assign(out, output, "log1p") def log2(input, out=None): input = cast_to_ms_tensor(input) output = ms.ops.log2(input) return _out_inplace_assign(out, output, "log2") def sin(input, out=None): input = cast_to_ms_tensor(input) output = ms.ops.sin(input) return _out_inplace_assign(out, output, "sin") def max(input, dim=None, keepdim=False, *, out=None): #TODO: not supprt GRAPH_MODE, not supper max(input, other) input = cast_to_ms_tensor(input) type = input.dtype input = input.astype(ms.float32) if dim is None: output = input.max(axis=dim, keepdims=keepdim).astype(type) if out is not None: ops.assign(out, output) return out return cast_to_adapter_tensor(output) output = list(ms.ops.max(input, axis=dim, keep_dims=keepdim)) value = output[1].astype(type) indice = output[0] point = collections.namedtuple('max', 'values,indices') rlt = point(cast_to_adapter_tensor(value), cast_to_adapter_tensor(indice)) if out is not None: if pynative_mode_condition(): if len(out) != 2 or not isinstance(out[0], adapter_tensor) or not isinstance(out[1], adapter_tensor): raise TypeError("In ms_adapter.torch.max(), `out` should be tuple of Tensors.") out[0].assign_value(value) out[1].assign_value(indice) return out else: raise ValueError('In MindSpore static graph mode, `out` in `max` shoud be None, ' 'please set out=None and use return value instead of `out`.') return rlt def min(input, dim=None, keepdim=False, *, out=None): # TODO: Right Now, not support 'min(input, other, *, out=None)' input = cast_to_ms_tensor(input) if dim is None: return cast_to_adapter_tensor(input.min()) indices, result = ms.ops.min(input, axis=dim, keep_dims=keepdim) if out is not None: if pynative_mode_condition(): if len(out) != 2 or not isinstance(out[0], adapter_tensor) or not isinstance(out[1], adapter_tensor): raise TypeError("In ms_adapter.torch.min(), `out` should be tuple of Tensors.") out[0].assign_value(result) out[1].assign_value(indices) return out else: raise ValueError('In MindSpore static graph mode, `out` in `min` shoud be None, ' 'please set out=None and use return value instead of `out`.') return cast_to_adapter_tensor(result), cast_to_adapter_tensor(indices) def mean(input, dim=None, keepdim=False, *, dtype=None, out=None): # TODO: not supprt GRAPH_MODE input = cast_to_ms_tensor(input) if dtype is not None: input = input.astype(dtype) if dim is not None: output = ms.ops.mean(input, axis=dim, keep_dims=keepdim) else: output = ms.ops.mean(input, keep_dims=keepdim) return _out_inplace_assign(out, output, "mean") def round(input, *, decimals=0, out=None): input = cast_to_ms_tensor(input) if decimals == 0: output = ms.ops.round(input) else: p = 10**decimals input = input*p output = ms.ops.round(input)/p return _out_inplace_assign(out, output, "round") def floor(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.floor(input) return _out_inplace_assign(out, output, "floor") def ceil(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.ceil(input) return _out_inplace_assign(out, output, "ceil") def sign(input, *, out=None): input = cast_to_ms_tensor(input) output = _get_cache_prim(ms.ops.Sign)()(input) return _out_inplace_assign(out, output, "sign") def pow(input, exponent, *, out=None): # TODO: not support input that is above 7-dimention on GPU and Ascend input = cast_to_ms_tensor(input) exponent = cast_to_ms_tensor(exponent) output = ms.ops.pow(input, exponent) return _out_inplace_assign(out, output, "pow") def exp(input, *, out=None): input = cast_to_ms_tensor(input) shape = input.shape if len(shape) > 7: input = input.flatten() if input.dtype != ms.float64: input = input.astype(ms.float32) output = ms.ops.exp(input) if len(shape) > 7: output = output.reshape(shape) return _out_inplace_assign(out, output, "exp") def ge(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.ge(input, other) return _out_inplace_assign(out, output, "ge") def gt(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.gt(input, other) return _out_inplace_assign(out, output, "gt") def le(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.le(input, other) return _out_inplace_assign(out, output, "le") def lt(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.less(input, other) return _out_inplace_assign(out, output, "lt") def sum(input, dim=None, keepdim=False, *, dtype=None, out=None): input = cast_to_ms_tensor(input) if dtype is not None: input = input.astype(dtype) elif input.dtype in (mstype.uint8, mstype.uint16, mstype.uint32, mstype.int8, mstype.int16, mstype.int32): dtype = mstype.int64 input = input.astype(dtype) output = input.sum(axis=dim, dtype=dtype, keepdims=keepdim) return _out_inplace_assign(out, output, "sum") def median(input, dim=None, keepdim=False, *, out=None): def _process_out(value, indice): point = collections.namedtuple('median', 'values,indices') rlt = point(cast_to_adapter_tensor(value), cast_to_adapter_tensor(indice)) if out is not None: if pynative_mode_condition(): if len(out) != 2 or not isinstance(out[0], adapter_tensor) or not isinstance(out[1], adapter_tensor): raise TypeError("In ms_adapter.torch.median(), `out` should be tuple of Tensors.") out[0].assign_value(value) out[1].assign_value(indice) return out else: raise ValueError('In MindSpore static graph mode, `out` in `median` shoud be None, ' 'please set out=None and use return value instead of `out`.') return rlt input = cast_to_ms_tensor(input) if is_under_ascend_context(): input_dtype = input.dtype input = _ascend_tensor_general_cast(input) if dim is None: _avg = ms.ops.mean(input) _distance = ms.ops.abs(input - _avg) _min_indice = _distance.argmin() out = input.flatten()[_min_indice] out = out.astype(input_dtype) return cast_to_adapter_tensor(out) else: _avg = ms.ops.mean(input, axis=dim, keep_dims=True) _distance = ms.ops.abs(input - _avg) _min_indice, _ = ms.ops.min(_distance, axis=dim, keep_dims=True) _min = ms.ops.gather_d(input, dim, _min_indice) if not keepdim: _min = ms.ops.squeeze(_min, dim) _min_indice = ms.ops.squeeze(_min_indice, dim) _min = _min.astype(input_dtype) return _process_out(_min, _min_indice) if dim is None: output, _ = ms.ops.median(input, keepdims=keepdim) return cast_to_adapter_tensor(output) else: value, indices = ms.ops.median(input, axis=dim, keepdims=keepdim) return _process_out(value, indices) def matmul(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.matmul(input, other) return _out_inplace_assign(out, output, "matmul") def norm(input, p='fro', dim=None, keepdim=False, out=None, dtype=None): input = cast_to_ms_tensor(input) if dtype is None: dtype = ms.float32 input = ms.numpy.array(input, dtype=dtype) output = ms.numpy.norm(input, ord=p, axis=dim, keepdims=keepdim) return _out_inplace_assign(out, output, "norm") def stft(input, n_fft, hop_length=None, win_length=None, window=None, center=True, pad_mode='reflect', normalized=False, onesided=None, return_complex=None): unsupported_attr(normalized) unsupported_attr(onesided) unsupported_attr(return_complex) input = cast_to_ms_tensor(input) input = input.asnumpy() if pad_mode == 'reflect': pad_mode = 'even' if window is None: window = 'hann' if hop_length is None: hop_length = floor(n_fft / 4) if win_length is None: win_length = n_fft output = signal.stft(input, window=window, nperseg=win_length, noverlap=hop_length, padded=center, boundary=pad_mode) return output def istft(): raise NotImplementedError def bartlett_window(window_length, periodic=True, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) input = tensor(window_length) output = ms.ops.bartlett_window(input, periodic=periodic, dtype=dtype) return cast_to_adapter_tensor(output) def hamming_window(window_length, periodic=True, alpha=0.54, beta=0.46, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(periodic) unsupported_attr(alpha) unsupported_attr(beta) unsupported_attr(dtype) unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) output = ms.numpy.hamming(window_length) return cast_to_adapter_tensor(output) def hann_window(window_length, periodic=False, dtype=None, layout=None, device=None, requires_grad=False): unsupported_attr(periodic) unsupported_attr(dtype) unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) if periodic is True: raise NotImplementedError("periodic is not supported to True.") output = ms.numpy.hanning(window_length) return cast_to_adapter_tensor(output) def cumsum(input, dim, dtype=None, out=None): input = cast_to_ms_tensor(input) output = ms.ops.cumsum(input, axis=dim, dtype=dtype) return _out_inplace_assign(out, output, "cumsum") def einsum(equation, *operands): output = _get_cache_prim(ms.ops.Einsum)(equation=equation)(operands) return cast_to_adapter_tensor(output) def histc(input, bins=100, min=0, max=0, out=None): input = cast_to_ms_tensor(input) type = input.dtype input = input.reshape(-1) if max == min: _, max = ms.ops.max(input) _, min = ms.ops.min(input) if max == min: out = [0]*bins out[bins//2] = ms.ops.size(input) return ms.Tensor(out, type) if max < min: raise RuntimeError("histc: max must be larger than min") input_tmp = input-max max_count = ms.ops.size(input) - ms.ops.count_nonzero(input_tmp) range = max - min scale = bins/range input = (input - min) * scale + 1 input = ms.ops.clip(input, 0) input = input.astype(ms.int32) out = ms.ops.bincount(input) out = out[1:bins+1] out[-1] += max_count return out.astype(type) def triu(input, diagonal=0, out=None): input = cast_to_ms_tensor(input) input = ms.numpy.array(input) output = ms.numpy.triu(input, diagonal) output = cast_to_adapter_tensor(output) return _out_inplace_assign(out, output, "triu") def unbind(input, dim=0): input = cast_to_ms_tensor(input) output = ms.ops.unbind(input, dim) return cast_to_adapter_tensor(output) def unsqueeze(input, dim): input = cast_to_ms_tensor(input) output = ms.ops.unsqueeze(input, dim) return cast_to_adapter_tensor(output) def reshape(input, shape): input = cast_to_ms_tensor(input) input_size = input.shape if input_size[0] == 0: # only support first element is 0 numel = ms.ops.size(input) shape = _infer_size(shape, numel) output = ms.ops.zeros(shape, input.dtype) else: shape = tuple(shape) output = ms.ops.reshape(input, shape) return cast_to_adapter_tensor(output) def isfinite(input): input_ms = cast_to_ms_tensor(input) output = ms.ops.isfinite(input_ms) return cast_to_adapter_tensor(output) def isnan(input): input_ms = cast_to_ms_tensor(input) return cast_to_adapter_tensor(input_ms.isnan()) def view_as_real(input): #Todo: not view warnings.warn("not support output as a view.") input = cast_to_ms_tensor(input) input = input.asnumpy() real = np.expand_dims(np.real(input), axis=-1) imag = np.expand_dims(np.imag(input), axis=-1) output_np = np.concatenate((real, imag), axis=-1) output = ms.Tensor(output_np) return cast_to_adapter_tensor(output) def bincount(input, weights=None, minlength=0): input = cast_to_ms_tensor(input) type = 'int64' if input.dtype == ms.uint8: input = input.astype(ms.int16) if weights is not None: weights = cast_to_ms_tensor(weights) type = weights.dtype output = ms.numpy.bincount(input, weights, minlength).astype(type) return cast_to_adapter_tensor(output) def mul(input, other, *, out=None): if not isinstance(input, (int, adapter_tensor)): raise TypeError(f"mul(): argument 'input' (position 1) must be Tensor, not {type(input)}") if not isinstance(other, (int, adapter_tensor)): raise TypeError(f"mul(): argument 'other' (position 2) must be Tensor, not {type(other)}") input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.mul(input, other) return _out_inplace_assign(out, output, "mul") def index_select(input, dim, index, *, out=None): _input_params = cast_to_ms_tensor(input) _axis = dim _input_indices = cast_to_ms_tensor(index) output = ms.ops.gather(_input_params, _input_indices, _axis) return _out_inplace_assign(out, output, "index_select") def sort(input, dim=-1, descending=False, stable=False, *, out=None): unsupported_attr(stable) input = cast_to_ms_tensor(input) # TODO: ops.sort() should be replaced. output = ms.ops.Sort(dim, descending)(input) return _out_inplace_assign(out, output, "sort") def msort(input, *, out=None): input = cast_to_ms_tensor(input) # TODO: ops.sort() should be replaced. output, _ = _get_cache_prim(ms.ops.Sort)(axis=0)(input) return _out_inplace_assign(out, output, "msort") def argsort(input, dim=-1, descending=False, stable=False): unsupported_attr(stable) input = cast_to_ms_tensor(input) # TODO: ops.sort() should be replaced. _, output= _get_cache_prim(ms.ops.Sort)(dim, descending)(input) return cast_to_adapter_tensor(output) def t(input): input_ms = cast_to_ms_tensor(input) if input_ms.ndim > 2: raise ValueError("t() expects a tensor with <= 2 dimensions, but self is {}D".format(input_ms.ndim)) dims = list(range(input_ms.ndim)).reverse() output = input_ms.transpose(dims) return cast_to_adapter_tensor(output) def squeeze(input, dim=None): input_ms = cast_to_ms_tensor(input) if dim is not None: if input_ms.shape[dim] != 1: output = input else: output = ms.ops.squeeze(input_ms, dim) else: output = ms.ops.squeeze(input_ms) return cast_to_adapter_tensor(output) def from_numpy(np_data): return cast_to_adapter_tensor(ms.Tensor.from_numpy(np_data)) def absolute(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.absolute(input) return _out_inplace_assign(out, output, "absolute") def acos(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.acos(input) return _out_inplace_assign(out, output, "acos") def arccos(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.acos(input) return _out_inplace_assign(out, output, "arccos") def acosh(input, *, out=None): input = cast_to_ms_tensor(input) shape = input.shape if len(shape) > 7: input = input.flatten() output = ms.ops.acosh(input) if len(shape) > 7: output = output.reshape(shape) return _out_inplace_assign(out, output, "acosh") def arccosh(input, *, out=None): input = cast_to_ms_tensor(input) shape = input.shape if len(shape) > 7: input = input.flatten() output = ms.ops.acosh(input) if len(shape) > 7: output = output.reshape(shape) return _out_inplace_assign(out, output, "arccosh") def add(input, other, *, alpha=1, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.add(input, other*alpha) return _out_inplace_assign(out, output, "add") def addcdiv(input, tensor1, tensor2, *, value=1, out=None): input = cast_to_ms_tensor(input) tensor1 = cast_to_ms_tensor(tensor1) tensor2 = cast_to_ms_tensor(tensor2) value = ms.Tensor(value) output = ms.ops.addcdiv(input, tensor1, tensor2, value) return _out_inplace_assign(out, output, "addcdiv") def addcmul(input, tensor1, tensor2, *, value=1, out=None): #Todo: use ms.ops.addcmul after it has been fixed input = cast_to_ms_tensor(input) tensor1 = cast_to_ms_tensor(tensor1) tensor2 = cast_to_ms_tensor(tensor2) value = ms.ops.scalar_to_tensor(value) mul = ms.ops.mul(tensor1, tensor2) * value output = ms.ops.add(input, mul) return _out_inplace_assign(out, output, "addcmul") def angle(input, *, out=None): input = cast_to_ms_tensor(input) shape = input.shape if len(shape)>7: input = input.flatten() real = _get_cache_prim(ms.ops.Real)()(input) imag = _get_cache_prim(ms.ops.Imag)()(input) #Todo: ms.ops.copysign is not same as torch.copysign when input is -0.0, # replace to ms.ops.copysign after it has been fixed imag_np = imag.asnumpy() sign_imag_np = np.copysign(np.ones_like(imag_np), imag_np) sign_imag = ms.Tensor(sign_imag_np) denom = ms.ops.sqrt(ms.ops.square(real) + ms.ops.square(imag)) div = ms.ops.div(real, denom) mask = ms.ops.equal(denom, ms.Tensor(0)) mask_array = ms.ops.ones_like(denom)*sign_imag div = ms.ops.select(mask, mask_array, div) output = ms.ops.mul(ms.ops.acos(div), sign_imag) if len(shape)>7: output = output.reshape(shape) return _out_inplace_assign(out, output, "angle") def asin(input, *, out=None): input = cast_to_ms_tensor(input) if input.dtype in all_int_type: input = input.astype(mstype.float32) output = ms.ops.asin(input) return _out_inplace_assign(out, output, "asin") def arcsin(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.asin(input) return _out_inplace_assign(out, output, "arcsin") def asinh(input, *, out=None): input = cast_to_ms_tensor(input) shape = input.shape if len(shape) > 7: input = input.flatten() output = ms.ops.asinh(input) if len(shape) > 7: output = output.reshape(shape) return _out_inplace_assign(out, output, "asinh") def arcsinh(input, *, out=None): input = cast_to_ms_tensor(input) shape = input.shape if len(shape) > 7: input = input.flatten() output = ms.ops.asinh(input) if len(shape) > 7: output = output.reshape(shape) return _out_inplace_assign(out, output, "arcsinh") def atan(input, *, out=None): shape = input.shape if len(shape) > 7: input = input.flatten() input = cast_to_ms_tensor(input) output = ms.ops.atan(input) if len(shape) > 7: output = output.reshape(shape) return _out_inplace_assign(out, output, "atan") def arctan(input, *, out=None): shape = input.shape if len(shape) > 7: input = input.flatten() input = cast_to_ms_tensor(input) output = ms.ops.atan(input) if len(shape) > 7: output = output.reshape(shape) return _out_inplace_assign(out, output, "arctan") def atanh(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.atanh(input) return _out_inplace_assign(out, output, "atanh") def arctanh(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.atanh(input) return _out_inplace_assign(out, output, "arctanh") def arctan2(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.atan2(input, other) return _out_inplace_assign(out, output, "arctan2") def bitwise_not(input, *, out=None): input = cast_to_ms_tensor(input) type = input.dtype if str(type) != 'Bool': output = 0 - input - 1 else: output = 1 - input output = output.astype(ms.bool_) return _out_inplace_assign(out, output, "bitwise_not") def bitwise_and(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input) input_is_bool = str(input.dtype) == 'Bool' else: input_is_bool = isinstance(input, bool) if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other) other_is_bool = str(other.dtype) == 'Bool' else: other_is_bool = isinstance(other, bool) if input_is_bool and other_is_bool: if isinstance(input, adapter_tensor): input = input.astype(ms.int8) else: other = other.astype(ms.int8) output = ms.ops.bitwise_and(input, other) if input_is_bool and other_is_bool: output = output.astype(ms.bool_) return _out_inplace_assign(out, output, "bitwise_and") def bitwise_or(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input) input_is_bool = str(input.dtype) == 'Bool' else: input_is_bool = isinstance(input, bool) if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other) other_is_bool = str(other.dtype) == 'Bool' else: other_is_bool = isinstance(other, bool) if input_is_bool and other_is_bool: if isinstance(input, adapter_tensor): input = input.astype(ms.int8) else: other = other.astype(ms.int8) output = ms.ops.bitwise_or(input, other) if input_is_bool and other_is_bool: output = output.astype(ms.bool_) return _out_inplace_assign(out, output, "bitwise_or") def bitwise_xor(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input) input_is_bool = str(input.dtype) == 'Bool' else: input_is_bool = isinstance(input, bool) if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other) other_is_bool = str(other.dtype) == 'Bool' else: other_is_bool = isinstance(other, bool) if input_is_bool and other_is_bool: if isinstance(input, adapter_tensor): input = input.astype(ms.int8) else: other = other.astype(ms.int8) output = ms.ops.bitwise_xor(input, other) if input_is_bool and other_is_bool: output = output.astype(ms.bool_) return _out_inplace_assign(out, output, "bitwise_xor") def bitwise_left_shift(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input).asnumpy() if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other).asnumpy() output = ms.Tensor(np.left_shift(input, other)) return _out_inplace_assign(out, output, "bitwise_left_shift") def bitwise_right_shift(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input).asnumpy() if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other).asnumpy() output = ms.Tensor(np.right_shift(input, other)) return _out_inplace_assign(out, output, "bitwise_right_shift") def split(tensor, split_size_or_sections, dim=0): tensor = cast_to_ms_tensor(tensor) output = ms.ops.split(tensor, split_size_or_sections, dim) return cast_to_adapter_tensor(output) def nonzero(input, *, out=None, as_tuple=False): input = cast_to_ms_tensor(input) if as_tuple: if input.ndim == 1: res = ms.ops.nonzero(input) output = (cast_to_adapter_tensor(res.flatten()), ) elif input.ndim > 1: output = [] res = ms.ops.nonzero(input) res = res.transpose(1,0) res = ms.ops.split(res, 1, axis=0) output = cast_to_adapter_tensor(res) elif input.ndim == 0: raise ValueError("Do not support input ndim == 0.") else: output = ms.ops.nonzero(input) return _out_inplace_assign(out, output, "nonzero") def clip(input, min=None, max=None, *, out=None): input = cast_to_ms_tensor(input) output = input.clip(min, max) return _out_inplace_assign(out, output, "clip") def conj_physical(input, *, out=None): input = cast_to_ms_tensor(input) # TODO:　use code below in the future mindspore on Ascend, because Ascend not support input as real number #if ms.ops.is_complex(input): # output = ms.ops.conj(input) #else: # output = input output = ms.ops.conj(input) return _out_inplace_assign(out, output, "conj_physical") def copysign(input, other, *, out=None): input = cast_to_ms_tensor(input) input_type = input.dtype input = input.asnumpy() is_num = True if isinstance(other, adapter_tensor): is_num = False other = cast_to_ms_tensor(other) other_type = other.dtype other = other.asnumpy() output = ms.Tensor(np.copysign(input, other)) if 'Int' in str(input_type): if is_num or 'Int' in str(other_type): output = output.astype(ms.float32) else: output = output.astype(other_type) elif is_num or 'Int' in str(other_type): output = output.astype(input_type) return _out_inplace_assign(out, output, "copysign") def cosh(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.cosh(input) return _out_inplace_assign(out, output, "cosh") def deg2rad(input, *, out=None): input = cast_to_ms_tensor(input) if input.dtype not in (ms.float16, ms.float32, ms.float64): input = input.astype(ms.float32) output = ms.ops.deg2rad(input) return _out_inplace_assign(out, output, "deg2rad") def devide(input, other, *, rounding_mode=None, out=None): _out_limit_pynative(out, "devide") return div(input, other, rounding_mode=rounding_mode, out=out) def erf(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.erf(input) return _out_inplace_assign(out, output, "erf") def erfc(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.erfc(input) return _out_inplace_assign(out, output, "erfc") def erfinv(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.erfinv(input) return _out_inplace_assign(out, output, "erfinv") def exp2(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.exp2(input) return _out_inplace_assign(out, output, "exp2") def expm1(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.expm1(input) return _out_inplace_assign(out, output, "expm1") def fake_quantize_per_channel_affine(input, scale, zero_point, axis, quant_min, quant_max): input = cast_to_ms_tensor(input) scale = cast_to_ms_tensor(scale) zero_point = cast_to_ms_tensor(zero_point) if axis not in range(0, input.ndim): raise IndexError("`axis` must be between 0 and number of dimensions of input") if input.shape[axis] != scale.shape[0] or input.shape[axis] != zero_point.shape[0]: raise RuntimeError("dimensions of scale or zero-point are not consistent with input tensor") i = axis + 1 while i < input.ndim: scale = scale.expand_dims(-1) zero_point = zero_point.expand_dims(-1) i += 1 output = ms.ops.round(input/scale + zero_point) output = ms.ops.clip_by_value(output, quant_min, quant_max) - zero_point output = output * scale return cast_to_adapter_tensor(output) def fake_quantize_per_tensor_affine(input, scale, zero_point, quant_min, quant_max): input = cast_to_ms_tensor(input) scale = cast_to_ms_tensor(scale) zero_point = cast_to_ms_tensor(zero_point) output = ms.ops.round(input/scale + zero_point) output = ms.ops.clip_by_value(output, quant_min, quant_max) - zero_point output = output * scale return cast_to_adapter_tensor(output) def fix(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.trunc(input) return _out_inplace_assign(out, output, "fix") def float_power(input, exponent, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input).asnumpy() if isinstance(exponent, adapter_tensor): exponent = cast_to_ms_tensor(exponent).asnumpy() output = ms.Tensor(np.float_power(input, exponent)) return _out_inplace_assign(out, output, "float_power") def floor_divide(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.floor_div(input, other) return _out_inplace_assign(out, output, "floor_divide") def frexp(input, *, out=None): _out_limit_pynative(out, "frexp") input = cast_to_ms_tensor(input).asnumpy() mantissa, exponent = np.frexp(input) out1 = ms.Tensor(mantissa) out2 = ms.Tensor(exponent) if out is not None and len(out) != 2: out[0].assign_value(out1) out[1].assign_value(out2) return out return cast_to_adapter_tensor(out1), cast_to_adapter_tensor(out2) def gradient(input, *, spacing=1, dim=None, edge_order=1): input = cast_to_ms_tensor(input) if isinstance(spacing, adapter_tensor): spacing = cast_to_ms_tensor(spacing) elif isinstance(spacing, tuple) and isinstance(spacing[0], adapter_tensor): spacing = cast_to_ms_tensor(spacing) output = ms.numpy.gradient(input, spacing, axis=dim, edge_order=edge_order) output = cast_to_adapter_tensor(output) if not isinstance(output, tuple): return (output,) else: return output def imag(input): input = cast_to_ms_tensor(input) output = ms.ops.imag(input) return cast_to_adapter_tensor(output) def ldexp(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.ldexp(input, other) return _out_inplace_assign(out, output, "ldexp") def lerp(input, end, weight, *, out=None): input = cast_to_ms_tensor(input) end = cast_to_ms_tensor(end) if isinstance(weight, adapter_tensor): weight = cast_to_ms_tensor(weight) elif not isinstance(weight, float): weight = float(weight) output = ms.ops.lerp(input, end, weight) return _out_inplace_assign(out, output, "lerp") def logaddexp(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.logaddexp(input, other) return _out_inplace_assign(out, output, "logaddexp") def logaddexp2(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.logaddexp2(input, other) return _out_inplace_assign(out, output, "logaddexp2") def logical_and(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input).astype(ms.bool_) if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other).astype(ms.bool_) output = ms.ops.logical_and(input, other) return _out_inplace_assign(out, output, "logical_and") def logical_not(input, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input).astype(ms.bool_) output = ms.ops.logical_not(input) return _out_inplace_assign(out, output, "logical_not") def logical_or(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input).astype(ms.bool_) if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other).astype(ms.bool_) output = ms.ops.logical_or(input, other) return _out_inplace_assign(out, output, "logical_or") def logical_xor(input, other, *, out=None): if isinstance(input, adapter_tensor): input = cast_to_ms_tensor(input).astype(ms.bool_) if isinstance(other, adapter_tensor): other = cast_to_ms_tensor(other).astype(ms.bool_) # TODO: ms.ops.logical_xor to supported GPU and Ascend if is_under_ascend_context() or is_under_gpu_context(): output = ms.numpy.logical_xor(input, other) else: output = ms.ops.logical_xor(input, other) return _out_inplace_assign(out, output, "logical_xor") def logit(input, eps=None, *, out=None): #TODO: ops.logit not support cpu input = cast_to_ms_tensor(input) if eps is not None: input = ms.ops.clip_by_value(input, eps, 1.0-eps) output = ms.ops.log(input/(1.0-input)) return _out_inplace_assign(out, output, "logit") def frombuffer(buffer, *, dtype = None, count=- 1, offset=0, requires_grad=False): unsupported_attr(requires_grad) np_dtype = _TypeDict[dtype] output = np.frombuffer(buffer=buffer, dtype=np_dtype, count=count, offset=offset) return adapter_tensor(output, dtype=dtype) def as_strided(input, size, stride, storage_offset=None): warnings.warn("not support output as a view.") input_ms = cast_to_ms_tensor(input) if len(size) != len(stride): raise RuntimeError("mismatch in length of strides and shape.") index = np.arange(0, size[0] * stride[0], stride[0]) for i in range(1, len(size)): tmp = np.arange(0, size[i] * stride[i], stride[i]) index = np.expand_dims(index, -1) index = index + tmp if storage_offset is not None: index = index + storage_offset input_indices = ms.Tensor(index) out = ms.ops.gather(input_ms.reshape(-1), input_indices, 0) return cast_to_adapter_tensor(out) def ne(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.ne(input, other) return _out_inplace_assign(out, output, "ne") def tanh(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.tanh(input) return _out_inplace_assign(out, output, "tanh") def maximum(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.maximum(input, other) return _out_inplace_assign(out, output, "maximum") def minimum(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.minimum(input, other) return _out_inplace_assign(out, output, "minimum") def sigmoid(input, *, out=None): #TODO: ms.ops.sigmoid() not support float64 input = cast_to_ms_tensor(input) # output = 1 / (ms.ops.exp(0 - input) + 1) sigmoid_op = _get_cache_prim(ms.ops.Sigmoid)() if input.dtype == ms.float64: input = input.astype(ms.float32) output = sigmoid_op(input) output = output.astype(ms.float64) else: output = sigmoid_op(input) return _out_inplace_assign(out, output, "sigmoid") def softmax(input, dim, dtype=None, *, out=None): input = cast_to_ms_tensor(input) if dtype is not None: input = input.astype(dtype) output = ms.ops.softmax(input, dim) return _out_inplace_assign(out, output, "softmax") def prod(input, dim=None, keepdim=False, *, dtype=None, out=None): input = cast_to_ms_tensor(input) if dtype is not None: input = input.astype(dtype) if dim is None: output = ms.ops.prod(input) else: output = ms.ops.prod(input, axis=dim, keep_dims=keepdim) return _out_inplace_assign(out, output, "prod") def eq(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.equal(input, other) return _out_inplace_assign(out, output, "eq") def hypot(input, other, *, out=None): #TODO: can't found ms.ops.hypot() input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.sqrt(input.square() + other.square()) #output = ms.ops.hypot(input, other) return _out_inplace_assign(out, output, "hypot") def i0(input, *, out=None): input = cast_to_ms_tensor(input) float_type = [ms.float16, ms.half, ms.float32, ms.single, ms.float64, ms.double] if input.dtype not in float_type: input = input.astype(ms.float32) output = ms.ops.bessel_i0(input) return _out_inplace_assign(out, output, "i0") def _set_type_gamma(input, other): float_type = [ms.float16, ms.half, ms.float32, ms.single, ms.float64, ms.double] is_float16 = False input_type_index = -1 other_type_index = -1 if input.dtype in float_type: input_type_index = float_type.index(input.dtype) if other.dtype in float_type: other_type_index = float_type.index(other.dtype) if input_type_index < other_type_index: input = input.astype(other.dtype) elif input_type_index > other_type_index: other = other.astype(input.dtype) if input.dtype == ms.float16: input = input.astype(ms.float32) other = other.astype(ms.float32) is_float16 = True return input, other, is_float16 def igamma(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) #TODO: delete here after ms.ops.igamma support those types as pytorch input, other, is_float16 = _set_type_gamma(input, other) output = ms.ops.igamma(input, other) if is_float16: output = output.astype(ms.float16) return _out_inplace_assign(out, output, "igamma") def igammac(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) # TODO: delete here after ms.ops.igammac support those types as pytorch input, other, is_float16 = _set_type_gamma(input, other) output = ms.ops.igammac(input, other) if is_float16: output = output.astype(ms.float16) return _out_inplace_assign(out, output, "igammac") def multiply(input, other, *, out=None): if not isinstance(input, (int, adapter_tensor)): raise TypeError(f"multiply(): argument 'input' (position 1) must be Tensor, not {type(input)}") if not isinstance(other, (int, adapter_tensor)): raise TypeError(f"multiply(): argument 'other' (position 2) must be Tensor, not {type(other)}") input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.mul(input, other) return _out_inplace_assign(out, output, "multiply") def mvlgamma(input, p, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.mvlgamma(input, p) return _out_inplace_assign(out, output, "mvlgamma") def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): #TODO: not found mindspore.ops.nan_to_num input_np = input.asnumpy() output_np = np.nan_to_num(input_np, nan=nan, posinf=posinf, neginf=neginf) output = ms.Tensor(output_np) return _out_inplace_assign(out, output, "nan_to_num") def neg(input, *, out=None): input = cast_to_ms_tensor(input) if 'Complex' in str(input.dtype): output = input-input-input else: output = 0 - input return _out_inplace_assign(out, output, "neg") def negative(input, *, out=None): input = cast_to_ms_tensor(input) if 'Complex' in str(input.dtype): output = input-input-input else: output = 0 - input return _out_inplace_assign(out, output, "negative") def nextafter(input, other, *, out=None): # TODO: not found mindspore.ops.nextafter input_np = input.asnumpy() other_np = other.asnumpy() output_np = np.nextafter(input_np, other_np) output = ms.Tensor(output_np) return _out_inplace_assign(out, output, "nextafter") def positive(input): return input def rad2deg(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.rad2deg(input) return _out_inplace_assign(out, output, "rad2deg") def real(input): #TODO: replace tp ms.ops.real() input = cast_to_ms_tensor(input) output = _get_cache_prim(ms.ops.Real)()(input) return cast_to_adapter_tensor(output) def reciprocal(input, *, out=None): #TODO: replace tp ms.ops.reciprocal() input = cast_to_ms_tensor(input) if 'Bool' in str(input.dtype) or 'Int' in str(input.dtype): input = input.astype(ms.float32) # TODO: future to support ms.ops.reciprocal output = ms.Tensor(1.0, dtype=input.dtype)/input return _out_inplace_assign(out, output, "reciprocal") def remainder(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.remainder(input, other) return _out_inplace_assign(out, output, "remainder") def rsqrt(input, *, out=None): input = cast_to_ms_tensor(input) if 'Bool' in str(input.dtype) or 'Int' in str(input.dtype): input = input.astype(ms.float32) output = _get_cache_prim(ms.ops.Rsqrt)()(input) return _out_inplace_assign(out, output, "rsqrt") def sgn(input, *, out=None): input = cast_to_ms_tensor(input) sign_op = _get_cache_prim(ms.ops.Sign)() if 'Bool' in str(input.dtype) or 'Int' in str(input.dtype): type = input.dtype input = input.astype(ms.float32) output = sign_op(input).astype(type) else: output = sign_op(input) return _out_inplace_assign(out, output, "sgn") def signbit(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.numpy.signbit(input) return _out_inplace_assign(out, output, "signbit") def sinc(input, *, out=None): input = cast_to_ms_tensor(input) pi = 3.141592653589793 div = ms.ops.sin(pi*input)/(pi*input) output = ms.numpy.where(input==0, 1.0, div) return _out_inplace_assign(out, output, "sinc") def sinh(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.sinh(input) return _out_inplace_assign(out, output, "sinh") def square(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.square(input) return _out_inplace_assign(out, output, "square") def sub(input, other, *, alpha=1, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.sub(input, other * alpha) return _out_inplace_assign(out, output, "sub") def subtract(input, other, *, alpha=1, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.sub(input, other * alpha) return _out_inplace_assign(out, output, "subtract") def tan(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.tan(input) return _out_inplace_assign(out, output, "tan") def _check_isint(input): if isinstance(input, int): return True if isinstance(input, (adapter_tensor, ms.Tensor)) and 'Int' in str(input.dtype): return True return False def _int_to_float(input): if isinstance(input, int): return float(input) return input.astype(ms.float32) def true_divide(dividend, divisor, *, out=None): input = cast_to_ms_tensor(dividend) other = cast_to_ms_tensor(divisor) is_input_int = _check_isint(input) is_other_int = _check_isint(other) if is_input_int and is_other_int: input = _int_to_float(input) other = _int_to_float(other) if isinstance(input, float) and isinstance(other, float): input = ms.Tensor(input) output = ms.ops.div(input, other, rounding_mode=None) return _out_inplace_assign(out, output, "true_divide") def trunc(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.trunc(input) return _out_inplace_assign(out, output, "trunc") def xlogy(input, other, *, out=None): if not isinstance(input, adapter_tensor) and not isinstance(other, adapter_tensor): raise TypeError("For xlogy: one of the input must be Tensor.") input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.xlogy(input, other) return _out_inplace_assign(out, output, "xlogy") def cov(input, *, correction=1, fweights=None, aweights=None, out=None): #TODO： need be replaced with mindspore.ops.cov input = cast_to_ms_tensor(input).float() if len(input.shape) > 2: raise ValueError("cov(): expected input to have two or fewer dimensions") if isinstance(input.dtype, ms.dtype.Bool): raise TypeError("cov(): bool dtype is not supported for input") if len(input.shape) < 2: input = ms.ops.reshape(input, (1, -1)) observations_dim = 1 num_observations = input.shape[observations_dim] # The product of frequencies (fweights) and weights (aweights). weights = None if fweights is not None: fweights = cast_to_ms_tensor(fweights) if len(fweights) <= 1: raise ValueError("cov(): expected fweights to have one or fewer dimensions") if not isinstance(fweights.dtype, ms.dtype.Int): raise TypeError("cov(): expected fweights to have integral dtype") if ms.ops.size(fweights) != num_observations: raise ValueError("cov(): expected fweights to have the same size as there are observations in the input") if num_observations != 0 and fweights.min() < 0: raise ValueError("cov(): fweights cannot be negative") weights = fweights if aweights is not None: aweights = cast_to_ms_tensor(aweights) if len(aweights) <= 1: raise ValueError("cov(): expected aweights to have one or fewer dimensions") if not isinstance(aweights.dtype, ms.dtype.Float): raise TypeError("cov(): expected aweights to have floating point dtype") if ms.ops.size(aweights) != num_observations: raise ValueError("cov(): expected aweights to have the same size as there are observations in the input") if num_observations != 0 and aweights.min() < 0: raise ValueError("cov(): aweights cannot be negative") if fweights is not None: weights = fweights * aweights else: weights = aweights # Compute a weighted average of the obervations if weights is not None: weights_sum = weights.sum() if weights_sum == 0: raise ValueError("cov(): weights sum to zero, can't be normalized") avg = (input * weights).sum(axis=observations_dim) / weights_sum else: weights_sum = ms.ops.scalar_to_tensor(num_observations, input.dtype) avg = input.sum(axis=observations_dim) / weights_sum # Compute the normalization factor if weights is not None and aweights is not None and correction != 0: norm_factor = weights_sum - correction * (weights * aweights).sum() / weights_sum else: norm_factor = weights_sum - correction # Compute covariance matrix input = input - avg[:, None] if weights is not None: input_T = (input * weights).T else: input_T = input.T # TODO: use code below in future mindspore version, because Ascend not support input as real number for ms.ops.conj # TODO: to support complex number in 'cov' # if ms.ops.is_complex(input_T): # input_T = ms.ops.conj(input_T) output = ms.numpy.dot(input, input_T) output = ms.numpy.true_divide(output, norm_factor) return _out_inplace_assign(out, output.squeeze(), "cov") def corrcoef(input, *, out=None): input = cast_to_ms_tensor(input) if len(input.shape) > 2: raise ValueError("corrcoef(): expected input to have two or fewer dimensions") output = cov(input) if len(output.shape) == 0: return output / output # normalize covariance d = ms.numpy.diag(output) # Clip real and imaginary parts to [-1, 1]. if input.dtype == ms.complex64: real_op = _get_cache_prim(ms.ops.Real)() imag_op = _get_cache_prim(ms.ops.Imag)() complex_op = _get_cache_prim(ms.ops.Complex)() d_real = real_op(d) stddev = ms.ops.sqrt(d_real) output = output / stddev[:, None] output = output / stddev[None, :] output_real = real_op(output) output_imag = imag_op(output) output_real = ms.ops.clip_by_value(output_real, -1, 1) output_imag = ms.ops.clip_by_value(output_imag, -1, 1) output = complex_op(output_real, output_imag) else: stddev = ms.ops.sqrt(d) output = output / stddev[:, None] output = output / stddev[None, :] output = ms.ops.clip_by_value(output, -1, 1) return _out_inplace_assign(out, output, "corrcoef") def cross(input, other, dim=None, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) if is_under_ascend_context() or is_under_gpu_context(): if dim is None: dim = -65530 _op = _get_cache_prim(ms.ops.Cross)(dim=dim) _op.set_device("CPU") output = _op(input, other) else: output = ms.ops.cross(input, other, dim) return _out_inplace_assign(out, output, "cross") def cummax(input, dim, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.cummax(input, axis=dim) return _out_inplace_assign(out, output, "cummax") def cummin(input, dim, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.cummin(input, dim) # the output dtype in ms.ops.cummin is different with ms.ops.cummax output[1] = output[1].astype(ms.common.dtype.int64) return _out_inplace_assign(out, output, "cummin") def cumprod(input, dim, *, dtype=None, out=None): input = cast_to_ms_tensor(input) output = ms.ops.cumprod(input, dim, dtype=dtype) return _out_inplace_assign(out, output, "cumprod") def diagflat(input, offset=0, *, out=None): input = cast_to_ms_tensor(input) output = ms.numpy.diagflat(input, offset) return _out_inplace_assign(out, output, "diagflat") def diagonal(input, offset=0, dim1=0, dim2=1): input = cast_to_ms_tensor(input) #TODO float64 not support if offset != 0 if offset != 0: input = input.astype(mstype.float32) output = ms.ops.diagonal(input, offset, dim1, dim2) return cast_to_adapter_tensor(output) def diff(input, n=1, dim=-1, prepend=None, append=None): input = cast_to_ms_tensor(input) output = ms.numpy.diff(input, n, dim, prepend, append) return cast_to_adapter_tensor(output) def flip(input, dims): input = cast_to_ms_tensor(input) if isinstance(dims, list): dims = tuple(dims) output = ms.ops.flip(input, dims) return cast_to_adapter_tensor(output) def fliplr(input): input = cast_to_ms_tensor(input) output = ms.ops.fliplr(input) return cast_to_adapter_tensor(output) def gather(input, dim, index, *, sparse_grad=False, out=None): if sparse_grad: raise ValueError("`sparse_grad` in `sparse_grad` can not be True.") input = cast_to_ms_tensor(input) index = cast_to_ms_tensor(index) output = ms.ops.gather_elements(input, dim, index) return _out_inplace_assign(out, output, "gather") def bmm(input, mat2, *, out=None) : input_x = cast_to_ms_tensor(input) mat2 = cast_to_ms_tensor(mat2) output = ms.ops.bmm(input_x, mat2) return _out_inplace_assign(out, output, "bmm") def equal(input, other): if not isinstance(input, adapter_tensor) or not isinstance(other, adapter_tensor): raise ValueError("`input` and `other` must be Tensor") x = cast_to_ms_tensor(input) y = cast_to_ms_tensor(other) if x.dtype != y.dtype: return False if x.shape == y.shape: size = x.size output = ms.ops.equal(x, y) output = output.sum() if output == size: return True return False def greater_equal(input, other, *, out=None): x = cast_to_ms_tensor(input) y = cast_to_ms_tensor(other) output = ms.ops.greater_equal(x, y) return _out_inplace_assign(out, output, "greater_equal") def greater(input, other, *, out=None): x = cast_to_ms_tensor(input) y = cast_to_ms_tensor(other) output = ms.ops.greater(x, y) return _out_inplace_assign(out, output, "greater") def less_equal(input, other, *, out=None): x = cast_to_ms_tensor(input) y = cast_to_ms_tensor(other) output = ms.ops.less_equal(x, y) return _out_inplace_assign(out, output, "less_equal") def less(input, other, *, out=None): x = cast_to_ms_tensor(input) y = cast_to_ms_tensor(other) output = ms.ops.less(x, y) return _out_inplace_assign(out, output, "less") def not_equal(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) output = ms.ops.ne(input, other) return _out_inplace_assign(out, output, "not_equal") def baddbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None): x = cast_to_ms_tensor(input) batch1 = cast_to_ms_tensor(batch1) batch2 = cast_to_ms_tensor(batch2) output = ms.ops.baddbmm(x, batch1, batch2, beta, alpha) return _out_inplace_assign(out, output, "baddbmm") def masked_select(input, mask, *, out=None): x = cast_to_ms_tensor(input) mask = cast_to_ms_tensor(mask) output = ms.ops.masked_select(x, mask) return _out_inplace_assign(out, output, "masked_select") def select(input, dim, index): input = cast_to_ms_tensor(input) _input_indices = ms.Tensor(index) output = ms.ops.gather(input, _input_indices, dim) @constexpr def _get_out_shape(input_shape, dim): shape = [input_shape[i] for i in range(len(input_shape)) if i != dim] return tuple(shape) output_shape = _get_out_shape(input.shape, dim) output = output.reshape(output_shape) return cast_to_adapter_tensor(output) def argmin(input, dim=None, keepdim=False): input = cast_to_ms_tensor(input) # TODO: output = ms.ops.argmin(input, axis=dim, keepdims=keepdim) if keepdim: raise NotImplementedError("keepdim is not supported.") output = ms.ops.argmin(input, axis=dim) return cast_to_adapter_tensor(output) def argmax(input, dim=None, keepdim=False): input = cast_to_ms_tensor(input) # TODO: output = ms.ops.argmax(input, axis=dim, keepdims=keepdim) if keepdim: raise NotImplementedError("keepdim is not supported.") output = ms.ops.argmax(input, axis=dim) return cast_to_adapter_tensor(output) def broadcast_to(input, shape): input = cast_to_ms_tensor(input) output = ms.ops.broadcast_to(input, shape) return cast_to_adapter_tensor(output) def ravel(input): x = cast_to_ms_tensor(input) output = ms.ops.reshape(x, (-1,)) return cast_to_adapter_tensor(output) def unique(input, sorted=True, return_inverse=False, return_counts=False, dim=None): unsupported_attr(dim) unsupported_attr(return_counts) input = cast_to_ms_tensor(input) data_type = input.dtype if sorted and return_inverse: raise ValueError("Don't support sorted=True and return_inverse=True.") res, idx = ms.ops.unique(input) if sorted: res = ms.ops.cast(res, ms.float32) res, _ = ms.ops.sort(res) res = ms.ops.cast(res, data_type) if return_inverse: res = cast_to_adapter_tensor(res) idx = cast_to_adapter_tensor(idx) return (res, idx) else: res = cast_to_adapter_tensor(res) return res def permute(input, dims): ms_input = cast_to_ms_tensor(input) output = ms_input.transpose(dims) return cast_to_adapter_tensor(output) def numel(input): input = cast_to_ms_tensor(input) return ms.ops.size(input) def logsumexp(input, dim, keepdim=False, *, out=None): ms_input = cast_to_ms_tensor(input) if ms_input.dtype != mstype.float32: ms_input = ms_input.astype(mstype.float32) output = ms.ops.logsumexp(ms_input, dim, keepdim) return _out_inplace_assign(out, output, "logsumexp") def addmv(input, mat, vec, *, beta=1, alpha=1, out=None): input = cast_to_ms_tensor(input) mat = cast_to_ms_tensor(mat) vec = cast_to_ms_tensor(vec) output = ms.ops.addmv(input, mat, vec, beta=beta, alpha=alpha) return _out_inplace_assign(out, output, "addmv") def dot(input, other, *, out=None): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) #TODO: ms.ops.tensor_dot only supports float16/float32 input_dtype = input.dtype if input_dtype in (mstype.float32, mstype.float16): output = ms.ops.tensor_dot(input, other, 1) else: input = input.astype(ms.float32) other = other.astype(ms.float32) output = ms.ops.tensor_dot(input, other, 1) output = output.astype(input_dtype) return _out_inplace_assign(out, output, "dot") def inverse(input, *, out=None): input = cast_to_ms_tensor(input) #TODO: Ascend has no ms.ops.inverse # output = ms.ops.inverse(input) # ops.inverse is not ready at 11-13 if input.dtype in all_int_type: input = input.astype(mstype.float32) output = _get_cache_prim(ms.ops.MatrixInverse)()(input) return _out_inplace_assign(out, output, "inverse") def count_nonzero(input, dim=None): input = cast_to_ms_tensor(input) if dim is None: dim = () output = ms.ops.count_nonzero(input, axis=dim) return cast_to_adapter_tensor(output) def all(input, dim=(), keepdim=False, *, out=None): input = cast_to_ms_tensor(input) if input.dtype != ms.bool_: input = input.astype(ms.bool_) output = input.all(axis=dim, keep_dims=keepdim) return _out_inplace_assign(out, output, "all") def scatter(input, dim, index, src): return input.scatter(dim, index, src) def topk(input, k, dim=None, largest=True, sorted=True, *, out=None): unsupported_attr(dim) unsupported_attr(largest) input_x = cast_to_ms_tensor(input) output = ms.ops.top_k(input_x, k, sorted=sorted) return _out_inplace_assign(out, output, "topk") def addbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None): _input, _batch1, _batch2 = cast_to_ms_tensor((input, batch1, batch2)) output = ms.ops.addbmm(_input, _batch1, _batch2, beta=beta, alpha=alpha) return _out_inplace_assign(out, output, "addbmm") def addmm(input, mat1, mat2, *, beta=1, alpha=1, out=None): _input, _mat1, _mat2 = cast_to_ms_tensor((input, mat1, mat2)) output = ms.ops.addmm(_input, _mat1, _mat2, beta=beta, alpha=alpha) return _out_inplace_assign(out, output, "addbmm") def addr(input, vec1, vec2, *, beta=1, alpha=1, out=None): _input, _vec1, _vec2 = cast_to_ms_tensor((input, vec1, vec2)) output = ms.ops.addr(_input, _vec1, _vec2, beta=beta, alpha=alpha) return _out_inplace_assign(out, output, "addr") def isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False): _input, _other = cast_to_ms_tensor((input, other)) output = ms.ops.isclose(_input, _other, rtol=rtol, atol=atol, equal_nan=equal_nan) return cast_to_adapter_tensor(output) def allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False): output = all(isclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan)).item() return output def cholesky(self, upper=False, *, out=None): input = cast_to_ms_tensor(self) output = input.cholesky(upper) return _out_inplace_assign(out, output, "cholesky") def cholesky_inverse(self, upper=False, *, out=None): input = cast_to_ms_tensor(self) output = input.cholesky_inverse(upper) return _out_inplace_assign(out, output, "cholesky_inverse") def dist(input, other, p=2): _input = cast_to_ms_tensor(input) _other = cast_to_ms_tensor(other) _input_dtype = _input.dtype if _input_dtype in (ms.float16, ms.float32): _other.astype(_input_dtype) output = ms.ops.dist(_input, _other, p=p) elif _input_dtype == ms.float64: _input = _input.astype(ms.float32) _other = _other.astype(_input.dtype) output = ms.ops.dist(_input, _other, p=p) output = output.astype(ms.float64) else: raise ValueError(f"For torch.dist, input should be floating Tensor, but got {_input_dtype}.") return cast_to_adapter_tensor(output) def aminmax(input, *, dim=None, keepdim=False, out=None): _input = cast_to_ms_tensor(input) _min = _input.min(axis=dim, keepdims=keepdim) _max = _input.max(axis=dim, keepdims=keepdim) return _out_inplace_assign(out, (_min, _max), "aminmax") def any(input, dim=(), keepdim=False, *, out=None): input = cast_to_ms_tensor(input) if input.dtype != ms.bool_: input = input.astype(ms.bool_) output = input.any(axis=dim, keep_dims=keepdim) return _out_inplace_assign(out, output, "any") def is_complex(input): input = cast_to_ms_tensor(input) return input.is_complex() def isinf(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.isinf(input) return _out_inplace_assign(out, output, "isinf") def isneginf(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.isneginf(input) return _out_inplace_assign(out, output, "isneginf") def isposinf(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.isposinf(input) return _out_inplace_assign(out, output, "isposinf") def isreal(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.isreal(input) return _out_inplace_assign(out, output, "isreal") def lgamma(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.lgamma(input) return _out_inplace_assign(out, output, "lgamma") def digamma(input, *, out=None): input = cast_to_ms_tensor(input) output = ms.ops.digamma(input) return _out_inplace_assign(out, output, "digamma") msadapter/ms_adapter/pytorch/nn/000077500000000000000000000000001436713554100172715ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/nn/__init__.py000066400000000000000000000003671436713554100214100ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch.nn.modules import * from ms_adapter.pytorch.nn import init from ms_adapter.pytorch.nn.parameter import Parameter, ParameterTuple from ms_adapter.pytorch.nn.functional import * msadapter/ms_adapter/pytorch/nn/functional.py000066400000000000000000002345511436713554100220170ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- """Functional interface""" import math import warnings import numpy as np import mindspore as ms import mindspore.nn as nn from mindspore.ops import constexpr from mindspore.ops.operations.nn_ops import TripletMarginLoss as TripletMarginLossOp from mindspore.ops._primitive_cache import _get_cache_prim from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.tensor import Tensor, cast_to_ms_tensor, cast_to_adapter_tensor from ms_adapter.pytorch.common._inner import _inplace_assign_pynative all = [ 'smooth_l1_loss', 'log_softmax', 'logsigmoid', 'elu', 'elu_', 'relu', 'relu_', 'upsample', 'rrelu', 'rrelu_', 'selu', 'celu', 'gelu', 'mish', 'softshrink', 'hardtanh', 'hardtanh_', 'hardswish', 'relu6', 'leaky_relu', 'softmax', 'softmin', 'softsign', 'tanh', 'tanhshrink', 'glu', 'softplus', 'sigmoid', 'hardsigmoid', 'silu', 'gumbel_softmax', 'threshold', 'threshold_', 'hardshrink', 'conv1d', 'conv2d', 'conv3d', 'normalize', 'local_response_norm', 'l1_loss', 'cross_entropy', 'ctc_loss', 'gaussian_nll_loss', 'hinge_embedding_loss', 'margin_ranking_loss', 'multilabel_margin_loss', 'multilabel_soft_margin_loss', 'nll_loss', 'kl_div', 'binary_cross_entropy', 'binary_cross_entropy_with_logits', 'upsample_nearest', 'poisson_nll_loss', 'triplet_margin_with_distance_loss', 'pairwise_distance', 'cosine_similarity', 'pdist', 'dropout1d', 'dropout2d', 'dropout3d', 'dropout', 'alpha_dropout', 'feature_alpha_dropout' 'huber_loss', 'soft_margin_loss', 'cosine_embedding_loss', 'pixel_shuffle', 'pixel_unshuffle', 'one_hot', 'embedding', 'max_pool2d', 'fold', 'unfold' ] @constexpr def _get_adaptive_pool_args(input_shape, output_size): _, _, h, w = input_shape if isinstance(output_size, int): output_size = [output_size, ] * 2 condition = [0, ] * 2 out_h = output_size[0] + condition[0] * h out_w = output_size[1] + condition[1] * w stride_h = math.floor(h / out_h) kernel_h = h - (out_h - 1) * stride_h stride_w = math.floor(w / out_w) kernel_w = w - (out_w - 1) * stride_w return kernel_h, kernel_w, stride_h, stride_w def adaptive_avg_pool2d(input, output_size): kernel_h, kernel_w, stride_h, stride_w = _get_adaptive_pool_args(input.shape, output_size) avg_pool = _get_cache_prim(ms.ops.AvgPool)(kernel_size=(kernel_h, kernel_w), strides=(stride_h, stride_w), pad_mode="valid", data_format="NCHW") input = cast_to_ms_tensor(input) out = avg_pool(input) return cast_to_adapter_tensor(out) def adaptive_avg_pool1d(input, output_size): input = cast_to_ms_tensor(input) ndim = input.ndim if ndim == 2: input = input.expand_dims(0) output = ms.ops.adaptive_avg_pool1d(input, output_size) output = output.squeeze(0) else: output = ms.ops.adaptive_avg_pool1d(input, output_size) return cast_to_adapter_tensor(output) # def adaptive_avg_pool2d(input, output_size): # TODO: This ops only supports the GPU # input = cast_to_ms_tensor(input) # output = ms.ops.adaptive_avg_pool2d(input, output_size) # return cast_to_adapter_tensor(output) def adaptive_avg_pool3d(input, output_size): input = cast_to_ms_tensor(input) output = ms.ops.adaptive_avg_pool3d(input, output_size) return cast_to_adapter_tensor(output) def adaptive_max_pool1d(input, output_size, return_indices=False): input = cast_to_ms_tensor(input) ndim = input.ndim if ndim == 2: input = input.expand_dims(0) if return_indices: raise ValueError('keyword argument return_indices is ont supported.') output = ms.ops.adaptive_max_pool1d(input, output_size) output = output.squeeze(0) else: if return_indices: raise ValueError('keyword argument return_indices is ont supported.') output = ms.ops.adaptive_max_pool1d(input, output_size) return cast_to_adapter_tensor(output) def adaptive_max_pool2d(input, output_size, return_indices=False): input = cast_to_ms_tensor(input) output = ms.ops.adaptive_max_pool2d(input, output_size, return_indices) return cast_to_adapter_tensor(output) def adaptive_max_pool3d(input, output_size, return_indices=False): input = cast_to_ms_tensor(input) output = ms.ops.adaptive_max_pool3d(input, output_size, return_indices) return cast_to_adapter_tensor(output) def pad(input, pad, mode="constant", value=0): if mode == "replicate": mode = "edge" value = ms.Tensor(value, dtype=input.dtype) dims = len(input.shape) list_pad = [pad[i:i+2] for i in range(0, len(pad), 2)] list_pad.reverse() new_pad = [[0,0],] * int((dims - len(pad) /2)) new_pad.extend(list_pad) input = cast_to_ms_tensor(input) # TODO: -> ms.ops.PadV3, Padv3 is not supported on Ascend now. # output = ms.ops.operations.nn_ops.PadV3(mode=mode)(input, pad, value) output = ms.numpy.pad(input, new_pad, mode=mode, constant_values=value) return cast_to_adapter_tensor(output) def log_softmax(input, dim=None, _stacklevel=3, dtype=None): unsupported_attr(_stacklevel) # MS dim default is -1 if dim is None: warnings.warn("Implicit dimension choice for log_softmax has been deprecated. " "Change the call to include dim=X as an argument") dim = -1 input = cast_to_ms_tensor(input) if dtype is not None: input = ms.ops.cast(input, dtype) out = ms.ops.log_softmax(input, dim) return cast_to_adapter_tensor(out) def logsigmoid(input): input = cast_to_ms_tensor(input) sigmoid_op = _get_cache_prim(ms.ops.Sigmoid)() sigmoid_out= sigmoid_op(input) ret = ms.ops.log(sigmoid_out) return cast_to_adapter_tensor(ret) def elu(input, alpha=1.0, inplace=False): input_ms = cast_to_ms_tensor(input) if alpha == 1: out = ms.ops.elu(input_ms, alpha) else: cond = ms.ops.gt(input_ms, 0) out = alpha * (ms.ops.exp(input_ms) - 1) out = ms.ops.select(cond, input_ms, out) return _inplace_assign_pynative(input, inplace, out, "elu") def rrelu(input, lower=1.0/8, upper=1.0/3, training=False, inplace=False): if training: raise ValueError("training '{}' is not currently supported.".format(training)) input_ms = cast_to_ms_tensor(input) #TODO： nn.RReLU should be replaced out = nn.RReLU(lower=lower, upper=upper)(input_ms) return _inplace_assign_pynative(input, inplace, out, "rrelu") def selu(input, inplace=False): input_ms = cast_to_ms_tensor(input) out = ms.ops.selu(input_ms) return _inplace_assign_pynative(input, inplace, out, "selu") def celu(input, alpha=1.0, inplace=False): input_ms = cast_to_ms_tensor(input) out = ms.ops.celu(input_ms, alpha) return _inplace_assign_pynative(input, inplace, out, "celu") def gelu(input, approximate='none'): input_x = cast_to_ms_tensor(input) if approximate == 'none': out = input * 0.5 * (1.0 + ms.ops.erf(input/math.sqrt(2.0))) else: out = ms.ops.gelu(input_x, approximate) return cast_to_adapter_tensor(out) def mish(input, inplace=False): input_ms = cast_to_ms_tensor(input) out = ms.ops.mish(input_ms) return _inplace_assign_pynative(input, inplace, out, "mish") def softshrink(input, lambd=0.5): input = cast_to_ms_tensor(input) out = ms.ops.soft_shrink(input, lambd) return cast_to_adapter_tensor(out) def relu(input, inplace=False): input_ms = cast_to_ms_tensor(input) out = ms.ops.relu(input_ms) return _inplace_assign_pynative(input, inplace, out, "relu") def hardtanh(input, min_val=-1.0, max_val=1.0, inplace=False): input_ms = cast_to_ms_tensor(input) out = nn.Hardtanh(min_val, max_val)(input_ms) return _inplace_assign_pynative(input, inplace, out, "hardtanh") def hardswish(input, inplace=False): input_ms = cast_to_ms_tensor(input) out = ms.ops.hardswish(input_ms) return _inplace_assign_pynative(input, inplace, out, "hardswish") def relu6(input, inplace=False): input_ms = cast_to_ms_tensor(input) out = ms.ops.relu6(input_ms) return _inplace_assign_pynative(input, inplace, out, "relu6") def leaky_relu(input, negative_slope=0.01, inplace=False): input_ms = cast_to_ms_tensor(input) #TODO：nn.LeakyReLU should be replaced. out = nn.LeakyReLU(alpha=negative_slope)(input_ms) return _inplace_assign_pynative(input, inplace, out, "leaky_relu") def upsample(input, size=None, scale_factor=None, mode='nearest', align_corners=False): if size is None and scale_factor is None: raise ValueError("either size or scale_factor should be defined") if size is not None and scale_factor is not None: raise ValueError("only one of size or scale_factor should be defined") def linear_func(input): if align_corners is True: trans_mode = 'align_corners' else: trans_mode = 'half_pixel' _size =_upsample_common_process_size(size=size, scale_factor=scale_factor, shape=input.shape) input = cast_to_ms_tensor(input) out = ms.ops.interpolate(input, scales=None, sizes=_size, coordinate_transformation_mode=trans_mode, mode=mode) return cast_to_adapter_tensor(out) def bllinear_func(input): return upsample_bilinear(input, size=size, scale_factor=scale_factor, align_corners=align_corners) def resize_nearest_neighbor_func(input): return upsample_nearest(input, size=size, scale_factor=scale_factor) mode_func = {'linear': linear_func, 'bilinear': bllinear_func, 'nearest': resize_nearest_neighbor_func} if mode not in mode_func: raise ValueError("Until now, `mode` beside 'linear', 'bilinear', 'nearest' are not supported") func = mode_func[mode] out = func(input) return out @constexpr def _get_softmax_dim(ndim): if ndim in (0, 1, 3): ret = 0 else: ret = 1 return ret def softmax(input, dim=None, _stacklevel=3, dtype=None): unsupported_attr(_stacklevel) # MS dim default is -1 if dim is None: dim = -1 input = cast_to_ms_tensor(input) if dtype is not None: input = ms.ops.cast(input, dtype) out = ms.ops.softmax(input, axis=dim) return cast_to_adapter_tensor(out) def softmin(input, dim=None, dtype=None): # MS dim default is -1 # TODO # ms.ops.softmax should be replaced by ms.ops.softmin if dim is None: dim = _get_softmax_dim(input.dim()) input = cast_to_ms_tensor(input) if dtype is not None: input = ms.ops.cast(input, dtype) x = -input out = ms.ops.softmax(x, axis=dim) return cast_to_adapter_tensor(out) def softsign(input): input = cast_to_ms_tensor(input) output = ms.ops.functional.softsign(input) return cast_to_adapter_tensor(output) def tanh(input): input = cast_to_ms_tensor(input) output = ms.ops.functional.tanh(input) return cast_to_adapter_tensor(output) def tanhshrink(input): input = cast_to_ms_tensor(input) ouput = input - ms.ops.functional.tanh(input) return cast_to_adapter_tensor(ouput) def glu(input, dim=-1): if input.dim() == 0: raise RuntimeError("glu does not support scalars because halving size must be even") if input.shape[dim] % 2 == 1: raise RuntimeError("Halving dimension must be even, but dimension {} is size {}".format(dim,input.shape[dim])) halflen = input.shape[dim]//2 data_a = input.narrow(axis=dim, start=0, length=halflen) data_b = input.narrow(axis=dim, start=halflen, length=halflen) data_a = cast_to_ms_tensor(data_a) data_b = cast_to_ms_tensor(data_b) sigmoid_data_b = ms.ops.sigmoid(data_b) out = ms.ops.mul(data_a, sigmoid_data_b) return cast_to_adapter_tensor(out) def normalize(input, p=2.0, dim=1, eps=1e-12, out=None): #the type of 'p' in ms.ops.functional.norm should be 'int' input = cast_to_ms_tensor(input) input_p = ms.ops.pow(abs(input), p) input_p_sum = input_p.sum(axis = dim, keepdims=True) norm = ms.ops.pow(input_p_sum, 1.0/p) min_value = ms.Tensor(eps, ms.float32) denom = ms.ops.clip_by_value(norm, min_value) denom = denom.expand_as(input) output = ms.ops.functional.div(input, denom) if out is not None: ms.ops.assign(out, output) return out return cast_to_adapter_tensor(output) def softplus(input, beta=1, threshold=20): input = cast_to_ms_tensor(input) input_x = beta * input dtype_op = _get_cache_prim(ms.ops.DType)() cast_op = _get_cache_prim(ms.ops.Cast)() alpha_array = cast_op(ms.ops.functional.scalar_to_tensor(threshold), dtype_op(input)) mask = ms.ops.less(alpha_array, input_x) input_mask = ms.ops.masked_fill(input_x, mask, 0) out_mask = ms.ops.exp(input_mask) out_mask_log = ms.ops.log1p(out_mask) ret_mask = out_mask_log/beta ret = ms.ops.select(mask, input, ret_mask) return cast_to_adapter_tensor(ret) def sigmoid(input): input = cast_to_ms_tensor(input) sigmoid_op = _get_cache_prim(ms.ops.Sigmoid)() out = sigmoid_op(input) return cast_to_adapter_tensor(out) def hardsigmoid(input, inplace=False): input_ms = cast_to_ms_tensor(input) hardsigmoid_op = _get_cache_prim(ms.ops.HSigmoid)() out = hardsigmoid_op(input_ms) return _inplace_assign_pynative(input, inplace, out, "hardsigmoid") def silu(input, inplace=False): input_ms = cast_to_ms_tensor(input) sigmoid_op = _get_cache_prim(ms.ops.Sigmoid)() out = sigmoid_op(input_ms) * input_ms return _inplace_assign_pynative(input, inplace, out, "silu") def gumbel_softmax(logits, tau=1.0, hard=False, eps=1e-10, dim=-1): if eps != 1e-10: warnings.warn("`eps` parameter is deprecated and has no effect.") logits = cast_to_ms_tensor(logits) out = ms.ops.gumbel_softmax(logits, tau, hard, dim) return cast_to_adapter_tensor(out) def threshold(input, threshold, value, inplace=False): input_ms = cast_to_ms_tensor(input) cond = ms.ops.gt(input_ms, threshold) value = ms.ops.fill(input_ms.dtype, input_ms.shape, value) out = ms.ops.select(cond, input_ms, value) return _inplace_assign_pynative(input, inplace, out, "threshold") rrelu_ = rrelu relu_ = relu elu_ = elu hardtanh_ = hardtanh leaky_relu_ = leaky_relu threshold_ = threshold @constexpr def _get_reduce_string(size_average, reduce): if size_average is None: size_average = True if reduce is None: reduce = True if size_average and reduce: ret = 'mean' elif reduce: ret = 'sum' else: ret = 'none' warning = "size_average and reduce args will be deprecated, please use reduction='{}' instead." warnings.warn(warning.format(ret)) return ret def smooth_l1_loss(input, target, size_average=None, reduce=None, reduction='mean', beta=1.0): if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) output = ms.ops.smooth_l1_loss(input, target, beta, reduction) if reduction != 'none' and output.shape == 1: return Tensor(output[0]) return cast_to_adapter_tensor(output) def _get_loss(x, reduction): """ Computes the loss. """ if reduction is None or reduction == 'none': return x def _get_axis(x): shape = ms.ops.shape(x) length = ms.ops.tuple_len(shape) perm = ms.ops.make_range(0, length) return perm input_dtype = x.dtype x = ms.ops.cast(x, ms.float32) if reduction == 'mean': reduce_mean_op = _get_cache_prim(ms.ops.ReduceMean)() x = reduce_mean_op(x, _get_axis(x)) if reduction == 'sum': x = ms.ops.reduce_sum(x, _get_axis(x)) x = ms.ops.cast(x, input_dtype) return x def l1_loss(input, target, size_average=None, reduce=None, reduction="mean"): """ Function that takes the mean element-wise absolute value difference. """ if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) def _l1_loss_function(input, target, reduction): x = ms.ops.abs(input - target) return _get_loss(x, reduction) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) # TODO: Replace with ms.ops.l1_loss result = _l1_loss_function(input, target, reduction) return cast_to_adapter_tensor(result) def mse_loss(input, target, size_average=None, reduce=None, reduction="mean"): """ Measures the element-wise mean squared error. """ if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) x = ms.ops.square(input - target) result = _get_loss(x, reduction) return cast_to_adapter_tensor(result) def cross_entropy(input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean", label_smoothing=0.0): """ This criterion computes the cross entropy loss between input logits and target. """ if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) weight = cast_to_ms_tensor(weight) result = ms.ops.cross_entropy(input, target, weight, ignore_index, reduction, label_smoothing) return cast_to_adapter_tensor(result) def ctc_loss(log_probs, targets, input_lengths, target_lengths, blank=0, reduction='mean', zero_infinity=False): log_probs = cast_to_ms_tensor(log_probs) targets = cast_to_ms_tensor(targets) #TODO: length do not support tuple if not isinstance(input_lengths, Tensor) or not isinstance(target_lengths, Tensor): raise TypeError("'input_lengths' and 'target_lengths' only support Tensor now") if isinstance(input_lengths, Tensor) and isinstance(target_lengths, Tensor): input_lengths = cast_to_ms_tensor(input_lengths) target_lengths = cast_to_ms_tensor(target_lengths) if targets.dtype not in {ms.int32, ms.int64} \ or not (targets.dtype == input_lengths.dtype and targets.dtype == target_lengths.dtype): targets = targets.astype(ms.int64) input_lengths = input_lengths.astype(ms.int64) target_lengths = target_lengths.astype(ms.int64) result, _ = ms.ops.ctc_loss(log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity) return cast_to_adapter_tensor(result) def gaussian_nll_loss(input, target, var, full=False, eps=1e-06, reduction='mean'): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) var = cast_to_ms_tensor(var) rlt = ms.ops.gaussian_nll_loss(input, target, var, full, eps, reduction) return cast_to_adapter_tensor(rlt) def hinge_embedding_loss(input, target, margin=1.0, size_average=None, reduce=None, reduction='mean'): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) rlt = ms.ops.hinge_embedding_loss(input, target, float(margin), reduction) return cast_to_adapter_tensor(rlt) def margin_ranking_loss(input1, input2, target, margin=0, size_average=None, reduce=None, reduction='mean'): input1 = cast_to_ms_tensor(input1) input2 = cast_to_ms_tensor(input2) target = cast_to_ms_tensor(target) if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) rlt = ms.ops.margin_ranking_loss(input1, input2, target, float(margin), reduction) return cast_to_adapter_tensor(rlt) def multilabel_margin_loss(input, target, size_average=None, reduce=None, reduction='mean'): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) if target.dtype != ms.int32: target = target.astype(ms.int32) if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) rlt = ms.ops.MultilabelMarginLoss(reduction)(input, target) return cast_to_adapter_tensor(rlt) def multilabel_soft_margin_loss(input, target, weight=None, size_average=None, reduce=None, reduction='mean'): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) if isinstance(weight, Tensor): weight = cast_to_ms_tensor(weight) if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) rlt = ms.ops.multilabel_soft_margin_loss(input, target, weight, reduction) return cast_to_adapter_tensor(rlt) def nll_loss(input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean"): """ The negative log likelihood loss. """ if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) weight = cast_to_ms_tensor(weight) result = ms.ops.nll_loss(input, target, weight, ignore_index, reduction, label_smoothing=0.0) return cast_to_adapter_tensor(result) def kl_div(input, target, size_average=None, reduce=None, reduction="mean", log_target=False): """ The `Kullback-Leibler divergence Loss. ` """ if size_average is not None or reduce is not None: reduction = _get_reduce_string(size_average, reduce) # TODO if log_target is True: raise ValueError('`log_target` in `{}` can not support True'.format(kl_div.__name__)) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) result = ms.ops.kl_div(input, target, reduction) return cast_to_adapter_tensor(result) def binary_cross_entropy(input, target, weight=None, size_average=None, reduce=None, reduction="mean"): """ Function that measures the Binary Cross Entropy between the target and input probabilities. """ if size_average is not None or reduce is not None: reduction = _get_reduce_string(size_average, reduce) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) weight = cast_to_ms_tensor(weight) result = ms.ops.binary_cross_entropy(input, target, weight, reduction) return cast_to_adapter_tensor(result) def binary_cross_entropy_with_logits(input, target, weight=None, size_average=None, reduce=None, reduction="mean", pos_weight=None): """ Function that measures Binary Cross Entropy between target and input logits. """ if size_average is not None or reduce is not None: reduction = _get_reduce_string(size_average, reduce) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) weight = cast_to_ms_tensor(weight) pos_weight = cast_to_ms_tensor(pos_weight) if weight is None or pos_weight is None: ones_input = ms.ops.ones_like(input) if weight is None: weight = ones_input if pos_weight is None: pos_weight = ones_input result = ms.ops.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction) return cast_to_adapter_tensor(result) def _upsample_common_check(size, scale_factor): if size is None and scale_factor is None: raise ValueError("either size or scale_factor should be defined.") if size is not None and scale_factor is not None: raise ValueError("only one of size or scale_factor should be defined.") def _upsample_common_process_size(size, scale_factor, shape): input_shape = list(shape) input_rank = len(shape) if scale_factor is not None: size_ = input_shape[2:] for i, _ in enumerate(size_): size_[i] *= scale_factor size_[i] = int(size_[i]) else: if not isinstance(size, (int, list, tuple)): raise TypeError("`size` should be in types of int, list and tuple.") if isinstance(size, int): size_ = [size for i in range(2, input_rank)] else: if len(size) != input_rank - 2: raise ValueError( "Input and output must have the same number of spatial dimensions, but got " f"input with spatial dimensions of {list(input_shape[2:])} and output size of {size}. " "Please provide input tensor in (N, C, d1, d2, ...,dK) format and " "output size in (o1, o2, ...,oK) format.") size_ = size return tuple(size_) def upsample_nearest(input, size=None, scale_factor=None): _upsample_common_check(size, scale_factor) input_shape = input.shape _input_rank = len(input_shape) size_ = _upsample_common_process_size(size, scale_factor, input_shape) if _input_rank == 4: _op = _get_cache_prim(ms.ops.ResizeNearestNeighbor)(size_, align_corners=False) elif _input_rank == 5: _op = _get_cache_prim(ms.ops.UpsampleNearest3D)(size_) else: raise ValueError(f"upsample_nearest only support 4D or 5D input, but got {_input_rank}D.") input = cast_to_ms_tensor(input) result = _op(input) return cast_to_adapter_tensor(result) def upsample_bilinear(input, size=None, scale_factor=None, *, align_corners=True): input_shape = input.shape if len(input_shape) != 4: raise ValueError("Until now, upsample_bilinear only support 4-D input.") _upsample_common_check(size, scale_factor) size_ = _upsample_common_process_size(size, scale_factor, input_shape) input = cast_to_ms_tensor(input) if align_corners is True: _cor_mode = "align_corners" else: _cor_mode = "half_pixel" result = ms.ops.interpolate(input, sizes=size_, coordinate_transformation_mode=_cor_mode, mode="bilinear") return cast_to_adapter_tensor(result) def pairwise_distance(x1, x2, p=2.0, eps=1e-06, keepdim=False): x1 = cast_to_ms_tensor(x1) x2 = cast_to_ms_tensor(x2) input = x1-x2+eps input_p = ms.ops.pow(ms.ops.abs(input), p) input_p_sum = input_p.sum(axis=-1, keepdims=keepdim) out = ms.ops.pow(input_p_sum, 1.0 / p) return cast_to_adapter_tensor(out) def cosine_similarity(x1, x2, dim=1, eps=1e-08): x1 = cast_to_ms_tensor(x1) x2 = cast_to_ms_tensor(x2) while x1.ndim < x2.ndim: x1 = x1.expand_dims(0) while x2.ndim < x1.ndim: x2 = x2.expand_dims(0) if x1.size < x2.size: x1 = ms.ops.broadcast_to(x1, x2.shape) if x2.size < x1.size: x2 = ms.ops.broadcast_to(x2, x1.shape) min_value = ms.Tensor(eps, ms.float32) x1_norm = ms.ops.pow(x1, 2) x1_norm = x1_norm.sum(axis=dim) x1_norm = ms.ops.pow(x1_norm, 1.0/2) x1_norm = ms.ops.clip_by_value(x1_norm, min_value) x2_norm = ms.ops.pow(x2, 2) x2_norm = x2_norm.sum(axis=dim) x2_norm = ms.ops.pow(x2_norm, 1.0/2) x2_norm = ms.ops.clip_by_value(x2_norm, min_value) denom = ms.ops.mul(x1_norm, x2_norm) out = ms.ops.mul(x1, x2).sum(axis=dim)/denom return cast_to_adapter_tensor(out) def pdist(input, p=2): inp_dim = input.dim() if inp_dim != 2: raise RuntimeError(f"pdist only supports 2D tensors, got: {inp_dim}D") if p < 0: raise RuntimeError("pdist only supports non-negative p values") input = cast_to_ms_tensor(input) n, m = input.shape x = input.broadcast_to((n, n, m)).astype(ms.float32) y = x.transpose(1, 0, 2) norm = ms.ops.pow(ms.ops.abs(x-y), p) norm = norm.sum(axis=-1) if p > 0: norm = ms.ops.pow(norm, 1.0/p) select = np.ones([n, n]) select = np.triu(select, 1).astype(np.bool8) select_t = ms.Tensor(select) out = ms.ops.masked_select(norm, select_t) return cast_to_adapter_tensor(out) def dropout1d(input, p = 0.5, training = True, inplace = False): if p < 0.0 or p > 1.0: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) inp_dim = input.dim() if inp_dim not in (2, 3): raise RuntimeError(f"dropout1d: Expected 2D or 3D input, but received a {inp_dim}D input. " "Note that dropout1d exists to provide channel-wise dropout on inputs with 1 " "spatial dimension, a channel dimension, and an optional batch dimension " "(i.e. 2D or 3D inputs).") # is_batched = inp_dim == 3 if not training: return input input_ms = cast_to_ms_tensor(input) out = ms.ops.dropout1d(input_ms, p) return _inplace_assign_pynative(input, inplace, out, "dropout1d") def dropout2d(input, p=0.5, training=True, inplace=False): if p < 0.0 or p > 1.0: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) inp_dim = input.dim() if inp_dim not in (3, 4): warn_msg = (f"dropout2d: Received a {inp_dim}-D input to dropout2d, which is deprecated " "and will result in an error in a future release. To retain the behavior " "and silence this warning, please use dropout instead. Note that dropout2d " "exists to provide channel-wise dropout on inputs with 2 spatial dimensions, " "a channel dimension, and an optional batch dimension (i.e. 3D or 4D inputs).") warnings.warn(warn_msg) if not training: return input if inp_dim == 3: warnings.warn("dropout2d: Received a 3D input to dropout2d and assuming that channel-wise " "1D dropout behavior is desired - input is interpreted as shape (N, C, L), where C " "is the channel dim. This behavior will change in a future release to interpret the " "input as one without a batch dimension, i.e. shape (C, H, W). To maintain the 1D " "channel-wise dropout behavior, please switch to using dropout1d instead.") return dropout1d(input, p, training, inplace) input_ms = cast_to_ms_tensor(input) out, _ = ms.ops.dropout2d(input_ms, p) return _inplace_assign_pynative(input, inplace, out, "dropout2d") def dropout3d(input, p=0.5, training=True, inplace=False): if p < 0.0 or p > 1.0: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) inp_dim = input.dim() if inp_dim not in (4, 5): warn_msg = (f"dropout3d: Received a {inp_dim}-D input to dropout3d, which is deprecated " "and will result in an error in a future release. To retain the behavior " "and silence this warning, please use dropout instead. Note that dropout3d " "exists to provide channel-wise dropout on inputs with 3 spatial dimensions, " "a channel dimension, and an optional batch dimension (i.e. 4D or 5D inputs).") warnings.warn(warn_msg) if not training: return input is_batched = inp_dim == 5 input_ms = cast_to_ms_tensor(input) if not is_batched: input_ms = ms.ops.expand_dims(input_ms, 0) out, _ = ms.ops.dropout3d(input_ms, p) if not is_batched: out = ms.ops.squeeze(out, 0) return _inplace_assign_pynative(input, inplace, out, "dropout3d") def dropout(input, p=0.5, training=True, inplace=False): if p < 0.0 or p > 1.0: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) if not training: return input input_ms = cast_to_ms_tensor(input) shape = input_ms.shape random_array_np = np.random.rand(input_ms.size).reshape(shape) random_array = ms.Tensor(random_array_np, ms.float32) mask = (random_array > ms.Tensor(p, ms.float32)) out = mask * 1.0 / (1.0-p) * input_ms return _inplace_assign_pynative(input, inplace, out, "dropout") def alpha_dropout(input, p=0.5, training=False, inplace=False): if p < 0.0 or p > 1.0: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) if not training: return input input_x = cast_to_ms_tensor(input) # mean = input.mean() # var = input.var() mean = 0.0 var = 1.0 scale = 1.0507009873554804934193349852946 alpha = 1.6732632423543772848170429916717 alpha_ = -scale * alpha q = 1.0 - p a = math.sqrt(var/(q*var + q*(1.0-q)*(alpha_-mean)*(alpha_-mean))) b = mean - a*(q*mean + (1.0-q)*alpha_) shape = input_x.shape random_array_np = np.random.rand(input_x.size).reshape(shape) random_array = ms.Tensor(random_array_np, ms.float32) mask = (random_array > ms.Tensor(p, ms.float32)) value = ms.ops.fill(input_x.dtype, shape, alpha_) out = input_x * mask out = ms.ops.select(mask, out, value) out = out * a + b return _inplace_assign_pynative(input, inplace, out, "alpha_dropout") def feature_alpha_dropout(input, p=0.5, training=False, inplace=False): if p < 0.0 or p > 1.0: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) if not training: return input input_x = cast_to_ms_tensor(input) # mean = input.mean() # var = input.var() mean = 0.0 var = 1.0 scale = 1.0507009873554804934193349852946 alpha = 1.6732632423543772848170429916717 alpha_ = -scale * alpha q = 1.0 - p a = math.sqrt(var/(q*var + q*(1.0-q)*(alpha_-mean)*(alpha_-mean))) b = mean - a*(q*mean + (1.0-q)*alpha_) shape = input_x.shape random_array_np = np.random.rand(shape[0], shape[1]) random_array = ms.Tensor(random_array_np, ms.float32) if input_x.dim() > 2: random_array = random_array.expand_dims(2) random_array = random_array.expand_as(input_x.reshape(shape[0], shape[1], -1)).reshape(shape) mask = (random_array > ms.Tensor(p, ms.float32)) value = ms.ops.fill(input_x.dtype, input_x.shape, alpha_) out = input_x * mask out = ms.ops.select(mask, out, value) out = out * a + b return _inplace_assign_pynative(input, inplace, out, "feature_alpha_dropout") def hardshrink(input, lambd=0.5): input = cast_to_ms_tensor(input) out = ms.ops.hardshrink(input, lambd) return cast_to_adapter_tensor(out) def huber_loss(input, target, reduction='mean', delta=1.0): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) delta_half = 0.5 * delta z = ms.ops.abs(ms.ops.sub(input, target)) condition = ms.ops.less(z, delta) l1 = ms.ops.mul(0.5, ms.ops.square(z)) l2 = ms.ops.mul(delta, ms.ops.sub(z, delta_half)) loss = ms.ops.select(condition, l1, l2) loss = _get_loss(loss, reduction) return cast_to_adapter_tensor(loss) def soft_margin_loss(input, target, size_average=None, reduce=None, reduction='mean'): if size_average is not None or reduce is not None: reduction = _get_reduce_string(size_average, reduce) input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) ops = ms.ops.SoftMarginLoss(reduction) loss = ops(input, target) return cast_to_adapter_tensor(loss) def cosine_embedding_loss( input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean", ): if margin < -1.0 or margin > 1.0: raise ValueError(f"'cosine_embedding_loss': `margin` should be from -1 to 1, but got {margin}") if size_average is not None or reduce is not None: reduction = _get_reduce_string(size_average, reduce) input1 = cast_to_ms_tensor(input1) input2 = cast_to_ms_tensor(input2) target = cast_to_ms_tensor(target) ndim = input1.ndim if ndim == 1: input1 = input1.expand_dims(0) input2 = input2.expand_dims(0) target = target.expand_dims(0) reduce_sum = _get_cache_prim(ms.ops.ReduceSum)() maximum = _get_cache_prim(ms.ops.Maximum)() prod_sum = reduce_sum(input1 * input2, (1,)) square1 = reduce_sum(ms.ops.square(input1), (1,)) square2 = reduce_sum(ms.ops.square(input2), (1,)) denom = ms.ops.sqrt(square1) * ms.ops.sqrt(square2) cosine = prod_sum / denom pos_value = 1.0 - cosine neg_value = maximum(cosine - margin, 0.0) zeros = ms.ops.zeros_like(cosine) pos_part = ms.ops.select(target == 1, pos_value, zeros) neg_part = ms.ops.select(target == -1, neg_value, zeros) output_unreduced = pos_part + neg_part loss = _get_loss(output_unreduced, reduction) if ndim == 1 and reduction == 'none': loss = loss.squeeze(0) return cast_to_adapter_tensor(loss) def triplet_margin_loss( anchor, positive, negative, margin=1.0, p=2, eps=1e-6, swap=False, size_average=None, reduce=None, reduction="mean", ): if size_average is not None or reduce is not None: reduction = _get_reduce_string(size_average, reduce) anchor, positive, negative = cast_to_ms_tensor((anchor, positive, negative)) ndim = anchor.ndim if ndim == 1: anchor = anchor.expand_dims(0) positive = positive.expand_dims(0) negative = negative.expand_dims(0) margin = ms.ops.scalar_to_tensor(margin) # TODO: 'TripletMarginLossOp' is a inner interface, should be change to public api in the future triplet_margin_loss = _get_cache_prim(TripletMarginLossOp)(p=p, swap=swap, eps=eps, reduction=reduction) loss = triplet_margin_loss(anchor, positive, negative, margin) if ndim == 1 and reduction == 'none': loss = loss.squeeze(0) return cast_to_adapter_tensor(loss) def multi_margin_loss( input, target, p=1, margin=1, weight=None, size_average=None, reduce=None, reduction="mean", ): if size_average is not None or reduce is not None: reduction = _get_reduce_string(size_average, reduce) if p not in (1, 2): raise ValueError("only p == 1 and p == 2 supported") input, target = cast_to_ms_tensor((input, target)) if weight is not None: if weight.dim() != 1: raise ValueError("weight must be one-dimensional") weight = cast_to_ms_tensor(weight) loss = ms.ops.multi_margin_loss(input, target, p=p, margin=margin, weight=weight, reduction=reduction) return cast_to_adapter_tensor(loss) loss = ms.ops.multi_margin_loss(input, target, p=p, margin=margin, weight=weight, reduction=reduction) return cast_to_adapter_tensor(loss) def avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None): unsupported_attr(ceil_mode) unsupported_attr(count_include_pad) unsupported_attr(divisor_override) if stride is None: stride = kernel_size padding = padding if isinstance(padding, tuple) else (padding, padding) pad_ops = ms.ops.Pad(((0, 0), (0, 0), (padding[0], padding[0]), (padding[1], padding[1]))) avg_pool_ops = ms.ops.AvgPool(kernel_size=kernel_size, strides=stride, pad_mode='valid') input = cast_to_ms_tensor(input) ndim = input.ndim if ndim == 3: input = input.expand_dims(0) input = pad_ops(input) out = avg_pool_ops(input) out = out.squeeze(0) else: input = pad_ops(input) out = avg_pool_ops(input) return cast_to_adapter_tensor(out) def local_response_norm(input, size, alpha=0.0001, beta=0.75, k=1.0): dim = input.dim() if dim < 3: raise ValueError( "Expected 3D or higher dimensionality \ input (got {} dimensions)".format( dim ) ) if input.size() == 0: return input input = cast_to_ms_tensor(input) div = ms.ops.mul(input, input).expand_dims(axis=1) if dim == 3: div = ms.ops.pad(div, ((0, 0), (0, 0), (size//2, (size-1)//2), (0, 0))) div = ms.ops.avg_pool2d(div, (size, 1), stride=1).squeeze(1) else: shape = input.shape div = div.view(shape[0], 1, shape[1], shape[2], -1) div = ms.ops.pad(div, ((0, 0), (0, 0), (size//2, (size-1)//2), (0, 0), (0, 0))) div = _get_cache_prim(ms.ops.AvgPool3D)((size, 1, 1), strides=1)(div).squeeze(1) div = div.view(shape) div = div * alpha + k div = ms.ops.pow(div, beta) output = input / div return cast_to_adapter_tensor(output) def one_hot(input, num_classes=-1): if num_classes == -1: depth = int(input.asnumpy().max()) + 1 else: depth = num_classes input = cast_to_ms_tensor(input) on_value = ms.Tensor(1.0, ms.float32) off_value = ms.Tensor(0.0, ms.float32) out = ms.ops.one_hot(input, depth, on_value, off_value).astype(ms.int64) return cast_to_adapter_tensor(out) def pixel_shuffle(input, upscale_factor): dim = input.dim() if dim < 3: raise RuntimeError("pixel_shuffle expects input to have at least 3 dimensions, " "but got input with {} dimension(s)".format(dim)) input = cast_to_ms_tensor(input) if dim == 3: input = input.expand_dims(0) shape_in = list(input.shape) tmp = input.reshape(-1, shape_in[-3], shape_in[-2], shape_in[-1]) c = int(tmp.shape[-3] / upscale_factor / upscale_factor) if c * upscale_factor * upscale_factor != tmp.shape[-3]: raise RuntimeError( "pixel_shuffle expects its input's 'channel' dimension to be divisible by the square of upscale_factor," "but input.size(-3)={} is not divisible by {}".format(tmp.shape[-3], upscale_factor*upscale_factor)) h = tmp.shape[-2] w = tmp.shape[-1] tmp = tmp.reshape(-1, c, upscale_factor, upscale_factor, h, w).transpose(0, 1, 4, 2, 5, 3) out = tmp.reshape(-1, c, h * upscale_factor, w * upscale_factor) shape_in[-3] = c shape_in[-2] = h * upscale_factor shape_in[-1] = w * upscale_factor out = out.reshape(shape_in) if dim == 3: out = out.squeeze(0) return cast_to_adapter_tensor(out) def pixel_unshuffle(input, downscale_factor): dim = input.dim() if dim < 3: raise RuntimeError("pixel_shuffle expects input to have at least 3 dimensions, " "but got input with {} dimension(s)".format(dim)) input = cast_to_ms_tensor(input) if dim == 3: input = input.expand_dims(0) shape_in = list(input.shape) tmp = input.reshape(-1, shape_in[-3], shape_in[-2], shape_in[-1]) c = tmp.shape[-3] h = int(tmp.shape[-2] / downscale_factor) w = int(tmp.shape[-1] / downscale_factor) if h * downscale_factor != tmp.shape[-2]: raise RuntimeError( "pixel_unshuffle expects height to be divisible by downscale_factor, " "but input.size(-2)={} is not divisible by {}".format(tmp.shape[-2], downscale_factor)) if w * downscale_factor != tmp.shape[-1]: raise RuntimeError( "pixel_unshuffle expects width to be divisible by downscale_factor, " "but input.size(-1)={} is not divisible by {}".format(tmp.shape[-1], downscale_factor)) tmp = tmp.reshape(-1, c, h, downscale_factor, w, downscale_factor).transpose(0, 1, 3, 5, 2, 4) out = tmp.reshape(-1, c * downscale_factor * downscale_factor, h, w) shape_in[-3] = c * downscale_factor * downscale_factor shape_in[-2] = h shape_in[-1] = w out = out.reshape(shape_in) if dim == 3: out = out.squeeze(0) return cast_to_adapter_tensor(out) def interpolate(input, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None, antialias=False): unsupported_attr(recompute_scale_factor) unsupported_attr(antialias) if mode in ("nearest", "area", "nearest-exact"): if align_corners is not None: raise ValueError( "align_corners option can only be set with the " "interpolating modes: linear | bilinear | bicubic | trilinear" ) align_corners = False else: if align_corners is None: align_corners = False if recompute_scale_factor is not None and recompute_scale_factor: # TODO: not support these two arguments until now pass if antialias: raise NotImplementedError("antialias in interpolate is not supported to True.") # TODO:　not support `antialias` until now. if antialias and not (mode in ("bilinear", "bicubic") and input.ndim == 4): raise ValueError("Anti-alias option is only supported for bilinear and bicubic modes") # TODO: 'nearest' only support 4D, 5D input. 3D is not support until now. if mode == 'nearest': if input.dim() not in (4, 5): raise NotImplementedError(f"For now, 'nearest' only 4D, 5D input are supported, but got {input.dim()}D") return upsample_nearest(input, size, scale_factor) # TODO: 'bilinear' only support 4D input. 3D, 5D are not support until now. if mode == 'bilinear': if input.dim() != 4: raise NotImplementedError(f"For now, 'bilinear' only 4D input is supported, but got {input.dim()}D") return upsample_bilinear(input, size, scale_factor, align_corners=align_corners) if mode == 'linear': if input.dim() != 3: raise ValueError(f"'linear' mode only support 3D input, but got {input.dim()}D") if align_corners is True: trans_mode = 'align_corners' else: trans_mode = 'half_pixel' _size =_upsample_common_process_size(size=size, scale_factor=scale_factor, shape=input.shape) input = cast_to_ms_tensor(input) out = ms.ops.interpolate(input, scales=None, sizes=_size, coordinate_transformation_mode=trans_mode, mode=mode) return cast_to_adapter_tensor(out) if mode in ['bicubic', 'trilinear', 'area', 'nearest-exact']: raise NotImplementedError(f"For interpolate: currently not support mode '{mode}'") raise NotImplementedError( "Input Error: Only 3D, 4D and 5D input Tensors supported" " (got {}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact" " (got {})".format(input.dim(), mode) ) def embedding( input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False ): unsupported_attr(scale_grad_by_freq) unsupported_attr(sparse) if padding_idx: raise NotImplementedError("nn.Embedding: `padding_idx` is not supported until now.") input = cast_to_ms_tensor(input) if padding_idx is not None: if padding_idx > 0: if padding_idx >= weight.shape[0]: raise ValueError("Padding_idx must be within num_embeddings") elif padding_idx < 0: if padding_idx < -weight.shape[0]: raise ValueError("Padding_idx must be within num_embeddings") padding_idx = weight.shape[0] + padding_idx # TODO: norm_type only support '2', others are not supported yet if norm_type != 2: raise NotImplementedError("`norm_type` beside 2 is not supported until now.") # TODO: Try to let 'weight[padding_idx]' not updating by gradient, but pynative didn't work. # Actually, when use "weight[padding_idx] = ...", it will create ops 'TensorScatterUpdate' # And 'TensorScatterUpdate''s backprop can meet that it would not pass gradient to weight[padding_idx]. # However, when directly use 'TensorScatterUpdate', ops will be eliminated in graph optimization. # So, that is the problem to solve, which means the 'padding_idx' will be supported in the future. if max_norm: weight = _get_cache_prim(ms.nn.ClipByNorm)(axis=1)(weight, clip_norm=ms.ops.scalar_to_tensor(max_norm)) out = ms.ops.gather(weight, input, axis=0) return cast_to_adapter_tensor(out) def grid_sample(input, grid, mode='bilinear', padding_mode='zeros', align_corners=None): input = cast_to_ms_tensor(input) grid = cast_to_ms_tensor(grid) if align_corners is None: align_corners = False output = ms.ops.grid_sample(input, grid, interpolation_mode=mode, padding_mode=padding_mode, align_corners=align_corners) output = cast_to_adapter_tensor(output) return output def conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): # TODO: not support float64, change to float32 now input_ms = cast_to_ms_tensor(input) weight_ms = cast_to_ms_tensor(weight) is_float64 = False if input_ms.dtype in (ms.float64, ms.double): input_ms = input_ms.astype(ms.float32) weight_ms = weight_ms.astype(ms.float32) is_float64 = True if isinstance(stride, tuple): stride = stride[0] pad_mode = "pad" if isinstance(padding, int): padding = (0, 0, padding, padding) elif isinstance(padding, tuple): padding = (0, 0, padding[0], padding[0]) else: pad_mode = padding padding = 0 if isinstance(dilation, tuple): dilation = dilation[0] input_shape = input_ms.shape if len(input_shape) != 3: raise ValueError(f"For 'conv1d', the dimension of input must be 3d, but got {len(input_shape)}.") input_ms = ms.ops.expand_dims(input_ms, 2) weight_ms = ms.ops.expand_dims(weight_ms, 2) output = ms.ops.conv2d(input_ms, weight_ms, pad_mode, padding, stride, dilation, groups) if bias is not None: # TODO: ms.ops.biasadd also not support float64 if bias.dtype != output.dtype: bias = bias.astype(output.dtype) output = ms.ops.bias_add(output, bias) output = ms.ops.squeeze(output, 2) if is_float64: output = output.astype(ms.float64) return cast_to_adapter_tensor(output) def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): # Todo: not support float64, change to float32 now input_ms = cast_to_ms_tensor(input) weight_ms = cast_to_ms_tensor(weight) is_float64 = False if input_ms.dtype in (ms.float64, ms.double): input_ms = input_ms.astype(ms.float32) weight_ms = weight_ms.astype(ms.float32) is_float64 = True if isinstance(stride, int): stride = (stride, stride) elif len(stride)==1: stride = (stride[0], stride[0]) pad_mode = "pad" if isinstance(padding, int): padding = (padding, padding, padding, padding) elif isinstance(padding, tuple): if len(padding)==1: padding = (padding[0], padding[0], padding[0], padding[0]) else: padding = (padding[0], padding[0], padding[1], padding[1]) else: pad_mode = padding padding = 0 if isinstance(dilation, int): dilation = (dilation, dilation) elif len(dilation) == 1: dilation = (dilation[0], dilation[0]) output = ms.ops.conv2d(input_ms, weight_ms, pad_mode, padding, stride, dilation, groups) if bias is not None: # TODO: ms.ops.biasadd also not support float64 if bias.dtype != output.dtype: bias = bias.astype(output.dtype) output = ms.ops.bias_add(output, bias) if is_float64: output = output.astype(ms.float64) return cast_to_adapter_tensor(output) def max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False): input = cast_to_ms_tensor(input) if input.dtype in (ms.int16, ms.int32, ms.int64): raise TypeError("'max_pool2d' not implemented for int.") if return_indices is True or dilation != 1 or ceil_mode is True: _kernel_size = kernel_size + (1,) if isinstance(kernel_size, tuple) else (kernel_size, kernel_size, 1) if stride is None: _stride = _kernel_size else: _stride = stride + (1,) if isinstance(stride, tuple) else (stride, stride, 1) _padding = padding + (0,) if isinstance(padding, tuple) else (padding, padding, 0) _dilation = dilation + (1,) if isinstance(dilation, tuple) else (dilation, dilation, 1) input = cast_to_ms_tensor(input) ndim = input.ndim if ndim == 3: input = input.expand_dims(0) input = input.expand_dims(-1) out, indices = ms.ops.max_pool3d(input, _kernel_size, _stride, _padding, _dilation, ceil_mode, True) out = out.squeeze(0) indices = indices.squeeze(0) else: input = input.expand_dims(-1) out, indices = ms.ops.max_pool3d(input, _kernel_size, _stride, _padding, _dilation, ceil_mode, True) out = out.squeeze(-1) indices = indices.squeeze(-1) if return_indices: return cast_to_adapter_tensor((out, indices)) return cast_to_adapter_tensor(out) else: _kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size) if stride is None: _stride = _kernel_size else: _stride = stride if isinstance(stride, tuple) else (stride, stride) _padding = padding if isinstance(padding, tuple) else (padding, padding) _dilation = dilation if isinstance(dilation, tuple) else (dilation, dilation) _extra_pad_h = 0 _extra_pad_w = 0 #TODO: this method do not support GRAPH_MODE # # if ceil_mode: # _input_shape = ms.ops.shape(input) # _valid_in_h = (_input_shape[-2] + _padding[0] * 2 - _dilation[0] * (_kernel_size[0] - 1) - 1) # _valid_out_h = _valid_in_h // _stride[0] # tmp = _valid_out_h * _stride[0] # if tmp < _valid_in_h: # _extra_pad_h = min(tmp + _stride[0] - _valid_in_h, kernel_size[0] - _padding[0] - 1) # _valid_in_w = (_input_shape[-1] + _padding[1] * 2 - _dilation[1] * (_kernel_size[1] - 1) - 1) # _valid_out_w = _valid_in_w // _stride[1] # tmp = _valid_out_w * _stride[1] # if tmp < _valid_in_w: # _extra_pad_w = min(tmp + _stride[1] - _valid_in_w, kernel_size[1] - _padding[1] - 1) # TODO: ms.ops.pad not support on ascend _pad = (((0, 0), (0, 0), (_padding[0], _padding[0]), (_padding[1], _padding[1] + _extra_pad_w))) _pad_op = _get_cache_prim(ms.ops.Pad)(_pad) #_pad = (_padding[1], _padding[1] + _extra_pad_w, _padding[0], _padding[0] + _extra_pad_h) #neg_inf = -float('inf') _max_pool = ms.ops.MaxPool(kernel_size=_kernel_size, strides=_stride, pad_mode='valid') ndim = input.ndim if ndim == 3: input = input.expand_dims(0) input = _pad_op(input) #input = ms.ops.pad(input, _pad, value=neg_inf) out = _max_pool(input) out = out.squeeze(0) else: input = _pad_op(input) #input = ms.ops.pad(input, _pad, value=neg_inf) out = _max_pool(input) return cast_to_adapter_tensor(out) def max_unpool1d(input, indices, kernel_size, stride, padding, output_size = None): input = cast_to_ms_tensor(input) indices = cast_to_ms_tensor(indices) if output_size is not None: output_size = tuple(output_size) out = ms.ops.max_unpool1d(input, indices, kernel_size, stride, padding, output_size) return out def max_unpool2d(input, indices, kernel_size, stride, padding, output_size = None): input = cast_to_ms_tensor(input) indices = cast_to_ms_tensor(indices) if output_size is not None: output_size = tuple(output_size) out = ms.ops.max_unpool2d(input, indices, kernel_size, stride, padding, output_size) return out def max_unpool3d(input, indices, kernel_size, stride, padding, output_size = None): input = cast_to_ms_tensor(input) indices = cast_to_ms_tensor(indices) if output_size is not None: output_size = tuple(output_size) out = ms.ops.max_unpool3d(input, indices, kernel_size, stride, padding, output_size) return cast_to_adapter_tensor(out) def linear(input, weight, bias=None): @constexpr def get_transpose_perm(shape): _rank = len(shape) perm = list(i for i in range(_rank)) _tmp = perm[-1] perm[-1] = perm[-2] perm[-2] = _tmp return tuple(perm) weight_shape = weight.shape weight_rank = len(weight_shape) if weight_rank not in (1, 2): raise ValueError("For nn.functional.linear, weight only support 2D or 1D input" f"but got {weight_rank}D input") if weight_rank == 2: weight = ms.ops.transpose(weight, get_transpose_perm(weight_shape)) input = cast_to_ms_tensor(input) output = ms.ops.matmul(input, weight) if bias is not None: output = ms.ops.add(output, bias) output = cast_to_adapter_tensor(output) return output def bilinear(input1, input2, weight, bias=None): input1 = cast_to_ms_tensor(input1) input2 = cast_to_ms_tensor(input2) weight = cast_to_ms_tensor(weight) x = ms.ops.matmul(input1.reshape(-1, input1.shape[-1]), weight.permute(1, 0, 2).reshape(weight.shape[1], -1)) x = ms.ops.mul(x, ms.ops.tile(input2.reshape(-1, input2.shape[-1]), (1, weight.shape[0]))) x = x.reshape(x.shape[0], weight.shape[0], -1) x = ms.ops.reduce_sum(x, -1) if bias is not None: bias = cast_to_ms_tensor(bias) x = ms.ops.bias_add(x, bias) output = x.reshape(*input1.shape[:-1], -1) return cast_to_adapter_tensor(output) def lp_pool1d(input, norm_type, kernel_size, stride = None, ceil_mode = False): input = cast_to_ms_tensor(input) output = ms.ops.lp_pool1d(input, norm_type, kernel_size, stride, ceil_mode) return cast_to_adapter_tensor(output) def lp_pool2d(input, norm_type, kernel_size, stride = None, ceil_mode = False): input = cast_to_ms_tensor(input) output = ms.ops.lp_pool2d(input, norm_type, kernel_size, stride, ceil_mode) return cast_to_adapter_tensor(output) def fractional_max_pool2d(input_x, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None): input_ms = cast_to_ms_tensor(input_x) out = ms.ops.fractional_max_pool2d(input_ms, kernel_size, output_size, output_ratio, return_indices, _random_samples) return cast_to_adapter_tensor(out) def fractional_max_pool3d(input_x, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None): input_ms = cast_to_ms_tensor(input_x) out = ms.ops.fractional_max_pool3d(input_ms, kernel_size, output_size, output_ratio, return_indices, _random_samples) return cast_to_adapter_tensor(out) def avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True): unsupported_attr(ceil_mode) unsupported_attr(count_include_pad) if isinstance(kernel_size, int): kernel_size = (kernel_size, 1) else: kernel_size = kernel_size + (1,) if stride is None: stride = (kernel_size, 1) elif isinstance(stride, int): stride = (stride, 1) else: stride = stride + (1,) if isinstance(padding, int): padding = (padding, 0) else: padding = padding + (0,) pad_ops = ms.ops.Pad(((0, 0), (0, 0), (padding[0], padding[0]), (padding[1], padding[1]))) avg_pool_ops = ms.ops.AvgPool(kernel_size=kernel_size, strides=stride, pad_mode='valid') input = cast_to_ms_tensor(input) ndim = input.ndim if ndim == 2: input = input.expand_dims(0) input = input.expand_dims(-1) input = pad_ops(input) out = avg_pool_ops(input) out = out.squeeze(-1) out = out.squeeze(0) else: input = input.expand_dims(-1) input = pad_ops(input) out = avg_pool_ops(input) out = out.squeeze(-1) return cast_to_adapter_tensor(out) def avg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None): input_ms = cast_to_ms_tensor(input) if input_ms.ndim == 4: _input_ms = input_ms[None,...] if stride is None: _stride = kernel_size else: _stride = stride if divisor_override is None: _divisor_override = 0 else: _divisor_override = divisor_override if isinstance(padding, tuple) and len(padding) == 3: _padding = (padding[0], padding[0], padding[1], padding[1], padding[2], padding[2]) else: _padding = padding out = ms.ops.avg_pool3d(_input_ms, kernel_size, _stride, _padding, ceil_mode, count_include_pad, _divisor_override) out = out.squeeze(0) else: _input_ms = input_ms if stride is None: _stride = kernel_size else: _stride = stride if divisor_override is None: _divisor_override = 0 else: _divisor_override = divisor_override if isinstance(padding, tuple) and len(padding) == 3: _padding = (padding[0], padding[0], padding[1], padding[1], padding[2], padding[2]) else: _padding = padding out = ms.ops.avg_pool3d(_input_ms, kernel_size, _stride, _padding, ceil_mode, count_include_pad, _divisor_override) return cast_to_adapter_tensor(out) def max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False): input = cast_to_ms_tensor(input) if input.dtype in {ms.int16, ms.int32, ms.int64}: raise TypeError("'max_pool1d' not implemented for int.") if return_indices is True or dilation != 1 or ceil_mode is True: if input.ndim == 2: _input_ms = input[None,...,None,None] elif input.ndim == 3: _input_ms = input[..., None, None] else: _input_ms = input if isinstance(kernel_size, int): _kernel_size = (kernel_size, 1, 1) elif isinstance(kernel_size, tuple): _kernel_size = kernel_size + (1, 1) else: _kernel_size = kernel_size if stride is None: _stride = (kernel_size, 1, 1) elif isinstance(stride, int): _stride = (stride, 1, 1) elif isinstance(stride, tuple): _stride = stride + (1, 1) else: _stride = stride _padding = (padding, 0, 0) _dilation = (dilation, 1, 1) out = ms.ops.max_pool3d(_input_ms, _kernel_size, _stride, _padding, _dilation, ceil_mode, return_indices) if isinstance(out, tuple): out = list(out) for id, value in enumerate(out): out[id] = value.squeeze(-1).squeeze(-1) if input.ndim == 2: out[id] = out[id].squeeze(0) out = tuple(out) else: out = out.squeeze(-1).squeeze(-1) if input.ndim == 2: out = out.squeeze(0) return cast_to_adapter_tensor(out) else: _kernel_size = (1, kernel_size) if stride is None: _stride = _kernel_size else: _stride = (1, stride) _padding = (0, padding) _dilation = (1, dilation) _extra_pad_w = 0 #TODO: this method do not support GRAPH_MODE # if ceil_mode: # _input_shape = ms.ops.shape(input) # _valid_in_w = (_input_shape[-1] + _padding[1] * 2 - _dilation[1] * (_kernel_size[1] - 1) - 1) # _valid_out_w = _valid_in_w // _stride[1] # tmp = _valid_out_w * _stride[1] # if tmp < _valid_in_w: # _extra_pad_w = min(tmp + _stride[1] - _valid_in_w, kernel_size[1] - _padding[1] - 1) # TODO: ms.ops.pad not support on ascend _pad = (((0, 0), (0, 0), (_padding[0], _padding[0]), (_padding[1], _padding[1] + _extra_pad_w))) _pad_op = _get_cache_prim(ms.ops.Pad)(_pad) # _pad = (_padding[1], _padding[1] + _extra_pad_w, _padding[0], _padding[0]) # neg_inf = -float('inf') _max_pool = ms.ops.MaxPool(kernel_size=_kernel_size, strides=_stride, pad_mode='valid') ndim = input.ndim if ndim == 2: input = input.expand_dims(0) input = input.expand_dims(2) input = _pad_op(input) # input = ms.ops.pad(input, _pad, value=neg_inf) out = _max_pool(input) out = out.squeeze(2) out = out.squeeze(0) else: input = input.expand_dims(2) input = _pad_op(input) # input = ms.ops.pad(input, _pad, value=neg_inf) out = _max_pool(input) out = out.squeeze(2) return cast_to_adapter_tensor(out) def max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False): input_ms = cast_to_ms_tensor(input) if input_ms.ndim == 4: _input_ms = input_ms[None, ...] else: _input_ms = input_ms out = ms.ops.max_pool3d(_input_ms, kernel_size, stride, padding, dilation, ceil_mode, return_indices) if input_ms.ndim == 4: if isinstance(out, tuple): out = list(out) for id, value in enumerate(out): out[id] = value.squeeze(0) out = tuple(out) else: out = out.squeeze(0) return cast_to_adapter_tensor(out) def conv_transpose1d(inputs, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1): inputs = cast_to_ms_tensor(inputs) weight = cast_to_ms_tensor(weight) has_bias = bias is not None bias = cast_to_ms_tensor(bias) if bias is not None else 'zeros' if len(inputs.shape) != 3: raise ValueError("the rank of inputs tensor should be 3.") if len(weight.shape) != 3: raise ValueError("the rank of weight tensor should be 3") in_channel = inputs.shape[1] out_channel = weight.shape[1] * groups kernel_size = weight.shape[2] if stride != 1 and padding == (kernel_size - 1) // 2 and output_padding == stride - 1: pad_mode = 'same' padding = 0 raise Warning("pad_mode = same is some thing wrong, please switch to others") elif stride != 1 and padding == 0 and output_padding == 0: pad_mode = 'valid' padding = 0 else: pad_mode = 'pad' _conv_1d_transpose = nn.Conv1dTranspose( in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size, stride=stride, pad_mode=pad_mode, padding=padding, dilation=dilation, group=groups, has_bias=has_bias, weight_init=weight, bias_init=bias) out = _conv_1d_transpose(inputs) return cast_to_adapter_tensor(out) def conv_transpose2d(inputs, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1): inputs = cast_to_ms_tensor(inputs) weight = cast_to_ms_tensor(weight) has_bias = bias is not None bias = cast_to_ms_tensor(bias) if bias is not None else 'zeros' if len(inputs.shape) != 4: raise ValueError("the rank of inputs tensor should be 4.") if len(weight.shape) != 4: raise ValueError("the rank of weight tensor should be 4") in_channel = inputs.shape[1] out_channel = weight.shape[1] * groups kernel_size = weight.shape[2:] if isinstance(padding, tuple): padding = list(np.repeat(padding, 2)) # convert default data type 'int64' to 'int' padding = tuple(map(int, padding)) if stride != 1 and padding == 0 and output_padding == 0: pad_mode = 'valid' padding = 0 else: pad_mode = 'pad' _conv_2d_transpose = nn.Conv2dTranspose( in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size, stride=stride, pad_mode=pad_mode, padding=padding, dilation=dilation, group=groups, has_bias=has_bias, weight_init=weight, bias_init=bias) out = _conv_2d_transpose(inputs) return cast_to_adapter_tensor(out) def conv_transpose3d(inputs, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1): inputs = cast_to_ms_tensor(inputs) weight = cast_to_ms_tensor(weight) has_bias = bias is not None bias = cast_to_ms_tensor(bias) if bias is not None else 'zeros' if len(inputs.shape) != 5: raise ValueError("the rank of inputs tensor should be 5.") if len(weight.shape) != 5: raise ValueError("the rank of weight tensor should be 5") in_channel = inputs.shape[1] out_channel = weight.shape[1] * groups kernel_size = weight.shape[2:] if isinstance(padding, tuple): padding = list(np.repeat(padding, 2)) # convert default data type 'int64' to 'int' padding = tuple(map(int, padding)) if stride != 1 and padding == 0 and output_padding == 0: pad_mode = 'valid' padding = 0 else: pad_mode = 'pad' _conv_3d_transpose = nn.Conv3dTranspose( in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size, stride=stride, pad_mode=pad_mode, padding=padding, dilation=dilation, group=groups, has_bias=has_bias, weight_init=weight, bias_init=bias) out = _conv_3d_transpose(inputs) return cast_to_adapter_tensor(out) def affine_grid(theta, size, align_corners=None): theta = cast_to_ms_tensor(theta) if align_corners is None: align_corners = False # TODO：the input argument[theta] must be a type of {Tensor[Float16], Tensor[Float32]} if theta.dtype == ms.float64: theta = theta.astype(ms.float32) output = ms.ops.affine_grid(theta, size, align_corners) return cast_to_adapter_tensor(output) def batch_norm(inputs, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-05): inputs = cast_to_ms_tensor(inputs) running_mean = cast_to_ms_tensor(running_mean) running_var = cast_to_ms_tensor(running_var) weight = cast_to_ms_tensor(weight) if weight is not None else weight bias = cast_to_ms_tensor(bias) if bias is not None else bias reduced_dim = tuple(i for i in range(inputs.dim()) if i != 1) normalized_shape = [1] * len(inputs.shape) normalized_shape[1] = inputs.shape[1] if training: mean = inputs.mean(axis=reduced_dim, keep_dims=True) var = inputs.var(reduced_dim, keepdims=True, ddof=False) mean_update = mean.squeeze() var_update = inputs.var(axis=reduced_dim, ddof=True) out = (inputs - mean) / ms.ops.sqrt(var + eps) # parameters updating reserved for future use running_mean = (1 - momentum) * running_mean + momentum * mean_update running_var = (1 - momentum) * running_var + momentum * var_update else: out = (inputs - running_mean.view(*normalized_shape)) / ms.ops.sqrt(running_var.view(*normalized_shape) + eps) if weight is not None: out = out * weight.view(*normalized_shape) if bias is not None: out = out + bias.view(*normalized_shape) return cast_to_adapter_tensor(out) def group_norm(inputs, num_groups, weight=None, bias=None, eps=1e-05): inputs = cast_to_ms_tensor(inputs) weight = cast_to_ms_tensor(weight) if weight is not None else weight bias = cast_to_ms_tensor(bias) if bias is not None else bias inputs_shape = list(inputs.shape) shape = [inputs_shape[0]] + [num_groups, inputs_shape[1] // num_groups] + inputs_shape[2:] normalized_shape = [1] * len(inputs.shape) normalized_shape[1] = inputs_shape[1] reduced_dim = tuple(i for i in range(len(shape) - 1, 1, -1)) inputs = inputs.reshape(*shape) mean = inputs.mean(axis=reduced_dim, keep_dims=True) var = inputs.var(axis=reduced_dim, keepdims=True, ddof=False) out = (inputs - mean) / ms.ops.sqrt(var + eps) out = out.reshape(*inputs_shape) if weight is not None: out = out * weight.view(*normalized_shape) if bias is not None: out = out + bias.view(*normalized_shape) return cast_to_adapter_tensor(out) def instance_norm(inputs, running_mean=None, running_var=None, weight=None, bias=None, use_input_stats=True, momentum=0.1, eps=1e-05): inputs = cast_to_ms_tensor(inputs) running_mean = cast_to_ms_tensor(running_mean) running_var = cast_to_ms_tensor(running_var) weight = cast_to_ms_tensor(weight) if weight is not None else weight bias = cast_to_ms_tensor(bias) if bias is not None else bias reduced_dim = tuple(i for i in range(inputs.dim()) if i not in [0, 1]) normalized_shape = [1] * len(inputs.shape) normalized_shape[1] = inputs.shape[1] shape = [1] * len(inputs.shape) shape[:2] = inputs.shape[:2] if use_input_stats: mean = inputs.mean(axis=reduced_dim) var = inputs.var(axis=reduced_dim, ddof=False) mean_update = mean.mean(0) var_update = inputs.var(axis=reduced_dim, ddof=True).mean(0) out = (inputs - mean.view(*shape)) / ms.ops.sqrt(var.view(*shape) + eps) running_mean = (1 - momentum) * running_mean + momentum * mean_update running_var = (1 - momentum) * running_var + momentum * var_update else: out = (inputs - running_mean.view(*normalized_shape)) \ / ms.ops.sqrt(running_var.view(*normalized_shape) + eps) if weight is not None: out = out * weight.view(*normalized_shape) if bias is not None: out = out + bias.view(*normalized_shape) return cast_to_adapter_tensor(out) def layer_norm(inputs, normalized_shape, weight=None, bias=None, eps=1e-05): inputs = cast_to_ms_tensor(inputs) if weight is not None: weight = cast_to_ms_tensor(weight) else: weight = ms.Tensor(np.ones(normalized_shape), inputs.dtype) if bias is not None: bias = cast_to_ms_tensor(bias) else: bias = ms.Tensor(np.zeros(normalized_shape), inputs.dtype) assert inputs.shape[-len(normalized_shape):] == normalized_shape _layer_norm = ms.ops.LayerNorm(epsilon=eps) out = _layer_norm(inputs, weight, bias) return cast_to_adapter_tensor(out[0]) def prelu(input, weight): #TODO:ms.ops.prelu only suports float16 and float32, not float64. input = cast_to_ms_tensor(input) weight = cast_to_ms_tensor(weight) output = ms.ops.prelu(input, weight) return cast_to_adapter_tensor(output) def poisson_nll_loss(input, target, log_input=True, full=False, size_average=None, eps=1e-08, reduce=None, reduction='mean'): input_ms = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) pi = 3.141592653589793 if reduce is not None or size_average is not None: reduction = _get_reduce_string(size_average, reduce) if reduction not in ('none', 'mean', 'sum'): raise ValueError(reduction + " is not valid") if log_input: ret = ms.ops.exp(input) - target * input else: ret = input_ms - target * ms.ops.log(input_ms + eps) if full: cond = ms.ops.gt(target, 1) out = target * ms.ops.log(target) - target + 0.5 * ms.ops.log(2 * pi * target) out = ms.ops.select(cond, out, ms.ops.zeros_like(input_ms)) ret = ret + out if reduction == "mean": ret = ms.ops.mean(ret) elif reduction == "sum": ret = ms.ops.sum(ret) return cast_to_adapter_tensor(ret) def triplet_margin_with_distance_loss(anchor, positive, negative, *, distance_function=None, margin=1.0, swap=False, reduction='mean'): distance_function = distance_function if distance_function is not None else pairwise_distance anchor = cast_to_ms_tensor(anchor) positive = cast_to_ms_tensor(positive) negative = cast_to_ms_tensor(negative) positive_dist = distance_function(anchor, positive) negative_dist = distance_function(anchor, negative) if swap: swap_dist = distance_function(positive, negative) negative_dist = ms.ops.minimum(negative_dist, swap_dist) output = ms.ops.clamp(positive_dist - negative_dist + margin, min=0.0) if reduction == "mean": ret = output.mean() elif reduction == "sum": ret = output.sum() else: ret = output return cast_to_adapter_tensor(ret) def conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): # Todo: not support float64, change to float32 now input_ms = cast_to_ms_tensor(input) weight_ms = cast_to_ms_tensor(weight) is_float64 = False if input_ms.dtype in (ms.float64, ms.double): input_ms = input_ms.astype(ms.float32) weight_ms = weight_ms.astype(ms.float32) is_float64 = True if isinstance(stride, int): stride = (stride, stride, stride) elif len(stride)==1: stride = (stride[0], stride[0], stride[0]) pad_mode = "pad" if isinstance(padding, int): padding = (padding, padding, padding, padding, padding, padding) elif isinstance(padding, tuple): if len(padding)==1: padding = (padding[0], padding[0], padding[0], padding[0], padding[0], padding[0]) else: padding = (padding[0], padding[0], padding[1], padding[1], padding[2], padding[2]) else: pad_mode = padding padding = 0 if isinstance(dilation, int): dilation = (dilation, dilation, dilation) elif len(dilation) == 1: dilation = (dilation[0], dilation[0], dilation[0]) output = ms.ops.conv3d(input_ms, weight_ms, pad_mode, padding, stride, dilation, groups) if bias is not None: # TODO: ms.ops.biasadd also not support float64 if bias.dtype != output.dtype: bias = bias.astype(output.dtype) output = ms.ops.bias_add(output, bias) if is_float64: output = output.astype(ms.float64) return cast_to_adapter_tensor(output) def unfold(input, kernel_size, dilation=1, padding=0, stride=1): # TODO: do not support on GPU input_ms = cast_to_ms_tensor(input) output = ms.ops.unfold(input_ms, kernel_size, dilation, padding, stride) output = output.reshape(output.shape[0], output.shape[1], -1) return cast_to_adapter_tensor(output) def fold(input, output_size, kernel_size, dilation=1, padding=0, stride=1): # TODO: do not support on Ascend input_ms = cast_to_ms_tensor(input) ndim = input_ms.ndim if ndim == 2: input_ms = input_ms.expand_dims(0) shape = input_ms.shape if isinstance(kernel_size, int): shape_tmp = kernel_size * kernel_size else: shape_tmp = kernel_size[0] * kernel_size[1] input_ms = input_ms.reshape(shape[0], -1, shape_tmp, shape[2]) output = ms.ops.fold(input_ms, ms.Tensor(output_size), kernel_size, dilation, padding, stride) if ndim == 2: output = output.squeeze(0) return cast_to_adapter_tensor(output) msadapter/ms_adapter/pytorch/nn/init.py000066400000000000000000000120211436713554100206020ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import warnings import math def _calculate_fan_in_and_fan_out(tensor): dimensions = tensor.dim() if dimensions < 2: raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions") num_input_fmaps = tensor.size(1) num_output_fmaps = tensor.size(0) receptive_field_size = 1 if tensor.dim() > 2: for s in tensor.shape[2:]: receptive_field_size *= s fan_in = num_input_fmaps * receptive_field_size fan_out = num_output_fmaps * receptive_field_size return fan_in, fan_out def calculate_gain(nonlinearity, param=None): linear_fns = ['linear', 'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d'] if nonlinearity in linear_fns or nonlinearity == 'sigmoid': res = 1 elif nonlinearity == 'tanh': res = 5.0 / 3 elif nonlinearity == 'relu': res = math.sqrt(2.0) elif nonlinearity == 'leaky_relu': if param is None: negative_slope = 0.01 elif not isinstance(param, bool) and isinstance(param, int) or isinstance(param, float): negative_slope = param else: raise ValueError("For 'HeUniform', 'negative_slope' {} is not a valid number." "When 'nonlinearity' has been set to " "'leaky_relu', 'negative_slope' should be int or float type, but got " "{}.".format(param, type(param))) res = math.sqrt(2.0 / (1 + negative_slope ** 2)) else: raise ValueError("For 'HeUniform', the argument 'nonlinearity' should be one of " "['sigmoid', 'tanh', 'relu' or 'leaky_relu'], " "but got {}.".format(nonlinearity)) return res def _no_grad_trunc_normal_(tensor, mean, std, a, b): # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf def norm_cdf(x): # Computes standard normal cumulative distribution function return (1. + math.erf(x / math.sqrt(2.))) / 2. if (mean < a - 2 * std) or (mean > b + 2 * std): warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " "The distribution of values may be incorrect.", stacklevel=2) # Values are generated by using a truncated uniform distribution and # then using the inverse CDF for the normal distribution. # Get upper and lower cdf values l = norm_cdf((a - mean) / std) u = norm_cdf((b - mean) / std) # Uniformly fill tensor with values from [l, u], then translate to # [2l-1, 2u-1]. tensor = tensor.uniform_(2 * l - 1, 2 * u - 1) # Use inverse cdf transform for normal distribution to get truncated # standard normal tensor.erfinv_() # Transform to proper mean, std tensor.mul_(std * math.sqrt(2.)) tensor.add_(mean) # Clamp to ensure it's in the proper range tensor.clamp_(min=a, max=b) return tensor def _calculate_correct_fan(tensor, mode): mode = mode.lower() valid_modes = ['fan_in', 'fan_out'] if mode not in valid_modes: raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes)) fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) return fan_in if mode == 'fan_in' else fan_out def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'): if 0 in tensor.shape: warnings.warn("Initializing zero-element tensors is a no-op") return tensor fan = _calculate_correct_fan(tensor, mode) gain = calculate_gain(nonlinearity, a) std = gain / math.sqrt(fan) bound = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation return tensor.uniform_(-bound, bound) def kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'): if 0 in tensor.shape: warnings.warn("Initializing zero-element tensors is a no-op") return tensor fan = _calculate_correct_fan(tensor, mode) gain = calculate_gain(nonlinearity, a) std = gain / math.sqrt(fan) return tensor.normal_(0, std) def xavier_normal_(tensor, gain = 1.): fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) std = gain * math.sqrt(2.0 / float(fan_in + fan_out)) return tensor.normal_(0., std) def xavier_uniform_(tensor, gain = 1.): fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) std = gain * math.sqrt(2.0 / float(fan_in + fan_out)) a = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation return tensor.uniform_(-a, a) def zeros_(tensor): return tensor.zero_() def ones_(tensor): return tensor.fill_(1.0) def constant_(tensor, val): return tensor.fill_(val) def trunc_normal_(tensor, mean = 0., std = 1., a = -2., b = 2.): return _no_grad_trunc_normal_(tensor, mean, std, a, b) def normal_(tensor, mean = 0., std = 1.): return tensor.normal_(mean, std) def uniform_(tensor, a = 0., b = 1.): return tensor.uniform_(a, b) msadapter/ms_adapter/pytorch/nn/modules/000077500000000000000000000000001436713554100207415ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/nn/modules/__init__.py000066400000000000000000000067041436713554100230610ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from .activation import * from .linear import * from .flatten import * from .conv import * from .distance import * from .batchnorm import * from .pooling import * from .unpooling import * from .loss import * from .padding import * from .rnn import * from .sparse import * from .module import Module from .container import Sequential, ModuleList, ModuleDict, ParameterList, ParameterDict from .dropout import Dropout, Dropout1d, Dropout2d, Dropout3d, AlphaDropout, FeatureAlphaDropout from .upsampling import * from .normalization import * from .pixel_shuffle import * from .channelshuffle import * from .fold import * from .adaptive import AdaptiveLogSoftmaxWithLoss __all__ = [ 'Linear', 'LazyLinear', 'Bilinear', 'Flatten', 'Unflatten', 'Conv1d', 'Conv2d', 'Conv3d', 'ConvTranspose1d', 'ConvTranspose2d', 'ConvTranspose3d', 'LazyConv1d', 'LazyConv2d', 'LazyConv3d', 'LazyConvTranspose1d', 'LazyConvTranspose2d', 'LazyConvTranspose3d', 'Fold', 'Unfold', 'BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'LazyBatchNorm1d', 'LazyBatchNorm2d', 'LazyBatchNorm3d', 'SyncBatchNorm', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d', 'LazyInstanceNorm1d', 'LazyInstanceNorm2d', 'LazyInstanceNorm3d', 'FractionalMaxPool2d', 'FractionalMaxPool3d', 'AdaptiveAvgPool1d', 'AdaptiveAvgPool2d', 'AdaptiveAvgPool3d', 'AdaptiveMaxPool1d', 'AdaptiveMaxPool2d', 'AdaptiveMaxPool3d', 'MaxPool1d', 'MaxPool2d', 'MaxPool3d', 'AvgPool1d', 'AvgPool2d', 'AvgPool3d', 'LPPool1d', 'LPPool2d', 'Identity', 'MaxUnpool1d', 'MaxUnpool2d', 'MaxUnpool3d', 'ReLU', 'ReLU6', 'Hardtanh', 'SiLU', 'Hardswish', 'LeakyReLU', 'Sigmoid', 'RReLU', 'PReLU', 'SELU', 'CELU', 'GELU', 'Mish', 'Softshrink', 'Hardshrink', 'Tanh', 'Tanhshrink', 'Threshold', 'Softplus', 'Softsign', 'Softmax', 'LogSoftmax', 'Softmax2d', 'Softmin', 'GLU', 'AdaptiveLogSoftmaxWithLoss', 'MultiheadAttention', 'Hardsigmoid', 'Module', 'Sequential', 'ModuleList', "ParameterList", "ParameterDict", "ModuleDict", 'Dropout', 'Dropout1d', 'Dropout2d', 'Dropout3d', 'AlphaDropout', 'FeatureAlphaDropout', 'Upsample', 'SmoothL1Loss', 'L1Loss', 'MSELoss', 'CrossEntropyLoss', 'NLLLoss', 'KLDivLoss', 'BCELoss', 'BCEWithLogitsLoss', 'HuberLoss', 'SoftMarginLoss', 'CosineEmbeddingLoss', 'MultiMarginLoss', 'TripletMarginLoss', 'PoissonNLLLoss', 'GaussianNLLLoss', 'HingeEmbeddingLoss', 'MultiLabelMarginLoss', 'MultiLabelSoftMarginLoss', 'TripletMarginWithDistanceLoss', 'MarginRankingLoss', 'CTCLoss', 'LogSigmoid', 'ELU', 'ConstantPad1d', 'ConstantPad2d', 'ConstantPad3d', 'ReflectionPad1d', 'ReflectionPad2d', 'ReflectionPad3d', 'ZeroPad2d', 'ReplicationPad1d', 'ReplicationPad2d', 'ReplicationPad3d', 'RNN', 'GRU', 'LSTM', 'RNNCell', 'GRUCell', 'LSTMCell', 'LayerNorm', 'GroupNorm', 'LocalResponseNorm', 'UpsamplingNearest2d', 'UpsamplingBilinear2d', 'PairwiseDistance', 'CosineSimilarity', 'Embedding', 'PixelShuffle', 'PixelUnshuffle', 'ChannelShuffle' ] msadapter/ms_adapter/pytorch/nn/modules/activation.py000066400000000000000000000470601436713554100234630ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import warnings import numpy as np from mindspore.ops import functional as F from mindspore.ops import operations as P from mindspore.common import dtype as mstype import mindspore as ms from mindspore import nn from mindspore._checkparam import Validator as validator from ms_adapter.pytorch.nn.parameter import Parameter import ms_adapter.pytorch.nn.functional as ms_torch_nn_func from ms_adapter.pytorch.tensor import Tensor, tensor, cast_to_ms_tensor, cast_to_adapter_tensor from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.common._inner import _inplace_assign, _inplace_limit_pynative from .module import Module __all__ = ['ReLU', 'Hardtanh', 'ReLU6', 'SiLU', 'Hardswish', 'LeakyReLU', 'Sigmoid', 'LogSigmoid', 'ELU', 'RReLU', 'SELU', 'CELU', 'GELU', 'Mish', 'Softshrink', 'Tanh', 'Tanhshrink','Threshold', 'Softmax', 'LogSoftmax', 'Softmin', 'Softsign', 'GLU', 'Hardshrink', 'MultiheadAttention', 'Hardsigmoid', 'PReLU', 'Softplus', 'Softmax2d'] class ReLU(Module): r"""Applies the rectified linear unit function element-wise: :math:`\text{ReLU}(x) = (x)^+ = \max(0, x)` Args: inplace: can optionally do the operation in-place. Default: ``False`` Shape: - Input: :math:`(*)`, where :math:`*` means any number of dimensions. - Output: :math:`(*)`, same shape as the input. .. image:: ../scripts/activation_images/ReLU.png Examples:: >>> import ms_adapter.pytorch as torch >>> import ms_adapter.pytorch.nn as nn >>> m = nn.ReLU() >>> input = torch.randn(2) >>> output = m(input) """ def __init__(self, inplace=False): super(ReLU, self).__init__() self.relu = P.ReLU() self.inplace = inplace _inplace_limit_pynative(inplace, "ReLU") def forward(self, input): input_ms = cast_to_ms_tensor(input) output = self.relu(input_ms) return _inplace_assign(input, self.inplace, output) def extra_repr(self): inplace_str = 'inplace=True' if self.inplace else '' return inplace_str class Hardtanh(Module): def __init__( self, min_val=-1., max_val=1., inplace=False, min_value=None, max_value=None ): super(Hardtanh, self).__init__() _inplace_limit_pynative(inplace, "Hardtanh") if min_value is not None: warnings.warn("keyword argument min_value is deprecated and rename to min_val") min_val = min_value if max_value is not None: warnings.warn("keyword argument max_value is deprecated and rename to max_val") max_val = max_value self.min_val = min_val self.max_val = max_val self.inplace = inplace if self.max_val <= self.min_val: raise ValueError('`max_val` must be larger than `min_val` in `{}`, but get `max_val`:{} and ' '`min_val`:{}'.format(self.__class__.__name__, self.max_val, self.min_val)) self.hardtanh = nn.Hardtanh(min_val, max_val) def forward(self, input): input_ms = cast_to_ms_tensor(input) output = self.hardtanh(input_ms) return _inplace_assign(input, self.inplace, output) def extra_repr(self): inplace_str = ', inplace=True' if self.inplace else '' return 'min_val={}, max_val={}{}'.format( self.min_val, self.max_val, inplace_str ) class ReLU6(Hardtanh): def __init__(self, inplace=False): _inplace_limit_pynative(inplace, "ReLU6") super(ReLU6, self).__init__(0., 6., inplace) def extra_repr(self): inplace_str = 'inplace=True' if self.inplace else '' return inplace_str class SiLU(Module): def __init__(self, inplace=False): super(SiLU, self).__init__() _inplace_limit_pynative(inplace, "SiLU") self.inplace = inplace self.sigmoid = P.Sigmoid() def forward(self, input): input_ms = cast_to_ms_tensor(input) output = self.sigmoid(input_ms) * input_ms return _inplace_assign(input, self.inplace, output) def extra_repr(self) -> str: inplace_str = 'inplace=True' if self.inplace else '' return inplace_str class Hardswish(Module): def __init__(self, inplace=False): super(Hardswish, self).__init__() _inplace_limit_pynative(inplace, "Hardswish") self.inplace = inplace self.hardswish = P.HSwish() def forward(self, input): input_ms = cast_to_ms_tensor(input) output = self.hardswish(input_ms) return _inplace_assign(input, self.inplace, output) class LeakyReLU(Module): def __init__(self, negative_slope=1e-2, inplace=False): super(LeakyReLU, self).__init__() _inplace_limit_pynative(inplace, "LeakyReLU") self.negative_slope = negative_slope self.inplace = inplace self.greater_equal = P.GreaterEqual() self.mul = P.Mul() self.select_op = P.Maximum() if self.negative_slope > 1: self.select_op = P.Minimum() self.cast = P.Cast() def forward(self, input): input_ms = cast_to_ms_tensor(input) alpha_array = self.cast(F.scalar_to_tensor(self.negative_slope), input_ms.dtype) output = self.select_op(alpha_array * input_ms, input_ms) return _inplace_assign(input, self.inplace, output) def extra_repr(self) -> str: inplace_str = ', inplace=True' if self.inplace else '' return 'negative_slope={}{}'.format(self.negative_slope, inplace_str) class Sigmoid(Module): def __init__(self): super(Sigmoid, self).__init__() self.sigmoid = P.Sigmoid() def forward(self, input): input = cast_to_ms_tensor(input) output = self.sigmoid(input) return cast_to_adapter_tensor(output) class LogSigmoid(Module): def __init__(self): super(LogSigmoid, self).__init__() self.logsigmoid = ms_torch_nn_func.logsigmoid def forward(self, input): return self.logsigmoid(input) class ELU(Module): def __init__(self, alpha: float=1., inplace: bool=False): super(ELU, self).__init__() _inplace_limit_pynative(inplace, "ELU") self.elu = ms_torch_nn_func.elu self.alpha = alpha self.inplace = inplace def forward(self, input): return self.elu(input, self.alpha, self.inplace) class RReLU(Module): def __init__( self, lower=1./8, upper=1./3, inplace=False ): super(RReLU, self).__init__() _inplace_limit_pynative(inplace, "RReLU") self.lower = lower self.upper = upper self.inplace = inplace self.rrelu = ms.nn.RReLU(lower=self.lower, upper=self.upper) def forward(self, input): input_ms = cast_to_ms_tensor(input) out = self.rrelu(input_ms) return _inplace_assign(input, self.inplace, out) def extra_repr(self): inplace_str = ', inplace=True' if self.inplace else '' return 'lower={}, upper={}{}'.format(self.lower, self.upper, inplace_str) class SELU(Module): def __init__(self, inplace=False): super(SELU, self).__init__() _inplace_limit_pynative(inplace, "SELU") self.inplace = inplace self.selu = P.SeLU() def forward(self, input): input_ms = cast_to_ms_tensor(input) out = self.selu(input_ms) return _inplace_assign(input, self.inplace, out) def extra_repr(self): inplace_str = 'inplace=True' if self.inplace else '' return inplace_str class CELU(Module): def __init__(self, alpha=1., inplace=False): super(CELU, self).__init__() _inplace_limit_pynative(inplace, "CELU") self.alpha = alpha self.inplace = inplace self.celu = P.CeLU(alpha=self.alpha) def forward(self, input): input_ms = cast_to_ms_tensor(input) out = self.celu(input_ms) return _inplace_assign(input, self.inplace, out) def extra_repr(self): inplace_str = ', inplace=True' if self.inplace else '' return 'alpha={}{}'.format(self.alpha, inplace_str) class GELU(Module): def __init__(self, approximate='none'): super(GELU, self).__init__() self.gelu = ms_torch_nn_func.gelu self.approximate = approximate def forward(self, input): out = self.gelu(input, self.approximate) return out class Mish(Module): def __init__(self, inplace=False): super(Mish, self).__init__() _inplace_limit_pynative(inplace, "Mish") self.inplace = inplace self.mish = P.Mish() def forward(self, input): input_ms = cast_to_ms_tensor(input) out = self.mish(input_ms) return _inplace_assign(input, self.inplace, out) def extra_repr(self): inplace_str = 'inplace=True' if self.inplace else '' return inplace_str class Softshrink(Module): def __init__(self, lambd=0.5): super(Softshrink, self).__init__() self.lambd = lambd self.softshrink = P.SoftShrink(lambd=self.lambd) def forward(self, input): input = cast_to_ms_tensor(input) out = self.softshrink(input) return cast_to_adapter_tensor(out) def extra_repr(self): return str(self.lambd) class Tanh(Module): def forward(self,input): return ms_torch_nn_func.tanh(input) class Tanhshrink(Module): def forward(self,input): return ms_torch_nn_func.tanhshrink(input) class Threshold(Module): def __init__(self, threshold, value, inplace=False): super(Threshold, self).__init__() _inplace_limit_pynative(inplace, "Threshold") self.threshold = threshold self.value = value self.inplace = inplace def forward(self, input): return ms_torch_nn_func.threshold(input, self.threshold, self.value, self.inplace) def extra_repr(self): inplace_str = ', inplace=True' if self.inplace else '' return 'threshold={}, value={}{}'.format(self.threshold, self.value, inplace_str) class Softmax(Module): def __init__(self, dim=None): super(Softmax, self).__init__() self.softmax = ms_torch_nn_func.softmax self.dim = dim def forward(self, input): return self.softmax(input, self.dim) def extra_repr(self): return 'dim={dim}'.format(dim=self.dim) class LogSoftmax(Module): def __init__(self, dim=None): super(LogSoftmax, self).__init__() self.logsoftmax = ms_torch_nn_func.log_softmax self.dim = dim def forward(self, input): return self.logsoftmax(input, self.dim) def extra_repr(self): return 'dim={dim}'.format(dim=self.dim) class Softmin(Module): def __init__(self, dim=None): super(Softmin, self).__init__() self.softmin = ms_torch_nn_func.softmin self.dim = dim def forward(self, input): return self.softmin(input, self.dim) def extra_repr(self): return 'dim={dim}'.format(dim=self.dim) class Softsign(Module): def __init__(self): super(Softsign, self).__init__() self.softsign = ms_torch_nn_func.softsign def forward(self, input): return self.softsign(input) class GLU(Module): def __init__(self, dim=-1): super(GLU, self).__init__() self.glu = ms_torch_nn_func.glu self.dim = dim def forward(self, input): return self.glu(input, self.dim) def extra_repr(self): return 'dim={dim}'.format(dim=self.dim) class Hardshrink(Module): def __init__(self, lambd: float=0.5): super(Hardshrink, self).__init__() self.lambd = lambd def forward(self, input): return ms_torch_nn_func.hardshrink(input, self.lambd) def extra_repr(self) -> str: return '{}'.format(self.lambd) class Hardsigmoid(Module): def __init__(self, inplace: bool=False): super(Hardsigmoid, self).__init__() _inplace_limit_pynative(inplace, "Hardsigmoid") self.inplace = inplace def forward(self, input): return ms_torch_nn_func.hardsigmoid(input, self.inplace) class MultiheadAttention(Module): def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, \ add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None): super(MultiheadAttention, self).__init__() if bias is not True: raise ValueError(f"`bias` can only be set to 'True', but got {bias}") if add_bias_kv: raise ValueError(f"`add_bias_kv` can only be set to 'False', but got {add_bias_kv}") if add_zero_attn: raise ValueError(f"`add_zero_attn` can only be set to 'False', but got {add_zero_attn}") unsupported_attr(kdim) unsupported_attr(vdim) unsupported_attr(device) self.embed_dim = embed_dim self.kdim = kdim if kdim is not None else embed_dim self.vdim = vdim if vdim is not None else embed_dim self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim self.num_heads = num_heads self.dropout = dropout self.add_bias_kv = add_bias_kv self.add_zero_attn = add_zero_attn self.kdim = kdim self.vdim = vdim self.batch_first = batch_first self.dtype = dtype self.reduce_mean = ms.ops.ReduceMean() def forward(self, query, key, value, key_padding_mask=None, need_weights: bool=True, attn_mask=None, average_attn_weights: bool=True): unsupported_attr(key_padding_mask) unsupported_attr(average_attn_weights) if need_weights is True: raise ValueError("Until now, `need_weights`='True' is not supported") query = self._batch_tensor(query, 'query') key = self._batch_tensor(key, 'key') value = self._batch_tensor(value, 'value') _batch_size = query.shape[0] _src_seq_length = query.shape[1] _tgt_seq_length = key.shape[1] if attn_mask: _attn_mask = self._process_mask(attn_mask, _batch_size) else: _attn_mask = ms.ops.ones((_batch_size, _src_seq_length, _tgt_seq_length), mstype.float32) self.ms_multihead_attention = ms.nn.transformer.MultiHeadAttention( batch_size = _batch_size, src_seq_length = _src_seq_length, tgt_seq_length = _tgt_seq_length, hidden_size=self.embed_dim, num_heads=self.num_heads, hidden_dropout_rate=self.dropout, attention_dropout_rate=self.dropout, compute_dtype=mstype.float32, softmax_compute_type=mstype.float32, param_init_type=mstype.float32, use_past=False) out, attn_output_weights = self.ms_multihead_attention(query, key, value, _attn_mask) if not self.batch_first: # ms default is (batch, seq, feature), batch_first out = ms.ops.transpose(out, (1, 0, 2)) # if need_weights: # if average_attn_weights: # attn_output_weights = self.reduce_mean(attn_output_weights, 1) # if _batch_size == 1: # attn_output_weights = self.reduce_mean(attn_output_weights, 0) # else: # attn_output_weights = None if _batch_size == 1: out = self.reduce_mean(out, 0) # TODO # Until Now, attn_output_weights is not the same as pytorch attn_output_weights = None return cast_to_adapter_tensor(out), cast_to_adapter_tensor(attn_output_weights) def _batch_tensor(self, x, x_name: str): x = cast_to_ms_tensor(x) _rank = ms.ops.rank(x) if _rank == 2: out = ms.ops.expand_dims(x, 0) return out if _rank == 3: if not self.batch_first: out = ms.ops.transpose(x, (1, 0 ,2)) else: out = x return out raise ValueError(f"For MultiheadAttention, rank of {x_name} should be 2 or 3, but got {_rank}") def _process_mask(self, mask, batch_size): mask = cast_to_ms_tensor(mask) _rank = ms.ops.rank(mask) if _rank == 2: out = ms.ops.expand_dims(mask, 0) return out if _rank == 3: if mask.shape[0] != batch_size: warnings.warn("Until now, `attn_mask` can only support shape (N, L, S)" "when `attn_mask` shape is (N * num_heads, L, S), pick the first (N, L, S) mask") mask = mask[:batch_size,:] return mask raise ValueError(f"For MultiheadAttention, rank of mask should be 2 or 3, but got {_rank}") class PReLU(Module): def __init__(self, num_parameters=1, init=0.25, device=None, dtype=None): super(PReLU, self).__init__() unsupported_attr(device) validator.check_positive_int(num_parameters, 'num_parameters', self.cls_name) if dtype is None: dtype = mstype.float32 w = init if isinstance(w, (float, np.float32)): tmp = np.empty((num_parameters,), dtype=np.float32) tmp.fill(w) w = tensor(tmp, dtype=dtype) elif isinstance(w, list): if len(w) != num_parameters: raise ValueError(f"For '{self.cls_name}', the length of 'init' must be equal to the 'num_parameters'" f"when the 'init' is a list, but got the length of 'num_parameters': {len(w)}, " f"the 'num_parameters': {num_parameters}.") for i in w: if not isinstance(i, (float, np.float32)): raise ValueError(f"For '{self.cls_name}', all elements in 'init' must be " f"float when the 'init' is a list, but got {i}.") w = tensor(w, dtype=dtype) elif isinstance(w, Tensor): if w.dtype not in (mstype.float16, mstype.float32): raise ValueError(f"For '{self.cls_name}', the dtype of 'init' must be float16 or " f"float32 when the 'init' is a tensor, but got {w.dtype}.") if len(w.shape) != 1 or w.shape[0] != num_parameters: raise ValueError(f"For '{self.cls_name}', the dimension of 'init' must be 1, and the elements number " f"should be equal to the 'num_parameters' when the 'init' is a tensor, " f"but got 'init' shape {w.shape}, the 'num_parameters' {num_parameters}.") else: raise TypeError(f"For '{self.cls_name}', the 'init' only supported float, list and tensor, " f"but got {type(w).__name__}.") self.weight = Parameter(w) self.num_parameters = num_parameters def forward(self, input): return ms_torch_nn_func.prelu(input, self.weight) def extra_repr(self) -> str: return 'num_parameters={}'.format(self.num_parameters) class Softplus(Module): def __init__(self, beta=1, threshold=20): super(Softplus, self).__init__() self.beta = beta self.threshold = threshold def forward(self, input): return ms_torch_nn_func.softplus(input, self.beta, self.threshold) def extra_repr(self): return 'beta={}, threshold={}'.format(self.beta, self.threshold) class Softmax2d(Module): def __init__(self): super(Softmax2d, self).__init__() self.softmax2d = ms.nn.Softmax2d() def forward(self, input): if input.dim() not in {3, 4}: raise RuntimeError("Softmax2d requires a 3D or 4D tensor as input") return self.softmax2d(input) msadapter/ms_adapter/pytorch/nn/modules/adaptive.py000066400000000000000000000171231436713554100231140ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from collections import namedtuple import mindspore as ms from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from ms_adapter.utils import unsupported_attr from .container import Sequential, ModuleList from .linear import Linear from .module import Module from ..functional import log_softmax _ASMoutput = namedtuple('_ASMoutput', ['output', 'loss']) class AdaptiveLogSoftmaxWithLoss(Module): def __init__(self, in_features, n_classes, cutoffs, div_value=4., head_bias=False, device=None, dtype=None): super(AdaptiveLogSoftmaxWithLoss, self).__init__() unsupported_attr(device) cutoffs = list(cutoffs) # #TODO: pylint # if (cutoffs != sorted(cutoffs)) \ # or (min(cutoffs) <= 0) \ # or (max(cutoffs) > (n_classes - 1)) \ # or (len(set(cutoffs)) != len(cutoffs)) \ # or any([int(c) != c for c in cutoffs]): # # raise ValueError("cutoffs should be a sequence of unique, positive " # "integers sorted in an increasing order, where " # "each value is between 1 and n_classes-1") self.in_features = in_features self.n_classes = n_classes self.cutoffs = cutoffs + [n_classes] self.div_value = div_value self.head_bias = head_bias self.dtype = dtype self.shortlist_size = self.cutoffs[0] self.n_clusters = len(self.cutoffs) - 1 self.head_size = self.shortlist_size + self.n_clusters self.head = Linear(self.in_features, self.head_size, bias=self.head_bias, dtype=self.dtype) self.tail = ModuleList() for i in range(self.n_clusters): hsz = int(self.in_features // (self.div_value ** (i + 1))) osz = self.cutoffs[i + 1] - self.cutoffs[i] projection = Sequential( Linear(self.in_features, hsz, bias=False, dtype=self.dtype), Linear(hsz, osz, bias=False, dtype=self.dtype), ) self.tail.append(projection) def reset_parameters(self): self.head.reset_parameters() for i2h, h2o in self.tail: i2h.reset_parameters() h2o.reset_parameters() def forward(self, input_, target_): input_ = cast_to_ms_tensor(input_) #target_ = cast_to_ms_tensor(target_) targ_dim = target_.dim() if targ_dim == 1: if input_.shape[0] != target_.shape[0]: raise RuntimeError('Input and target should have the same size ' 'in the batch dimension.') if input_.dim() != 2: raise RuntimeError('1D target tensor expects 2D input tensors, ' 'but found inputs with size', input_.shape()) elif targ_dim == 0: if input_.dim() != 1: raise RuntimeError('0D target tensor expects 1D input tensors, ' 'but found inputs with size', input_.shape()) else: raise RuntimeError('0D or 1D target tensor expected, ' 'multi-target not supported') is_batched = targ_dim > 0 input = input_ if is_batched else input_.unsqueeze(0) target = target_ if is_batched else target_.unsqueeze(0) used_rows = 0 batch_size = target.shape[0] output = input.new_zeros(batch_size) #gather_inds = ms.numpy.empty(batch_size, target.dtype) gather_inds = target.new_empty(batch_size) cutoff_values = [0] + self.cutoffs for i in range(len(cutoff_values) - 1): low_idx = cutoff_values[i] high_idx = cutoff_values[i + 1] target_mask = (target >= low_idx) & (target < high_idx) row_indices = target_mask.nonzero().squeeze() if row_indices.numel() == 0: continue if i == 0: #gather_inds.index_copy_(0, row_indices, target[target_mask]) gather_inds = index_copy_0dim(gather_inds, row_indices, target[target_mask]) else: relative_target = target[target_mask] - low_idx #input_subset = input.index_select(0, row_indices) input_subset = ms.ops.gather(input, row_indices, 0) cluster_output = self.tail[i - 1](input_subset) cluster_index = self.shortlist_size + i - 1 gather_inds.index_fill_(0, row_indices, cluster_index) cluster_logprob = log_softmax(cluster_output, dim=1) local_logprob = cluster_logprob.gather(1, relative_target.unsqueeze(1)) #output.index_copy_(0, row_indices, local_logprob.squeeze(1)) output = index_copy_0dim(output, row_indices, local_logprob.squeeze(1)) used_rows += row_indices.numel() if used_rows != batch_size: raise RuntimeError("Target values should be in [0, {}], " "but values in range [{}, {}] " "were found. ".format(self.n_classes - 1, target.min().item(), target.max().item())) head_output = self.head(input) head_logprob = log_softmax(head_output, dim=1) output += head_logprob.gather(1, gather_inds.unsqueeze(1)).squeeze() loss = (-output).mean() if not is_batched: output = output.squeeze(0) output = cast_to_adapter_tensor(output) loss = cast_to_adapter_tensor(loss) return _ASMoutput(output, loss) def _get_full_log_prob(self, input, head_output): input = cast_to_ms_tensor(input) head_output = cast_to_ms_tensor(head_output) out = input.new_empty((head_output.shape[0], self.n_classes)) head_logprob = log_softmax(head_output, dim=1) out[:, :self.shortlist_size] = head_logprob[:, :self.shortlist_size] for i, (start_idx, stop_idx) in enumerate(zip(self.cutoffs, self.cutoffs[1:])): cluster_output = self.tail[i](input) cluster_logprob = log_softmax(cluster_output, dim=1) output_logprob = cluster_logprob + head_logprob[:, self.shortlist_size + i].unsqueeze(1) out[:, start_idx:stop_idx] = output_logprob return cast_to_adapter_tensor(out) def log_prob(self, input): input = cast_to_ms_tensor(input) head_output = self.head(input) out = self._get_full_log_prob(input, head_output) return cast_to_adapter_tensor(out) def predict(self, input): input = cast_to_ms_tensor(input) head_output = self.head(input) cast_to_adapter_tensor() output = ms.ops.argmax(head_output, axis=1) not_in_shortlist = (output >= self.shortlist_size) any_in_shortlist = (output < self.shortlist_size) if not not_in_shortlist: return cast_to_adapter_tensor(output) elif not any_in_shortlist: log_prob = self._get_full_log_prob(input, head_output) return cast_to_adapter_tensor(ms.ops.argmax(log_prob, axis=1)) else: log_prob = self._get_full_log_prob(input[not_in_shortlist], head_output[not_in_shortlist]) output[not_in_shortlist] = ms.ops.argmax(log_prob, axis=1) return cast_to_adapter_tensor(output) def index_copy_0dim(input, index, tensor): for i in range(len(index)): input[index[i]] = tensor[i] return input msadapter/ms_adapter/pytorch/nn/modules/batchnorm.py000066400000000000000000000401441436713554100232730ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import itertools import mindspore.ops as P from mindspore.ops.operations import _inner_ops as inner from mindspore.communication.management import get_group_size, get_rank from mindspore._checkparam import Validator as validator from mindspore._checkparam import Rel from mindspore.communication import management import mindspore.context as context from ms_adapter.pytorch.nn import init from ms_adapter.pytorch.functional import empty from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from .module import Module __all__ = ['BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'LazyBatchNorm1d', 'LazyBatchNorm2d', 'LazyBatchNorm3d', 'SyncBatchNorm', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d', 'LazyInstanceNorm1d', 'LazyInstanceNorm2d', 'LazyInstanceNorm3d'] class _NormBase(Module): """Common base of _InstanceNorm and _BatchNorm""" def __init__( self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, device=None, dtype=None ): unsupported_attr(device) unsupported_attr(dtype) super(_NormBase, self).__init__() self.num_features = num_features self.eps = eps self.momentum = momentum self.affine = affine self.track_running_stats = track_running_stats self.weight = Parameter(empty(num_features), requires_grad=affine) self.bias = Parameter(empty(num_features), requires_grad=affine) self.running_mean = Parameter(empty(num_features), requires_grad=False) self.running_var = Parameter(empty(num_features), requires_grad=False) self.reset_parameters() if not self.track_running_stats: self.momentum = 0.0 def reset_running_stats(self): self.running_mean.zero_() self.running_var.fill_(1) def reset_parameters(self): self.reset_running_stats() init.ones_(self.weight) init.zeros_(self.bias) def extra_repr(self): return ( "{num_features}, eps={eps}, momentum={momentum}, affine={affine}, " "track_running_stats={track_running_stats}".format(**self.__dict__) ) class _BatchNorm(_NormBase): """Common base of BatchNorm""" def __init__( self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, device=None, dtype=None ): factory_kwargs = {'device': device, 'dtype': dtype} super(_BatchNorm, self).__init__(num_features, eps, momentum, affine, track_running_stats, **factory_kwargs) self.bn_train = P.BatchNorm(is_training=True, epsilon=self.eps, momentum=self.momentum, data_format='NCHW') self.bn_infer = P.BatchNorm(is_training=False, epsilon=self.eps, data_format='NCHW') def _check_input_dim(self, input): raise NotImplementedError def _check_rank_ids(self, process_groups, rank_size): seen = set() for rid in itertools.chain(*process_groups): validator.check_int_range(rid, 0, rank_size, Rel.INC_LEFT, "rank id in process_groups", self.cls_name) if rid in seen: raise ValueError(f"For '{self.cls_name}', rank id in 'process_groups' must not be duplicated, " f"but got {process_groups}.") seen.add(rid) def _create_sync_groups(self): for i in range(len(self.process_groups)): validator.check_isinstance("process_groups[%d]" % i, self.process_groups[i], list) self.group_device_num = len(self.process_groups[i]) if self.rank_id in self.process_groups[i] and self.group_device_num > 1: self.is_global = True global SYNC_BN_GROUP_NAME if SYNC_BN_GROUP_NAME == "": SYNC_BN_GROUP_NAME = "sync_bn_group%d" % i management.create_group(SYNC_BN_GROUP_NAME, self.process_groups[i]) def forward(self, input): self._check_input_dim(input) input = cast_to_ms_tensor(input) # TODO cast Parameter if self.training or (not self.training and not self.track_running_stats): output = self.bn_train(input, self.weight, self.bias, self.running_mean, self.running_var)[0] else: output = self.bn_infer(input, self.weight, self.bias, self.running_mean, self.running_var)[0] return cast_to_adapter_tensor(output) class BatchNorm1d(_BatchNorm): def _check_input_dim(self, input): if len(input.shape) not in {2, 3}: raise ValueError( "expected 2D or 3D input (got {}D input)".format(input.dim()) ) return True def forward(self, input): self._check_input_dim(input) input = cast_to_ms_tensor(input) shape = input.shape if len(shape) == 3: input = input.transpose(0, 2, 1).reshape(-1, shape[1]) # TODO cast Parameter if self.training or (not self.training and not self.track_running_stats): output = self.bn_train(input, self.weight, self.bias, self.running_mean, self.running_var)[0] else: output = self.bn_infer(input, self.weight, self.bias, self.running_mean, self.running_var)[0] output = output.reshape(shape[0], shape[2], shape[1]).transpose(0, 2, 1) else: if self.training or (not self.training and not self.track_running_stats): output = self.bn_train(input, self.weight, self.bias, self.running_mean, self.running_var)[0] else: output = self.bn_infer(input, self.weight, self.bias, self.running_mean, self.running_var)[0] return cast_to_adapter_tensor(output) class BatchNorm2d(_BatchNorm): r"""Applies Batch Normalization over a 4D input (a mini-batch of 2D inputs with additional channel dimension) as described in the paper `Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift `__ . .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta The mean and standard-deviation are calculated per-dimension over the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0. The standard-deviation is calculated via the biased estimator, equivalent to `ms_adapter.pytorch.var(input, unbiased=False)`. Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default :attr:`momentum` of 0.1. If :attr:`track_running_stats` is set to ``False``, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well. .. note:: This :attr:`momentum` argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`, where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value. Because the Batch Normalization is done over the `C` dimension, computing statistics on `(N, H, W)` slices, it's common terminology to call this Spatial Batch Normalization. Args: num_features: :math:`C` from an expected input of size :math:`(N, C, H, W)` eps: a value added to the denominator for numerical stability. Default: 1e-5 momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` Shape: - Input: :math:`(N, C, H, W)` - Output: :math:`(N, C, H, W)` (same shape as input) Examples:: >>> # With Learnable Parameters >>> m = nn.BatchNorm2d(100) >>> # Without Learnable Parameters >>> m = nn.BatchNorm2d(100, affine=False) >>> input = ms_adapter.pytorch.randn(20, 100, 35, 45) >>> output = m(input) """ def _check_input_dim(self, input): if len(input.shape) != 4: raise ValueError("expected 4D input (got {}D input)".format(input.dim())) return True class BatchNorm3d(_BatchNorm): def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, device=None, dtype=None): factory_kwargs = {'device': device, 'dtype': dtype} super(BatchNorm3d, self).__init__(num_features, eps, momentum, affine, track_running_stats, **factory_kwargs) self.reshape = P.Reshape() def forward(self, input): self._check_input_dim(input) input = cast_to_ms_tensor(input) x_shape = input.shape input = self.reshape(input, (x_shape[0], x_shape[1], x_shape[2] * x_shape[3], x_shape[4])) if self.training or (not self.training and not self.track_running_stats): bn2d_out = self.bn_train(input, self.weight, self.bias, self.running_mean, self.running_var)[0] else: bn2d_out = self.bn_infer(input, self.weight, self.bias, self.running_mean, self.running_var)[0] bn3d_out = self.reshape(bn2d_out, x_shape) return cast_to_adapter_tensor(bn3d_out) def _check_input_dim(self, input): if len(input.shape) != 5: raise ValueError("expected 5D input (got {}D input)".format(input.dim())) return True class SyncBatchNorm(_BatchNorm): def __init__( self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, device=None, dtype=None ): factory_kwargs = {'device': device, 'dtype': dtype} super(SyncBatchNorm, self).__init__( num_features, eps, momentum, affine, track_running_stats, **factory_kwargs ) self.process_groups = process_group global SYNC_BN_GROUP_NAME if self.process_groups != 0: self.rank_id = get_rank() self.rank_size = get_group_size() if self.process_groups is not None: validator.check_isinstance("process_groups", self.process_groups, list) self._check_rank_ids(self.process_groups, self.rank_size) self._create_sync_groups() elif self.rank_size > 1: self.is_global = True self.group_device_num = self.rank_size self.device_list = list(range(0, self.rank_size)) if context.get_context("device_target") == "Ascend": if SYNC_BN_GROUP_NAME == "": SYNC_BN_GROUP_NAME = "sync_bn_group0" management.create_group(SYNC_BN_GROUP_NAME, self.device_list) elif context.get_context("device_target") == "GPU": if SYNC_BN_GROUP_NAME == "": SYNC_BN_GROUP_NAME = "nccl_world_group" self.bn_train = inner.SyncBatchNorm(epsilon=self.eps, momentum=self.momentum, group=SYNC_BN_GROUP_NAME, device_num=self.group_device_num) def _check_input_dim(self, input): if len(input.shape) < 4: raise ValueError( "expected at least 2D input (got {}D input)".format(input.dim()) ) return True class _InstanceNorm(_NormBase): def __init__( self, num_features, eps=1e-5, momentum=0.1, affine=False, track_running_stats=False, device=None, dtype=None ): factory_kwargs = {'device': device, 'dtype': dtype} super(_InstanceNorm, self).__init__( num_features, eps, momentum, affine, track_running_stats, **factory_kwargs) self.instance_bn = P.operations.InstanceNorm(epsilon=self.eps, momentum=self.momentum) def _check_input_dim(self, input): raise NotImplementedError def forward(self, input): self._check_input_dim(input) input = cast_to_ms_tensor(input) output = self.instance_bn(input, self.weight, self.bias, self.running_mean, self.running_var)[0] output = cast_to_adapter_tensor(output) return output class InstanceNorm1d(_InstanceNorm): def _check_input_dim(self, input): if input.dim() == 2: raise ValueError( 'InstanceNorm1d returns 0-filled tensor to 2D tensor.' 'This is because InstanceNorm1d reshapes inputs to' '(1, N * C, ...) from (N, C,...) and this makes' 'variances 0.' ) if input.dim() != 3: raise ValueError('expected 3D input (got {}D input)' .format(input.dim())) return True class InstanceNorm2d(_InstanceNorm): def _check_input_dim(self, input): if input.dim() != 4: raise ValueError('expected 4D input (got {}D input)' .format(input.dim())) return True class InstanceNorm3d(_InstanceNorm): def _check_input_dim(self, input): if input.dim() != 5: raise ValueError('expected 5D input (got {}D input)' .format(input.dim())) return True LazyBatchNorm1d = BatchNorm1d LazyBatchNorm2d = BatchNorm2d LazyBatchNorm3d = BatchNorm3d LazyInstanceNorm1d = InstanceNorm1d LazyInstanceNorm2d = InstanceNorm2d LazyInstanceNorm3d = InstanceNorm3d msadapter/ms_adapter/pytorch/nn/modules/channelshuffle.py000066400000000000000000000012031436713554100242740ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore.nn as nn from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from .module import Module __all__ = ['ChannelShuffle'] class ChannelShuffle(Module): def __init__(self, groups): super(ChannelShuffle, self).__init__() self.groups = groups self.channel_shuffle = nn.ChannelShuffle(self.groups) def forward(self, input): input = cast_to_ms_tensor(input) out = self.channel_shuffle(input) return cast_to_adapter_tensor(out) def extra_repr(self): return 'groups={}'.format(self.groups) msadapter/ms_adapter/pytorch/nn/modules/container.py000066400000000000000000001016511436713554100233010ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from abc import abstractmethod import operator from typing import Any, Dict, Iterable, Iterator, Mapping, Optional, Tuple, Union from collections import OrderedDict, abc as container_abcs from mindspore.nn.layer.container import _get_prefix_and_index, _valid_index, _valid_cell from ms_adapter.pytorch.tensor import Tensor, cast_to_adapter_tensor from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.pytorch._ref import typename from .module import Module class Sequential(Module): """ Sequential Module container. For more details about Module, please refer to A list of Cells will be added to it in the order they are passed in the constructor. Alternatively, an ordered dict of cells can also be passed in. Note: Sequential and nn.ModuleList are different, ModuleList is a list for storing modules. However, the layers in a Sequential are connected in a cascading way. Args: args (list, OrderedDict): List or OrderedDict of subclass of Module. Inputs: - **x** (Tensor) - Tensor with shape according to the first Module in the sequence. Outputs: Tensor, the output Tensor with shape depending on the input `x` and defined sequence of Cells. Raises: TypeError: If the type of the `args` is not list or OrderedDict. Supported Platforms: ``Ascend`` ``GPU`` ``CPU`` Examples: >>> conv = nn.Conv2d(3, 2, 3, pad_mode='valid', weight_init="ones") >>> relu = nn.ReLU() >>> seq = nn.Sequential([conv, relu]) >>> x = Tensor(np.ones([1, 3, 4, 4]), dtype=mindspore.float32) >>> output = seq(x) >>> print(output) [[[[27. 27.] [27. 27.]] [[27. 27.] [27. 27.]]]] >>> from collections import OrderedDict >>> d = OrderedDict() >>> d["conv"] = conv >>> d["relu"] = relu >>> seq = nn.Sequential(d) >>> x = Tensor(np.ones([1, 3, 4, 4]), dtype=mindspore.float32) >>> output = seq(x) >>> print(output) [[[[27. 27.] [27. 27.]] [[27. 27.] [27. 27.]]]] """ def __init__(self, *args): """Initialize Sequential.""" super(Sequential, self).__init__() self._is_dynamic_name = [] if len(args) == 1: cells = args[0] if isinstance(cells, list): for index, cell in enumerate(cells): self.insert_child_to_cell(str(index), cell) cell.update_parameters_name(str(index) + ".") self._is_dynamic_name.append(True) elif isinstance(cells, OrderedDict): for name, cell in cells.items(): self.insert_child_to_cell(name, cell) cell.update_parameters_name(name + ".") self._is_dynamic_name.append(False) elif isinstance(cells, Module): for index, cell in enumerate(args): self.insert_child_to_cell(str(index), cell) cell.update_parameters_name(str(index) + ".") self._is_dynamic_name.append(True) else: raise TypeError(f"For '{self.__class__.__name__}', the 'args[0]' must be list or orderedDict, " f"but got {type(cells).__name__}") else: for index, cell in enumerate(args): self.insert_child_to_cell(str(index), cell) cell.update_parameters_name(str(index) + ".") self._is_dynamic_name.append(True) self.cell_list = list(self._cells.values()) def __getitem__(self, index): if isinstance(index, slice): return self.__class__( OrderedDict(list(self._cells.items())[index])) index = _valid_index(len(self), index, self.__class__.__name__) return list(self._cells.values())[index] def __setitem__(self, index, module): cls_name = self.__class__.__name__ if _valid_cell(module, cls_name): prefix, _ = _get_prefix_and_index(self._cells) index = _valid_index(len(self), index, cls_name) key = list(self._cells.keys())[index] self._cells[key] = module module.update_parameters_name(prefix + key + ".") self.cell_list = list(self._cells.values()) def __delitem__(self, index): cls_name = self.__class__.__name__ if isinstance(index, int): index = _valid_index(len(self), index, cls_name) key = list(self._cells.keys())[index] del self._cells[key] del self._is_dynamic_name[index] elif isinstance(index, slice): keys = list(self._cells.keys())[index] for key in keys: del self._cells[key] del self._is_dynamic_name[index] else: raise TypeError(f"For '{cls_name}', the type of index must be int type or slice type, " f"but got {type(index).__name__}") prefix, key_index = _get_prefix_and_index(self._cells) temp_dict = OrderedDict() for idx, key in enumerate(self._cells.keys()): cell = self._cells[key] if self._is_dynamic_name[idx]: for _, param in cell.parameters_and_names(): param.name = prefix + str(idx) + "." + ".".join(param.name.split(".")[key_index+1:]) temp_dict[str(idx)] = cell else: temp_dict[key] = cell self._cells = temp_dict self.cell_list = list(self._cells.values()) def __len__(self): return len(self._cells) @property def _modules(self): return self._cells def set_grad(self, flag=True): self.requires_grad = flag for cell in self._cells.values(): cell.set_grad(flag) def append(self, module): """ Appends a given Module to the end of the list. Args: module(Module): The Module to be appended. Examples: >>> conv = nn.Conv2d(3, 2, 3, pad_mode='valid', weight_init="ones") >>> bn = nn.BatchNorm2d(2) >>> relu = nn.ReLU() >>> seq = nn.Sequential([conv, bn]) >>> seq.append(relu) >>> x = Tensor(np.ones([1, 3, 4, 4]), dtype=mindspore.float32) >>> output = seq(x) >>> print(output) [[[[26.999863 26.999863] [26.999863 26.999863]] [[26.999863 26.999863] [26.999863 26.999863]]]] """ if _valid_cell(module, self.__class__.__name__): prefix, _ = _get_prefix_and_index(self._cells) module.update_parameters_name(prefix + str(len(self)) + ".") self._is_dynamic_name.append(True) self._cells[str(len(self))] = module self.cell_list = list(self._cells.values()) def add_module(self, name, module): if not isinstance(module, Module) and module is not None: raise TypeError("{} is not a Module subclass".format( module.__name__)) elif hasattr(self, name) and name not in self._cells: raise KeyError("attribute '{}' already exists".format(name)) elif '.' in name: raise KeyError("module name can't contain \".\", got: {}".format(name)) elif name == '': raise KeyError("module name can't be empty string \"\"") if _valid_cell(module, self.__class__.__name__): module.update_parameters_name(name + ".") self._is_dynamic_name.append(False) self._cells[name] = module self.cell_list = list(self._cells.values()) def forward(self, input): for cell in self.cell_list: input = cell(input) return cast_to_adapter_tensor(input) class _ModuleListBase: """ An interface for base the Module as list. The sequential Module may be iterated using the construct method using for-in statement. But there are some scenarios that the construct method built-in does not fit. For convenience, we provide an interface that indicates the sequential Module may be interpreted as list of Cells, so it can be accessed using iterator or subscript when a sequential Module instantiate is accessed by iterator or subscript, it will be interpreted as a list of Cells. """ def __init__(self): """Initialize _ModuleListBase.""" self.__cell_as_list__ = True @abstractmethod def __len__(self): pass @abstractmethod def __getitem__(self, index): pass def construct(self): raise NotImplementedError class ModuleList(_ModuleListBase, Module): """ Holds Cells in a list. For more details about Module, please refer to ModuleList can be used like a regular Python list, the Cells it contains have been initialized. Args: args (list, optional): List of subclass of Module. Supported Platforms: ``Ascend`` ``GPU`` ``CPU`` Examples: >>> import mindspore.nn as nn >>> >>> conv = nn.Conv2d(100, 20, 3) >>> bn = nn.BatchNorm2d(20) >>> relu = nn.ReLU() >>> cell_ls = nn.ModuleList([bn]) >>> cell_ls.insert(0, conv) >>> cell_ls.append(relu) >>> cell_ls.extend([relu, relu]) """ def __init__(self, *args, **kwargs): """Initialize ModuleList.""" auto_prefix = kwargs["auto_prefix"] if "auto_prefix" in kwargs.keys() else True _ModuleListBase.__init__(self) Module.__init__(self, auto_prefix) if len(args) == 1: self.extend(args[0]) def __getitem__(self, index): cls_name = self.__class__.__name__ if isinstance(index, slice): return self.__class__(list(self._cells.values())[index]) if isinstance(index, int): index = _valid_index(len(self), index, cls_name) return self._cells[str(index)] raise TypeError(f"For '{cls_name}', the type of 'index' must be int or slice, " f"but got {type(index).__name__}.") def __setitem__(self, index, module): cls_name = self.__class__.__name__ if not isinstance(index, int) and _valid_cell(module, cls_name): raise TypeError(f"For '{cls_name}', the type of 'index' must be int, " f"but got {type(index).__name__}.") index = _valid_index(len(self), index, cls_name) if self._auto_prefix: prefix, _ = _get_prefix_and_index(self._cells) module.update_parameters_name(prefix + str(index) + ".") self._cells[str(index)] = module def __delitem__(self, index): cls_name = self.__class__.__name__ if isinstance(index, int): index = _valid_index(len(self), index, cls_name) del self._cells[str(index)] elif isinstance(index, slice): keys = list(self._cells.keys())[index] for key in keys: del self._cells[key] else: raise TypeError(f"For '{cls_name}', the type of 'index' must be int or slice, " f"but got {type(index).__name__}.") # adjust orderedDict prefix, key_index = _get_prefix_and_index(self._cells) temp_dict = OrderedDict() for idx, cell in enumerate(self._cells.values()): if self._auto_prefix: for _, param in cell.parameters_and_names(): param.name = prefix + str(idx) + "." + ".".join(param.name.split(".")[key_index+1:]) temp_dict[str(idx)] = cell self._cells = temp_dict def __len__(self): return len(self._cells) def __iter__(self): return iter(self._cells.values()) def __iadd__(self, modules): self.extend(modules) return self def insert(self, index, module): """ Inserts a given Module before a given index in the list. Args: index(int): The Insert index in the ModuleList. module(Module): The Module to be inserted. """ cls_name = self.__class__.__name__ idx = _valid_index(len(self), index, cls_name) _valid_cell(module, cls_name) length = len(self) prefix, key_index = _get_prefix_and_index(self._cells) while length > idx: if self._auto_prefix: tmp_cell = self._cells[str(length-1)] for _, param in tmp_cell.parameters_and_names(): param.name = prefix + str(length) + "." + ".".join(param.name.split(".")[key_index+1:]) self._cells[str(length)] = self._cells[str(length - 1)] length -= 1 self._cells[str(idx)] = module if self._auto_prefix: module.update_parameters_name(prefix + str(idx) + ".") def extend(self, modules): """ Appends Cells from a Python iterable to the end of the list. Args: cells(list): The Cells to be extended. Raises: TypeError: If the argument cells are not a list of Cells. """ cls_name = self.__class__.__name__ if not isinstance(modules, container_abcs.Iterable): raise TypeError("ModuleList.extend should be called with an " "iterable, but got " + type(modules).__name__) prefix, _ = _get_prefix_and_index(self._cells) for module in modules: if _valid_cell(module, cls_name): if self._auto_prefix: module.update_parameters_name(prefix + str(len(self)) + ".") self._cells[str(len(self))] = module return self def append(self, module): """ Appends a given Module to the end of the list. Args: module(Module): The subcell to be appended. """ if _valid_cell(module, self.__class__.__name__): if self._auto_prefix: prefix, _ = _get_prefix_and_index(self._cells) module.update_parameters_name(prefix + str(len(self)) + ".") self._cells[str(len(self))] = module def set_grad(self, flag=True): self.requires_grad = flag for cell in self._cells.values(): cell.set_grad(flag) def construct(self, *inputs): raise NotImplementedError class ModuleDict(Module): r"""Holds submodules in a dictionary. :class:`nn.ModuleDict` can be indexed like a regular Python dictionary, but modules it contains are properly registered, and will be visible by all :class:`nn.Module` methods. :class:`nn.ModuleDict` is an **ordered** dictionary that respects * the order of insertion, and * in :meth:`nn.ModuleDict.update`, the order of the merged ``OrderedDict``, ``dict`` (started from Python 3.6) or another :class:`nn.ModuleDict` (the argument to :meth:`nn.ModuleDict.update`). Note that :meth:`nn.ModuleDict.update` with other unordered mapping types (e.g., Python's plain ``dict`` before Python version 3.6) does not preserve the order of the merged mapping. Args: modules (iterable, optional): a mapping (dictionary) of (string: module) or an iterable of key-value pairs of type (string, module) Example:: class MyModule(nn.Module): def __init__(self): super(MyModule, self).__init__() self.choices = nn.ModuleDict({ 'conv': nn.Conv2d(10, 10, 3), 'pool': nn.MaxPool2d(3) }) self.activations = nn.ModuleDict([ ['lrelu', nn.LeakyReLU()], ['prelu', nn.PReLU()] ]) def forward(self, x, choice, act): x = self.choices[choice](x) x = self.activations[act](x) return x """ def __init__(self, modules: Optional[Mapping[str, Module]] = None) -> None: super(ModuleDict, self).__init__() if modules is not None: self.update(modules) def __getitem__(self, key: str) -> Module: return self._cells[key] def __setitem__(self, key: str, module: Module) -> None: self.add_module(key, module) def __delitem__(self, key: str) -> None: del self._cells[key] def __len__(self) -> int: return len(self._cells) def __iter__(self) -> Iterator[str]: return iter(self._cells) def __contains__(self, key: str) -> bool: return key in self._cells def clear(self) -> None: """Remove all items from the ModuleDict. """ self._cells.clear() def pop(self, key: str) -> Module: r"""Remove key from the ModuleDict and return its module. Args: key (str): key to pop from the ModuleDict """ v = self[key] del self[key] return v def keys(self) -> Iterable[str]: r"""Return an iterable of the ModuleDict keys. """ return self._cells.keys() def items(self) -> Iterable[Tuple[str, Module]]: r"""Return an iterable of the ModuleDict key/value pairs. """ return self._cells.items() def values(self) -> Iterable[Module]: r"""Return an iterable of the ModuleDict values. """ return self._cells.values() def update(self, modules: Mapping[str, Module]) -> None: r"""Update the :class:`nn.ModuleDict` with the key-value pairs from a mapping or an iterable, overwriting existing keys. .. note:: If :attr:`modules` is an ``OrderedDict``, a :class:`nn.ModuleDict`, or an iterable of key-value pairs, the order of new elements in it is preserved. Args: modules (iterable): a mapping (dictionary) from string to :class:`nn.Module`, or an iterable of key-value pairs of type (string, :class:`nn.Module`) """ if not isinstance(modules, container_abcs.Iterable): raise TypeError("ModuleDict.update should be called with an " "iterable of key/value pairs, but got " + type(modules).__name__) if isinstance(modules, (OrderedDict, ModuleDict, container_abcs.Mapping)): for key, module in modules.items(): self[key] = module else: # modules here can be a list with two items for j, m in enumerate(modules): if not isinstance(m, container_abcs.Iterable): raise TypeError("ModuleDict update sequence element " "#" + str(j) + " should be Iterable; is" + type(m).__name__) if not len(m) == 2: raise ValueError("ModuleDict update sequence element " "#" + str(j) + " has length " + str(len(m)) + "; 2 is required") # modules can be Mapping (what it's typed at), or a list: [(name1, module1), (name2, module2)] # that's too cumbersome to type correctly with overloads, so we add an ignore here self[m[0]] = m[1] # type: ignore[assignment] # remove forward alltogether to fallback on Module's _forward_unimplemented class ParameterList(Module): """Holds parameters in a list. :class:`nn.ParameterList` can be used like a regular Python list, but Tensors that are :class:`nn.Parameter` are properly registered, and will be visible by all :class:`nn.Module` methods. Note that the constructor, assigning an element of the list, the :meth:`nn.ParameterDict.append` method and the :meth:`nn.ParameterDict.extend` method will convert any :class:`Tensor` into :class:`nn.Parameter`. Args: parameters (iterable, optional): an iterable of elements to add to the list. Example:: class MyModule(nn.Module): def __init__(self): super(MyModule, self).__init__() self.params = nn.ParameterList([nn.Parameter(ms_torch.randn(10, 10)) for i in range(10)]) def forward(self, x): # ParameterList can act as an iterable, or be indexed using ints for i, p in enumerate(self.params): x = self.params[i // 2].mm(x) + p.mm(x) return x """ def __init__(self, values): super(ParameterList, self).__init__() self._size = 0 if values is not None: self += values def _get_abs_string_index(self, idx): """Get the absolute index for the list of modules""" idx = operator.index(idx) if not -len(self) <= idx < len(self): raise IndexError('index {} is out of range'.format(idx)) if idx < 0: idx += len(self) return str(idx) def __getitem__(self, idx): if isinstance(idx, slice): start, stop, step = idx.indices(len(self)) out = self.__class__() for i in range(start, stop, step): out.append(self[i]) return out else: idx = self._get_abs_string_index(idx) return getattr(self, str(idx)) def __setitem__(self, idx: int, param: Any) -> None: # Note that all other function that add an entry to the list part of # the ParameterList end up here. So this is the only place where we need # to wrap things into Parameter if needed. # Objects added via setattr() are not in the list part and thus won't # call into this function. idx = self._get_abs_string_index(idx) if isinstance(param, Tensor) and not isinstance(param, Parameter): param = Parameter(param) return setattr(self, str(idx), param) def __len__(self) -> int: return self._size def __iter__(self) -> Iterator[Any]: return iter(self[i] for i in range(len(self))) def __iadd__(self, parameters: Iterable[Any]) -> 'ParameterList': return self.extend(parameters) def __dir__(self): keys = super(ParameterList, self).__dir__() keys = [key for key in keys if not key.isdigit()] return keys def append(self, value: Any) -> 'ParameterList': """Appends a given value at the end of the list. Args: value (Any): value to append """ new_idx = len(self) self._size += 1 self[new_idx] = value return self def extend(self, values: Iterable[Any]) -> 'ParameterList': """Appends values from a Python iterable to the end of the list. Args: values (iterable): iterable of values to append """ # Tensor is an iterable but we never want to unpack it here if not isinstance(values, container_abcs.Iterable) or isinstance(values, Tensor): raise TypeError("ParameterList.extend should be called with an " "iterable, but got " + type(values).__name__) for value in values: self.append(value) return self def extra_repr(self) -> str: child_lines = [] for k, p in enumerate(self): if isinstance(p, Tensor): size_str = 'x'.join(str(size) for size in p.size()) device_str = '' if not p.is_cuda else ' (GPU {})'.format(p.get_device()) parastr = '{} containing: [{} of size {}{}]'.format( "Parameter" if isinstance(p, Parameter) else "Tensor", p.dtype, size_str, device_str) child_lines.append(' (' + str(k) + '): ' + parastr) else: child_lines.append(' (' + str(k) + '): Object of type: ' + type(p).__name__) tmpstr = '\n'.join(child_lines) return tmpstr def __call__(self, *args, **kwargs): raise RuntimeError('ParameterList should not be called.') class ParameterDict(Module): """Holds parameters in a dictionary. ParameterDict can be indexed like a regular Python dictionary, but Parameters it contains are properly registered, and will be visible by all Module methods. Other objects are treated as would be done by a regular Python dictionary :class:`nn.ParameterDict` is an **ordered** dictionary. :meth:`nn.ParameterDict.update` with other unordered mapping types (e.g., Python's plain ``dict``) does not preserve the order of the merged mapping. On the other hand, ``OrderedDict`` or another :class:`nn.ParameterDict` will preserve their ordering. Note that the constructor, assigning an element of the dictionary and the :meth:`nn.ParameterDict.update` method will convert any :class:`Tensor` into :class:`nn.Parameter`. Args: values (iterable, optional): a mapping (dictionary) of (string : Any) or an iterable of key-value pairs of type (string, Any) Example:: class MyModule(nn.Module): def __init__(self): super(MyModule, self).__init__() self.params = nn.ParameterDict({ 'left': nn.Parameter(ms_torch.randn(5, 10)), 'right': nn.Parameter(ms_torch.randn(5, 10)) }) def forward(self, x, choice): x = self.params[choice].mm(x) return x """ def __init__(self, parameters: Any = None) -> None: super(ParameterDict, self).__init__() self._keys: Dict[str, None] = {} if parameters is not None: self.update(parameters) def _key_to_attr(self, key: str) -> str: if not isinstance(key, str): raise TypeError("Index given to ParameterDict cannot be used as a key as it is " f"not a string (type is '{type(key).__name__}'). Open an issue on " "github if you need non-string keys.") else: # Use the key as-is so that `.named_parameters()` returns the right thing return key def __getitem__(self, key: str) -> Any: attr = self._key_to_attr(key) return getattr(self, attr) def __setitem__(self, key: str, value: Any) -> None: # Note that all other function that add an entry to the dictionary part of # the ParameterDict end up here. So this is the only place where we need # to wrap things into Parameter if needed. # Objects added via setattr() are not in the dictionary part and thus won't # call into this function. self._keys[key] = None attr = self._key_to_attr(key) if isinstance(value, Tensor) and not isinstance(value, Parameter): value = Parameter(value) setattr(self, attr, value) def __delitem__(self, key: str) -> None: del self._keys[key] attr = self._key_to_attr(key) delattr(self, attr) def __len__(self) -> int: return len(self._keys) def __iter__(self) -> Iterator[str]: return iter(self._keys) def __reversed__(self) -> Iterator[str]: return reversed(list(self._keys)) def copy(self) -> 'ParameterDict': """Returns a copy of this :class:`nn.ParameterDict` instance. """ # We have to use an OrderedDict because the ParameterDict constructor # behaves differently on plain dict vs OrderedDict return ParameterDict(OrderedDict((k, self[k]) for k in self._keys)) def __contains__(self, key: str) -> bool: return key in self._keys def setdefault(self, key: str, default: Optional[Any] = None) -> Any: """If key is in the ParameterDict, return its value. If not, insert `key` with a parameter `default` and return `default`. `default` defaults to `None`. Args: key (str): key to set default for default (Any): the parameter set to the key """ if key not in self: self[key] = default return self[key] def clear(self) -> None: """Remove all items from the ParameterDict. """ for k in self._keys.copy(): del self[k] def pop(self, key: str) -> Any: r"""Remove key from the ParameterDict and return its parameter. Args: key (str): key to pop from the ParameterDict """ v = self[key] del self[key] return v def popitem(self) -> Tuple[str, Any]: """Remove and return the last inserted `(key, parameter)` pair from the ParameterDict """ k, _ = self._keys.popitem() # We need the key in the _keys to be able to access/del self._keys[k] = None val = self[k] del self[k] return k, val def get(self, key: str, default: Optional[Any] = None) -> Any: r"""Return the parameter associated with key if present. Otherwise return default if provided, None if not. Args: key (str): key to get from the ParameterDict default (Parameter, optional): value to return if key not present """ return self[key] if key in self else default def fromkeys(self, keys: Iterable[str], default: Optional[Any] = None) -> 'ParameterDict': r"""Return a new ParameterDict with the keys provided Args: keys (iterable, string): keys to make the new ParameterDict from default (Parameter, optional): value to set for all keys """ return ParameterDict(((k, default) for k in keys)) def keys(self) -> Iterable[str]: r"""Return an iterable of the ParameterDict keys. """ return self._keys.keys() def items(self) -> Iterable[Tuple[str, Any]]: r"""Return an iterable of the ParameterDict key/value pairs. """ return ((k, self[k]) for k in self._keys) def values(self) -> Iterable[Any]: r"""Return an iterable of the ParameterDict values. """ return (self[k] for k in self._keys) def update(self, parameters: Union[Mapping[str, Any], 'ParameterDict']) -> None: r"""Update the :class:`~nn.ParameterDict` with the key-value pairs from a mapping or an iterable, overwriting existing keys. .. note:: If :attr:`parameters` is an ``OrderedDict``, a :class:`~nn.ParameterDict`, or an iterable of key-value pairs, the order of new elements in it is preserved. Args: parameters (iterable): a mapping (dictionary) from string to :class:`~nn.Parameter`, or an iterable of key-value pairs of type (string, :class:`~nn.Parameter`) """ if not isinstance(parameters, container_abcs.Iterable): raise TypeError("ParametersDict.update should be called with an " "iterable of key/value pairs, but got " + type(parameters).__name__) if isinstance(parameters, (OrderedDict, ParameterDict)): for key, parameter in parameters.items(): self[key] = parameter elif isinstance(parameters, container_abcs.Mapping): for key, parameter in sorted(parameters.items()): self[key] = parameter else: for j, p in enumerate(parameters): if not isinstance(p, container_abcs.Iterable): raise TypeError("ParameterDict update sequence element " "#" + str(j) + " should be Iterable; is" + type(p).__name__) if not len(p) == 2: raise ValueError("ParameterDict update sequence element " "#" + str(j) + " has length " + str(len(p)) + "; 2 is required") # parameters as length-2 list too cumbersome to type, see ModuleDict.update comment self[p[0]] = p[1] # type: ignore[assignment] def extra_repr(self) -> str: child_lines = [] for k, p in self.items(): if isinstance(p, Tensor): size_str = 'x'.join(str(size) for size in p.size()) device_str = '' if not p.is_cuda else ' (GPU {})'.format(p.get_device()) parastr = '{} containing: [{} of size {}{}]'.format( "Parameter" if isinstance(p, Parameter) else "Tensor", typename(p), size_str, device_str) child_lines.append(' (' + str(k) + '): ' + parastr) else: child_lines.append(' (' + str(k) + '): Object of type: ' + type(p).__name__) tmpstr = '\n'.join(child_lines) return tmpstr def __call__(self, input): raise RuntimeError('ParameterDict should not be called.') def __or__(self, other: 'ParameterDict') -> 'ParameterDict': copy = self.copy() copy.update(other) return copy def __ror__(self, other: 'ParameterDict') -> 'ParameterDict': copy = other.copy() copy.update(self) return copy def __ior__(self, other : 'ParameterDict') -> 'ParameterDict': self.update(other) return self msadapter/ms_adapter/pytorch/nn/modules/conv.py000066400000000000000000000707101436713554100222650ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import math from mindspore.ops import operations as P from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.pytorch.nn import init from ms_adapter.pytorch.functional import empty from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from ms_adapter.pytorch.nn.functional import conv2d from .utils import _triple, _pair, _single, _reverse_repeat_tuple from .module import Module __all__ = ['Conv1d', 'Conv2d', 'Conv3d', 'ConvTranspose1d', 'ConvTranspose2d', 'ConvTranspose3d', 'LazyConv1d', 'LazyConv2d', 'LazyConv3d', 'LazyConvTranspose1d', 'LazyConvTranspose2d', 'LazyConvTranspose3d'] class _ConvNd(Module): def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, output_padding, groups, bias, padding_mode, device=None, dtype=None, transposed=False ): """Initialize _Conv.""" unsupported_attr(device) unsupported_attr(dtype) super(_ConvNd, self).__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = kernel_size self.stride = stride self.padding = padding self.dilation = dilation self.transposed = transposed self.output_padding = output_padding self.groups = groups self.padding_mode = padding_mode # MS add self.pad_mode = 'same' self.data_format = 'NCHW' if in_channels % groups != 0: raise ValueError('in_channels must be divisible by groups') if out_channels % groups != 0: raise ValueError('out_channels must be divisible by groups') valid_padding_strings = {'same', 'valid'} if isinstance(padding, str): if padding not in valid_padding_strings: raise ValueError( "Invalid padding string {!r}, should be one of {}".format( padding, valid_padding_strings)) if padding == 'same' and any(s != 1 for s in stride): raise ValueError("padding='same' is not supported for strided convolutions") if isinstance(self.padding, str): self._reversed_padding_repeated_twice = [0, 0] * len(kernel_size) if padding == 'same': for d, k, i in zip(dilation, kernel_size, range(len(kernel_size) - 1, -1, -1)): total_padding = d * (k - 1) left_pad = total_padding // 2 self._reversed_padding_repeated_twice[2 * i] = left_pad self._reversed_padding_repeated_twice[2 * i + 1] = ( total_padding - left_pad) else: self._reversed_padding_repeated_twice = _reverse_repeat_tuple(self.padding, 2) if transposed: self.weight = Parameter(empty((in_channels, out_channels // groups, *kernel_size))) else: self.weight = Parameter(empty((out_channels, in_channels // groups, *kernel_size))) if bias: self.bias = Parameter(empty(out_channels)) else: self.bias = None self.reset_parameters() def reset_parameters(self): init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) if fan_in != 0: bound = 1 / math.sqrt(fan_in) init.uniform_(self.bias, -bound, bound) def extra_repr(self): s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}' ', stride={stride}') if self.padding != (0,) * len(self.padding): s += ', padding={padding}' if self.dilation != (1,) * len(self.dilation): s += ', dilation={dilation}' if self.output_padding != (0,) * len(self.output_padding): s += ', output_padding={output_padding}' if self.groups != 1: s += ', groups={groups}' if self.bias is None: s += ', bias=False' if self.padding_mode != 'zeros': s += ', padding_mode={padding_mode}' return s.format(**self.__dict__) class Conv1d(_ConvNd): r""" 1D convolution layer. Calculates the 1D convolution on the input tensor which is typically of shape :math:`(N, C_{in}, L_{in})`, where :math:`N` is batch size, :math:`C_{in}` is a number of channels and :math:`L_{in}` is a length of sequence. For the tensor of each batch, its shape is :math:`(C_{in}, L_{in})`, the formula is defined as: Supported Platforms: ``Ascend`` ``GPU`` ``CPU`` Examples: >>> net = nn.Conv1d(120, 240, 4, has_bias=False, weight_init='normal') >>> x = Tensor(np.ones([1, 120, 640]), mindspore.float32) >>> output = net(x).shape >>> print(output) (1, 240, 640) """ def __init__( self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None ): factory_kwargs = {'device': device, 'dtype': dtype, 'transposed': False} self.has_bias = False if bias: self.has_bias=True kernel_size_ = (1, kernel_size) if isinstance(kernel_size, int) else (1, kernel_size[0]) stride_ = (1, stride if isinstance(stride, int) else stride[0]) dilation_ = (1, dilation) if isinstance(dilation, int) else (1, dilation[0]) padding_ = padding if isinstance(padding, str) else _single(padding) super(Conv1d, self).__init__(in_channels, out_channels, kernel_size_, stride_, padding_, dilation_, _pair(0), groups, bias, padding_mode, **factory_kwargs) #TODO pad_mode in ['zeros', 'reflect', 'replicate', 'circular'] if padding_mode in {'reflect', 'replicate', 'circular'}: raise ValueError("Pad mode '{}' is not currently supported.".format(padding_mode)) if padding == 0: self.pad_mode = 'valid' self.padding =(0, 0, 0, 0) elif isinstance(self.padding, str): self.pad_mode = self.padding self.padding = 0 elif padding_mode == 'zeros': self.pad_mode = "pad" self.padding =(0, 0, padding, padding) self.conv2d = P.Conv2D(out_channel=self.out_channels, kernel_size=self.kernel_size, mode=1, pad_mode=self.pad_mode, pad=self.padding, stride=self.stride, dilation=self.dilation, group=groups) self.bias_add = P.BiasAdd() self.expand_dims = P.ExpandDims() self.squeeze = P.Squeeze(2) def forward(self, input): input = cast_to_ms_tensor(input) ndim = input.ndim if ndim == 2: input = input.expand_dims(0) x = self.expand_dims(input, 2) output = self.conv2d(x, self.weight) if self.has_bias: output = self.bias_add(output, self.bias) output = self.squeeze(output) output = output.squeeze(0) else: x = self.expand_dims(input, 2) output = self.conv2d(x, self.weight) if self.has_bias: output = self.bias_add(output, self.bias) output = self.squeeze(output) return cast_to_adapter_tensor(output) class Conv2d(_ConvNd): def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None): """Initialize Conv2d.""" factory_kwargs = {'device': device, 'dtype': dtype, 'transposed': False} kernel_size_ = _pair(kernel_size) stride_ = _pair(stride) padding_ = padding if isinstance(padding, str) else _pair(padding) dilation_ = _pair(dilation) super(Conv2d, self).__init__(in_channels, out_channels, kernel_size_, stride_, padding_, dilation_, _pair(0), groups, bias, padding_mode, **factory_kwargs) #TODO pad_mode in ['zeros', 'reflect', 'replicate', 'circular'] if padding_mode in {'reflect', 'replicate', 'circular'}: raise ValueError("Pad mode '{}' is not currently supported.".format(padding_mode)) def forward(self, x): x = cast_to_ms_tensor(x) ndim = x.ndim if ndim == 3: x = x.expand_dims(0) output = conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups) output = output.squeeze(0) else: output = conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups) return cast_to_adapter_tensor(output) class Conv3d(_ConvNd): r""" 3D convolution layer. Calculates the 3D convolution on the input tensor which is typically of shape Supported Platforms: ``Ascend`` ``GPU`` ``CPU`` Examples: >>> x = Tensor(np.ones([16, 3, 10, 32, 32]), mindspore.float32) >>> conv3d = nn.Conv3d(in_channels=3, out_channels=32, kernel_size=(4, 3, 3)) >>> output = conv3d(x) >>> print(output.shape) (16, 32, 10, 32, 32) """ def __init__( self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None ): factory_kwargs = {'device': device, 'dtype': dtype, 'transposed': False} self.has_bias = False if bias: self.has_bias=True kernel_size_ = _triple(kernel_size) stride_ = _triple(stride) padding_ = padding if isinstance(padding, str) else _triple(padding) dilation_ = _triple(dilation) super(Conv3d, self).__init__(in_channels, out_channels, kernel_size_, stride_, padding_, dilation_, _pair(0), groups, bias, padding_mode, **factory_kwargs) #TODO pad_mode in ['zeros', 'reflect', 'replicate', 'circular'] if padding_mode in {'reflect', 'replicate', 'circular'}: raise ValueError("Pad mode '{}' is not currently supported.".format(padding_mode)) if padding == 0: self.pad_mode = 'valid' self.padding =(self.padding[0], self.padding[0], self.padding[1], self.padding[1], self.padding[2], self.padding[2]) elif isinstance(self.padding, str): self.pad_mode = self.padding self.padding = 0 elif padding_mode == 'zeros': self.pad_mode = "pad" self.padding =(self.padding[0], self.padding[0], self.padding[1], self.padding[1], self.padding[2], self.padding[2]) self.conv3d = P.Conv3D(out_channel=self.out_channels, kernel_size=self.kernel_size, mode=1, pad_mode=self.pad_mode, pad=self.padding, stride=self.stride, dilation=self.dilation, group=groups, data_format='NCDHW') self.bias_add = P.BiasAdd(data_format='NCDHW') self.shape = P.Shape() def forward(self, input): input = cast_to_ms_tensor(input) ndim = input.ndim if ndim == 4: input = input.expand_dims(0) output = self.conv3d(input, self.weight) if self.has_bias: output = self.bias_add(output, self.bias) output = output.squeeze(0) else: output = self.conv3d(input, self.weight) if self.has_bias: output = self.bias_add(output, self.bias) return cast_to_adapter_tensor(output) class ConvTranspose1d(_ConvNd): r""" 1D transposed convolution layer. Calculates a 1D transposed convolution, which can be regarded as Conv1d for the gradient of the input. It also called deconvolution (although it is not an actual deconvolution). Supported Platforms: ``Ascend`` ``GPU`` ``CPU`` Examples: >>> net = nn.ConvTranspose1d(3, 64, 4, has_bias=False) >>> x = Tensor(np.ones([1, 3, 50]), mindspore.float32) >>> output = net(x).shape >>> print(output) (1, 64, 53) """ def __init__( self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros', device=None, dtype=None, ): factory_kwargs = {'device': device, 'dtype': dtype, 'transposed': True} self.has_bias = False if bias: self.has_bias=True _padding = _single(padding) _kernel_size = (1, kernel_size) if isinstance(kernel_size, int) else (1, kernel_size[0]) _stride = (1, stride) if isinstance(stride, int) else (1, stride[0]) _dilation = (1, dilation) if isinstance(dilation, int) else (1, dilation[0]) super(ConvTranspose1d, self).__init__(in_channels, out_channels, _kernel_size, _stride, _padding, _dilation, output_padding, groups, bias, padding_mode, **factory_kwargs) self.shape = P.Shape() if padding_mode in {'reflect', 'replicate', 'circular'}: raise ValueError("Pad mode '{}' is not currently supported.".format(padding_mode)) if output_padding > 0: raise ValueError("output_padding '{}' is not currently supported.".format(output_padding)) if (isinstance(padding, int) and padding == 0) or (isinstance(padding, tuple) and padding[0] == 0): self.pad_mode = 'valid' self.padding = (0, 0, 0, 0) elif padding_mode == 'zeros': self.pad_mode = "pad" if isinstance(padding, int): self.padding = (0, 0, padding, padding) else: self.padding = (0, 0, padding[0], padding[0]) self.is_valid = self.pad_mode == 'valid' self.is_same = self.pad_mode == 'same' self.is_pad = self.pad_mode == 'pad' # cause Conv2DBackpropInput's out_channel refers to Conv2D's out_channel. self.conv2d_transpose = P.Conv2DBackpropInput(out_channel=self.in_channels, kernel_size=self.kernel_size, mode=1, pad_mode=self.pad_mode, pad=self.padding, stride=self.stride, dilation=self.dilation, group=groups) self.bias_add = P.BiasAdd() self.expand_dims = P.ExpandDims() self.squeeze = P.Squeeze(2) def construct(self, input, output_size=None): if output_size is not None: raise ValueError("output_size '{}' is not currently supported.".format(output_size)) x = cast_to_ms_tensor(input) ndim = x.ndim if ndim == 2: x = x.expand_dims(0) x = self.expand_dims(x, 2) n, _, h, w = self.shape(x) h_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, h, self.kernel_size[0], self.stride[0], self.dilation[0], self.padding[0] + self.padding[1]) w_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, w, self.kernel_size[1], self.stride[1], self.dilation[1], self.padding[2] + self.padding[3]) output = self.conv2d_transpose(x, self.weight, (n, self.out_channels, h_out, w_out)) if self.has_bias: output = self.bias_add(output, self.bias) output = self.squeeze(output) output = output.squeeze(0) else: x = self.expand_dims(x, 2) n, _, h, w = self.shape(x) h_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, h, self.kernel_size[0], self.stride[0], self.dilation[0], self.padding[0] + self.padding[1]) w_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, w, self.kernel_size[1], self.stride[1], self.dilation[1], self.padding[2] + self.padding[3]) output = self.conv2d_transpose(x, self.weight, (n, self.out_channels, h_out, w_out)) if self.has_bias: output = self.bias_add(output, self.bias) output = self.squeeze(output) return cast_to_adapter_tensor(output) def extend_repr(self): s = 'input_channels={}, output_channels={}, kernel_size={}, ' \ 'stride={}, pad_mode={}, padding={}, dilation={}, ' \ 'group={}, has_bias={}, ' \ 'weight_init={}'.format(self.in_channels, self.out_channels, self.kernel_size, self.stride, self.pad_mode, self.padding, self.dilation, self.group, self.bias, self.weight_init, ) return s class ConvTranspose2d(_ConvNd): r""" 2D transposed convolution layer. Calculates a 2D transposed convolution, which can be regarded as Conv2d for the gradient of the input. It also called deconvolution (although it is not an actual deconvolution). Supported Platforms: ``Ascend`` ``GPU`` ``CPU`` Examples: >>> net = nn.ConvTranspose2d(3, 64, 4, has_bias=False) >>> x = Tensor(np.ones([1, 3, 16, 50]), mindspore.float32) >>> output = net(x).shape >>> print(output) (1, 64, 19, 53) """ def __init__( self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros', device=None, dtype=None ): factory_kwargs = {'device': device, 'dtype': dtype, 'transposed': True} self.has_bias = False if bias: self.has_bias=True _kernel_size = _pair(kernel_size) _stride = _pair(stride) _padding = _pair(padding) _dilation = _pair(dilation) output_padding = _pair(output_padding) super(ConvTranspose2d, self).__init__(in_channels, out_channels, _kernel_size, _stride, _padding, _dilation, output_padding, groups, bias, padding_mode, **factory_kwargs) self.shape = P.Shape() if padding == 0: self.pad_mode = 'valid' self.padding =(self.padding[0], self.padding[0], self.padding[1], self.padding[1]) elif isinstance(self.padding, str): self.pad_mode = self.padding self.padding = 0 elif padding_mode == 'zeros': self.pad_mode = "pad" self.padding =(self.padding[0], self.padding[0], self.padding[1], self.padding[1]) if self.padding_mode != 'zeros': raise ValueError('Only `zeros` padding mode is supported for ConvTranspose2d') self.is_valid = self.pad_mode == 'valid' self.is_same = self.pad_mode == 'same' self.is_pad = self.pad_mode == 'pad' # cause Conv2DTranspose's out_channel refers to Conv2D's out_channel. self.conv2d_transpose = P.Conv2DTranspose(out_channel=self.in_channels, kernel_size=self.kernel_size, mode=1, pad_mode=self.pad_mode, pad=self.padding, stride=self.stride, dilation=self.dilation, group=groups) self.bias_add = P.BiasAdd() if isinstance(self.padding, int): self.padding_top, self.padding_bottom, self.padding_left, self.padding_right = (self.padding,) * 4 else: self.padding_top, self.padding_bottom, self.padding_left, self.padding_right = self.padding def forward(self, input, output_size = None): if output_size is not None: raise ValueError("output_size '{}' is not currently supported.".format(output_size)) x = cast_to_ms_tensor(input) ndim = x.ndim if ndim == 3: x = x.expand_dims(0) n, _, h, w = self.shape(x) h_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, h, self.kernel_size[0], self.stride[0], self.dilation[0], self.padding_top + self.padding_bottom) w_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, w, self.kernel_size[1], self.stride[1], self.dilation[1], self.padding_left + self.padding_right) output = self.conv2d_transpose(x, self.weight, (n, self.out_channels, h_out, w_out)) if self.has_bias: output = self.bias_add(output, self.bias) output = output.squeeze(0) else: n, _, h, w = self.shape(x) h_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, h, self.kernel_size[0], self.stride[0], self.dilation[0], self.padding_top + self.padding_bottom) w_out = _deconv_output_length(self.is_valid, self.is_same, self.is_pad, w, self.kernel_size[1], self.stride[1], self.dilation[1], self.padding_left + self.padding_right) output = self.conv2d_transpose(x, self.weight, (n, self.out_channels, h_out, w_out)) if self.has_bias: output = self.bias_add(output, self.bias) return cast_to_adapter_tensor(output) def extend_repr(self): s = 'input_channels={}, output_channels={}, kernel_size={}, ' \ 'stride={}, pad_mode={}, padding={}, dilation={}, ' \ 'group={}, has_bias={}'.format(self.in_channels, self.out_channels, self.kernel_size, self.stride, self.pad_mode, self.padding, self.dilation, self.group, self.has_bias) return s class ConvTranspose3d(_ConvNd): r""" 3D transposed convolution layer. Calculates a 3D transposed convolution, which can be regarded as Conv3d for the gradient of the input. It also called deconvolution (although it is not an actual deconvolution). Examples: >>> x = Tensor(np.ones([32, 16, 10, 32, 32]), mindspore.float32) >>> conv3d_transpose = nn.ConvTranspose3d(in_channels=16, out_channels=3, kernel_size=(4, 6, 2), ... pad_mode='pad') >>> output = conv3d_transpose(x) >>> print(output.shape) (32, 3, 13, 37, 33) """ def __init__( self, in_channels, out_channels, kernel_size, stride = 1, padding = 0, output_padding = 0, groups = 1, bias = True, dilation = 1, padding_mode = 'zeros', device=None, dtype=None ): factory_kwargs = {'device': device, 'dtype': dtype, 'transposed': True} _kernel_size = _triple(kernel_size) _stride = _triple(stride) _padding = _triple(padding) _dilation = _triple(dilation) output_padding = _triple(output_padding) super(ConvTranspose3d, self).__init__(in_channels, out_channels, _kernel_size, _stride, _padding, _dilation, output_padding, groups, bias, padding_mode, **factory_kwargs) if padding == 0: self.pad_mode = 'valid' self.padding =(self.padding[0], self.padding[0], self.padding[1], self.padding[1],self.padding[2], self.padding[2]) elif isinstance(self.padding, str): self.pad_mode = self.padding self.padding = 0 elif padding_mode == 'zeros': self.pad_mode = "pad" self.padding =(self.padding[0], self.padding[0], self.padding[1], self.padding[1], self.padding[2], self.padding[2]) if self.padding_mode != 'zeros': raise ValueError('Only `zeros` padding mode is supported for ConvTranspose3d') self.conv3d_transpose = P.Conv3DTranspose(in_channel=self.in_channels, out_channel=self.out_channels, kernel_size=self.kernel_size, mode=1, pad_mode=self.pad_mode, pad=self.padding, stride=self.stride, dilation=self.dilation, group=groups, output_padding=self.output_padding, data_format='NCDHW') self.bias_add = P.BiasAdd(data_format='NCDHW') self.has_bias = False if bias: self.has_bias = True def forward(self, input, output_size = None): x = cast_to_ms_tensor(input) ndim = x.ndim if ndim == 4: x = x.expand_dims(0) if output_size is not None: raise ValueError("output_size '{}' is not currently supported.".format(output_size)) output = self.conv3d_transpose(x, self.weight) if self.has_bias: output = self.bias_add(output, self.bias) output = output.squeeze(0) else: if output_size is not None: raise ValueError("output_size '{}' is not currently supported.".format(output_size)) output = self.conv3d_transpose(x, self.weight) if self.has_bias: output = self.bias_add(output, self.bias) return cast_to_adapter_tensor(output) def _deconv_output_length(is_valid, is_same, is_pad, input_length, filter_size, stride_size, dilation_size, padding): """Calculate the width and height of output.""" length = 0 filter_size = filter_size + (filter_size - 1) * (dilation_size - 1) if is_valid: if filter_size - stride_size > 0: length = input_length * stride_size + filter_size - stride_size else: length = input_length * stride_size elif is_same: length = input_length * stride_size elif is_pad: length = input_length * stride_size - padding + filter_size - stride_size return length LazyConv1d = Conv1d LazyConv2d = Conv2d LazyConv3d = Conv3d LazyConvTranspose1d = ConvTranspose1d LazyConvTranspose2d = ConvTranspose2d LazyConvTranspose3d = ConvTranspose3d msadapter/ms_adapter/pytorch/nn/modules/distance.py000066400000000000000000000014451436713554100231110ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import ms_adapter.pytorch.nn.functional as ms_torch_nn_func from .module import Module __all__ = ['PairwiseDistance', 'CosineSimilarity'] class PairwiseDistance(Module): def __init__(self, p=2.0, eps=1e-6, keepdim=False): super(PairwiseDistance, self).__init__() self.norm = p self.eps = eps self.keepdim = keepdim def forward(self, x1, x2): return ms_torch_nn_func.pairwise_distance(x1, x2, self.norm, self.eps, self.keepdim) class CosineSimilarity(Module): def __init__(self, dim=1, eps=1e-8): super(CosineSimilarity, self).__init__() self.dim = dim self.eps = eps def forward(self, x1, x2): return ms_torch_nn_func.cosine_similarity(x1, x2, self.dim, self.eps) msadapter/ms_adapter/pytorch/nn/modules/dropout.py000066400000000000000000000207101436713554100230070ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from mindspore.common.seed import _get_graph_seed from mindspore.ops import operations as P from ms_adapter.pytorch.tensor import cast_to_ms_tensor import ms_adapter.pytorch.nn.functional as ms_torch_nn_func from ms_adapter.pytorch.common._inner import _inplace_assign, _inplace_limit_pynative from .module import Module __all__ = ['Dropout', 'Dropout2d', 'Dropout3d', 'AlphaDropout', 'FeatureAlphaDropout'] class _DropoutNd(Module): def __init__(self, p=0.5, inplace=False): super(_DropoutNd, self).__init__() self.p = p self.inplace = inplace def extra_repr(self) -> str: return 'p={}, inplace={}'.format(self.p, self.inplace) class Dropout(_DropoutNd): r"""During training, randomly zeroes some of the elements of the input tensor with probability :attr:`p` using samples from a Bernoulli distribution. Each channel will be zeroed out independently on every forward call. This has proven to be an effective technique for regularization and preventing the co-adaptation of neurons as described in the paper `Improving neural networks by preventing co-adaptation of feature detectors`_ . Furthermore, the outputs are scaled by a factor of :math:`\frac{1}{1-p}` during training. This means that during evaluation the module simply computes an identity function. Args: p: probability of an element to be zeroed. Default: 0.5 inplace: If set to ``True``, will do this operation in-place. Default: ``False`` Shape: - Input: :math:`(*)`. Input can be of any shape - Output: :math:`(*)`. Output is of the same shape as input Examples:: >>> m = nn.Dropout(p=0.2) >>> input = ms_adapter.pytorch.randn(20, 16) >>> output = m(input) .. _Improving neural networks by preventing co-adaptation of feature detectors: https://arxiv.org/abs/1207.0580 """ def __init__(self, p=0.5, inplace=False): _inplace_limit_pynative(inplace, "Dropout") super(Dropout, self).__init__(p, inplace) if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) self.keep_prob = 1.0 - self.p seed0, seed1 = _get_graph_seed(0, "dropout") self.seed0 = seed0 self.seed1 = seed1 self.dropout = P.Dropout(self.keep_prob, seed0, seed1) def forward(self, input): if not self.training: return input if self.keep_prob == 1.0: return input input_ms = cast_to_ms_tensor(input) output, _ = self.dropout(input_ms) return _inplace_assign(input, self.inplace, output) class Dropout2d(_DropoutNd): r"""Randomly zero out entire channels (a channel is a 2D feature map, e.g., the :math:`j`-th channel of the :math:`i`-th sample in the batched input is a 2D tensor :math:`\text{input}[i, j]`). Each channel will be zeroed out independently on every forward call with probability :attr:`p` using samples from a Bernoulli distribution. Usually the input comes from :class:`nn.Conv2d` modules. As described in the paper `Efficient Object Localization Using Convolutional Networks`_ , if adjacent pixels within feature maps are strongly correlated (as is normally the case in early convolution layers) then i.i.d. dropout will not regularize the activations and will otherwise just result in an effective learning rate decrease. In this case, :func:`nn.Dropout2d` will help promote independence between feature maps and should be used instead. Args: p (float, optional): probability of an element to be zero-ed. inplace (bool, optional): If set to ``True``, will do this operation in-place Shape: - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`. - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input). Examples:: >>> m = nn.Dropout2d(p=0.2) >>> input = ms_adapter.randn(20, 16, 32, 32) >>> output = m(input) .. _Efficient Object Localization Using Convolutional Networks: https://arxiv.org/abs/1411.4280 """ def __init__(self, p=0.5, inplace=False): _inplace_limit_pynative(inplace, "Dropout2d") super(Dropout2d, self).__init__(p, inplace) if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) self.keep_prob = 1.0 - self.p self.dropout2d = P.Dropout2D(self.keep_prob) def forward(self, input): if not self.training: return input if self.keep_prob == 1: return input input_ms = cast_to_ms_tensor(input) output, _ = self.dropout2d(input_ms) return _inplace_assign(input, self.inplace, output) class Dropout3d(_DropoutNd): r"""Randomly zero out entire channels (a channel is a 3D feature map, e.g., the :math:`j`-th channel of the :math:`i`-th sample in the batched input is a 3D tensor :math:`\text{input}[i, j]`). Each channel will be zeroed out independently on every forward call with probability :attr:`p` using samples from a Bernoulli distribution. Usually the input comes from :class:`nn.Conv3d` modules. As described in the paper `Efficient Object Localization Using Convolutional Networks`_ , if adjacent pixels within feature maps are strongly correlated (as is normally the case in early convolution layers) then i.i.d. dropout will not regularize the activations and will otherwise just result in an effective learning rate decrease. In this case, :func:`nn.Dropout3d` will help promote independence between feature maps and should be used instead. Args: p (float, optional): probability of an element to be zeroed. inplace (bool, optional): If set to ``True``, will do this operation in-place Shape: - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`. - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input). Examples:: >>> m = nn.Dropout3d(p=0.2) >>> input = ms_adapter.randn(20, 16, 4, 32, 32) >>> output = m(input) .. _Efficient Object Localization Using Convolutional Networks: https://arxiv.org/abs/1411.4280 """ def __init__(self, p=0.5, inplace=False): _inplace_limit_pynative(inplace, "Dropout3d") super(Dropout3d, self).__init__(p, inplace) if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) self.keep_prob = 1.0 - self.p self.dropout3d = P.Dropout3D(self.keep_prob) def forward(self, input): if not self.training: return input if self.keep_prob == 1: return input input_ms = cast_to_ms_tensor(input) output, _ = self.dropout3d(input_ms) return _inplace_assign(input, self.inplace, output) class AlphaDropout(_DropoutNd): def __init__(self, p=0.5, inplace=False): _inplace_limit_pynative(inplace, "AlphaDropout") super(AlphaDropout, self).__init__(p, inplace) if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) def forward(self, input): return ms_torch_nn_func.alpha_dropout(input, self.p, self.training, self.inplace) class FeatureAlphaDropout(_DropoutNd): def __init__(self, p=0.5, inplace=False): _inplace_limit_pynative(inplace, "FeatureAlphaDropout") super(FeatureAlphaDropout, self).__init__(p, inplace) if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) def forward(self, input): return ms_torch_nn_func.feature_alpha_dropout(input, self.p, self.training, self.inplace) class Dropout1d(_DropoutNd): def __init__(self, p=0.5, inplace=False): _inplace_limit_pynative(inplace, "Dropout1d") super(Dropout1d, self).__init__(p, inplace) if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) def forward(self, input): return ms_torch_nn_func.dropout1d(input, self.p, self.training, self.inplace) msadapter/ms_adapter/pytorch/nn/modules/flatten.py000066400000000000000000000021251436713554100227500ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import ms_adapter.pytorch.functional as adapter_F from .module import Module __all__ = ['Flatten', 'Unflatten'] class Flatten(Module): def __init__(self, start_dim=1, end_dim=-1): super(Flatten, self).__init__() self.start_dim = start_dim self.end_dim = end_dim def forward(self, input): return adapter_F.flatten(input, self.start_dim, self.end_dim) def extra_repr(self) -> str: return 'start_dim={}, end_dim={}'.format( self.start_dim, self.end_dim ) class Unflatten(Module): def __init__(self, dim, unflattened_size): super(Unflatten, self).__init__() if isinstance(dim, str): raise TypeError("Until Now, `dim` not support type of str in `Unflatten`") self.dim = dim self.unflattened_size = unflattened_size def forward(self, input): return adapter_F.unflatten(input, self.dim, self.unflattened_size) def extra_repr(self) -> str: return 'dim={}, unflattened_size={}'.format(self.dim, self.unflattened_size) msadapter/ms_adapter/pytorch/nn/modules/fold.py000066400000000000000000000026651436713554100222500ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch.nn.functional import fold, unfold from .module import Module __all__ = ['Fold', 'Unfold'] class Fold(Module): #TODO: do not support on Ascend def __init__(self, output_size, kernel_size, dilation=1, padding=0, stride=1): super(Fold, self).__init__() self.output_size = output_size self.kernel_size = kernel_size self.dilation = dilation self.padding = padding self.stride = stride def forward(self, input): return fold(input, self.output_size, self.kernel_size, self.dilation, self.padding, self.stride) def extra_repr(self): return 'output_size={output_size}, kernel_size={kernel_size}, ' \ 'dilation={dilation}, padding={padding}, stride={stride}'.format( **self.__dict__ ) class Unfold(Module): # TODO: do not support on GPU def __init__(self, kernel_size, dilation=1, padding=0, stride=1): super(Unfold, self).__init__() self.kernel_size = kernel_size self.dilation = dilation self.padding = padding self.stride = stride def forward(self, input): return unfold(input, self.kernel_size, self.dilation, self.padding, self.stride) def extra_repr(self): return 'kernel_size={kernel_size}, dilation={dilation}, padding={padding},' \ ' stride={stride}'.format(**self.__dict__) msadapter/ms_adapter/pytorch/nn/modules/linear.py000066400000000000000000000130361436713554100225700ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import math import mindspore.ops as P from ms_adapter.pytorch.nn import init from ms_adapter.pytorch.nn.functional import linear from ms_adapter.pytorch.functional import empty from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from .module import Module __all__ = ['Linear', 'LazyLinear', 'Identity', 'Bilinear'] class Linear(Module): r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b` Args: in_features: size of each input sample out_features: size of each output sample bias: If set to ``False``, the layer will not learn an additive bias. Default: ``True`` Shape: - Input: :math:`(*, H_{in})` where :math:`*` means any number of dimensions including none and :math:`H_{in} = \text{in\_features}`. - Output: :math:`(*, H_{out})` where all but the last dimension are the same shape as the input and :math:`H_{out} = \text{out\_features}`. Attributes: weight: the learnable weights of the module of shape :math:`(\text{out\_features}, \text{in\_features})`. The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in\_features}}` bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. If :attr:`bias` is ``True``, the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where :math:`k = \frac{1}{\text{in\_features}}` Examples:: >>> import ms_adapter.pytorch as torch >>> import ms_adapter.pytorch.nn as nn >>> m = nn.Linear(20, 30) >>> input = torch.randn(128, 20) >>> output = m(input) >>> print(output.size()) torch.Size([128, 30]) """ def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): super(Linear, self).__init__() self.in_features = in_features self.out_features = out_features # self.matmul = P.MatMul(transpose_b=True) self.has_bias = False self.bias = None self.weight = Parameter(empty((self.out_features, self.in_features)), requires_grad=True) if bias: # self.bias_add = P.BiasAdd() self.bias = Parameter(empty(self.out_features), requires_grad=True) self.has_bias = True self.reset_parameters() unsupported_attr(device) unsupported_attr(dtype) def reset_parameters(self): # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.has_bias: fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 init.uniform_(self.bias, -bound, bound) def forward(self, input): return linear(input, self.weight, self.bias) def extra_repr(self): return 'in_features={}, out_features={}, bias={}'.format( self.in_features, self.out_features, self.has_bias is not None ) class Identity(Module): def __init__(self, *args, **kwargs): super(Identity, self).__init__() unsupported_attr(args) unsupported_attr(kwargs) self.identity = P.Identity() def forward(self, input): input = cast_to_ms_tensor(input) output = self.identity(input) return cast_to_adapter_tensor(output) LazyLinear = Linear class Bilinear(Module): def __init__(self, in1_features, in2_features, out_features, bias=True, device=None, dtype=None): super(Bilinear, self).__init__() self.in1_features = in1_features self.in2_features = in2_features self.out_features = out_features self.matmul = P.MatMul() self.mul = P.Mul() self.tile = P.Tile() self.reducesum = P.ReduceSum() self.has_bias = False self.weight = Parameter(empty((self.out_features, self.in1_features, self.in2_features)), requires_grad=True) if bias: self.bias_add = P.BiasAdd() self.bias = Parameter(empty(self.out_features), requires_grad=True) self.has_bias = True self.reset_parameters() unsupported_attr(device) unsupported_attr(dtype) def reset_parameters(self): bound = 1 / math.sqrt(self.weight.shape[1]) init.uniform_(self.weight, -bound, bound) if self.has_bias: init.uniform_(self.bias, -bound, bound) def forward(self, input1, input2): input1 = cast_to_ms_tensor(input1) input2 = cast_to_ms_tensor(input2) x = self.matmul(input1.reshape(-1, input1.shape[-1]), self.weight.permute(1, 0, 2).reshape(self.weight.shape[1], -1)) x = self.mul(x, self.tile(input2.reshape(-1, input2.shape[-1]), (1, self.out_features))) x = x.reshape(x.shape[0], self.out_features, -1) x = self.reducesum(x, -1) if self.has_bias: x = self.bias_add(x, self.bias) x = x.reshape(*input1.shape[:-1], -1) return cast_to_adapter_tensor(x) def extra_repr(self): return 'in1_features={}, in2_features={}, out_features={}, bias={}'.format( self.in1_features, self.in2_features, self.out_features, self.has_bias is not None ) msadapter/ms_adapter/pytorch/nn/modules/loss.py000066400000000000000000000305601436713554100222770ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import warnings import mindspore as ms import mindspore.nn as nn from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.tensor import Tensor, cast_to_ms_tensor, cast_to_adapter_tensor import ms_adapter.pytorch.nn.functional as F from .module import Module __all__ = [ 'SmoothL1Loss', 'L1Loss', 'MSELoss', 'CrossEntropyLoss', 'CTCLoss', 'NLLLoss', 'KLDivLoss', 'BCELoss', 'BCEWithLogitsLoss', 'HuberLoss', 'SoftMarginLoss', 'CosineEmbeddingLoss', 'MultiMarginLoss', 'TripletMarginLoss', 'PoissonNLLLoss', 'GaussianNLLLoss', 'HingeEmbeddingLoss', 'MarginRankingLoss', 'MultiLabelMarginLoss', 'MultiLabelSoftMarginLoss', 'TripletMarginWithDistanceLoss', ] class _Loss(Module): def __init__(self, size_average=None, reduce=None, reduction='mean'): super(_Loss, self).__init__() if size_average is not None or reduce is not None: self.reduction: str = self._get_reduce_string(size_average, reduce) else: self.reduction = reduction def _get_reduce_string(self, size_average, reduce): if size_average is None: size_average = True if reduce is None: reduce = True if size_average and reduce: ret = 'mean' elif reduce: ret = 'sum' else: ret = 'none' warning = "size_average and reduce args will be deprecated, please use reduction='{}' instead." warnings.warn(warning.format(ret)) return ret def cuda(self, *args, **kwargs): ## TODO No practical effect unsupported_attr(args) unsupported_attr(kwargs) return self class _WeightedLoss(_Loss): def __init__(self, weight=None, size_average=None, reduce=None, reduction='mean'): super(_WeightedLoss, self).__init__(size_average, reduce, reduction) if weight is not None: self.register_buffer('weight', weight) else: self.weight = None class SmoothL1Loss(_Loss): """Creates a criterion that uses a squared term if the absolute element-wise error falls below beta and an L1 term otherwise. """ def __init__(self, size_average=None, reduce=None, reduction='mean', beta=1.0): super(SmoothL1Loss, self).__init__(size_average, reduce, reduction) self.beta = beta def forward(self, logits, lables): return F.smooth_l1_loss(logits, lables, reduction=self.reduction, beta=self.beta) class L1Loss(_Loss): """Creates a criterion that measures the mean absolute error (MAE) between each element in the input and target. """ def forward(self, input, target): return F.l1_loss(input, target, reduction=self.reduction) class MSELoss(_Loss): """Creates a criterion that measures the mean squared error (squared L2 norm) between each element in the input and target. """ def forward(self, input, target): return F.mse_loss(input, target, reduction=self.reduction) class KLDivLoss(_Loss): """The Kullback-Leibler divergence loss.""" def __init__(self, size_average=None, reduce=None, reduction='mean', log_target=False): super(KLDivLoss, self).__init__(size_average, reduce, reduction) self.log_target = log_target def forward(self, input, target): return F.kl_div(input, target, reduction=self.reduction, log_target=self.log_target) class CrossEntropyLoss(_WeightedLoss): """ This criterion computes the cross entropy loss between input logits and target. """ def __init__(self, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean', label_smoothing=0.0): super(CrossEntropyLoss, self).__init__(weight, size_average, reduce, reduction) self.ignore_index = ignore_index self.label_smoothing = label_smoothing def forward(self, input, target): return F.cross_entropy(input, target, weight=self.weight, ignore_index=self.ignore_index, reduction=self.reduction, label_smoothing=self.label_smoothing) class NLLLoss(_WeightedLoss): """ The negative log likelihood loss. It is useful to train a classification problem with `C` classes. """ def __init__(self, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean'): super(NLLLoss, self).__init__(weight, size_average, reduce, reduction) self.ignore_index = ignore_index def forward(self, input, target): return F.nll_loss(input, target, weight=self.weight, ignore_index=self.ignore_index, reduction=self.reduction) class BCELoss(_WeightedLoss): """Creates a criterion that measures the Binary Cross Entropy between the target and the input probabilities: """ def forward(self, input, target): return F.binary_cross_entropy(input, target, weight=self.weight, reduction=self.reduction) class BCEWithLogitsLoss(_WeightedLoss): """This loss combines a `Sigmoid` layer and the `BCELoss` in one single class. This version is more numerically stable than using a plain `Sigmoid` followed by a `BCELoss` as, by combining the operations into one layer, we take advantage of the log-sum-exp trick for numerical stability. """ def __init__(self, weight=None, size_average=None, reduce=None, reduction='mean', pos_weight= None): super(BCEWithLogitsLoss, self).__init__(weight, size_average, reduce, reduction) if pos_weight is not None: self.register_buffer('pos_weight', pos_weight) else: self.pos_weight = pos_weight def forward(self, input, target): return F.binary_cross_entropy_with_logits(input, target, self.weight, pos_weight=self.pos_weight, reduction=self.reduction) class HuberLoss(_Loss): def __init__(self, reduction='mean', delta=1.0): super().__init__(reduction=reduction) self.delta = delta def forward(self, input, target): return F.huber_loss(input, target, reduction=self.reduction, delta=self.delta) class SoftMarginLoss(_Loss): def forward(self, input, target): return F.soft_margin_loss(input, target, reduction=self.reduction) class CosineEmbeddingLoss(_Loss): def __init__(self, margin=0., size_average=None, reduce=None, reduction='mean'): super(CosineEmbeddingLoss, self).__init__(size_average, reduce, reduction) self.margin = margin def forward(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor: return F.cosine_embedding_loss(input1, input2, target, margin=self.margin, reduction=self.reduction) class MultiMarginLoss(_WeightedLoss): # TODO: Until Now, mindspore do not support `margin` is 'float' type def __init__(self, p=1, margin=1, weight=None, size_average=None, reduce=None, reduction: str='mean'): super(MultiMarginLoss, self).__init__(weight, size_average, reduce, reduction) if p not in (1, 2): raise ValueError("only p == 1 and p == 2 supported") assert weight is None or weight.dim() == 1 self.p = p self.margin = margin def forward(self, input: Tensor, target: Tensor) -> Tensor: return F.multi_margin_loss(input, target, p=self.p, margin=self.margin, weight=self.weight, reduction=self.reduction) class TripletMarginLoss(_Loss): def __init__(self, margin=1.0, p=2., eps=1e-6, swap=False, size_average=None, reduce=None, reduction: str='mean'): super(TripletMarginLoss, self).__init__(size_average, reduce, reduction) self.margin = margin self.p = p self.eps = eps self.swap = swap def forward(self, anchor, positive, negative): return F.triplet_margin_loss(anchor, positive, negative, margin=self.margin, p=self.p, eps=self.eps, swap=self.swap, reduction=self.reduction) class PoissonNLLLoss(_Loss): def __init__(self, log_input=True, full=False, size_average=None, eps=1e-8, reduce=None, reduction='mean'): super(PoissonNLLLoss, self).__init__(size_average, reduce, reduction) self.log_input = log_input self.full = full self.eps = eps def forward(self, log_input, target): return F.poisson_nll_loss(log_input, target, log_input=self.log_input, full=self.full, eps=self.eps, reduction=self.reduction) class GaussianNLLLoss(_Loss): def __init__(self, *, full=False, eps=1e-6, reduction='mean'): super(GaussianNLLLoss, self).__init__(None, None, reduction) self.full = full self.eps = eps self.gaussian_nll_loss = nn.GaussianNLLLoss(full=self.full, eps=self.eps, reduction=self.reduction) def forward(self, input, target, var): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) var = cast_to_ms_tensor(var) out = self.gaussian_nll_loss(input, target, var) return cast_to_adapter_tensor(out) class MarginRankingLoss(_Loss): def __init__(self, margin=0., size_average=None, reduce=None, reduction='mean'): super(MarginRankingLoss, self).__init__(size_average, reduce, reduction) self.margin = margin self.margin_ranking_loss = nn.MarginRankingLoss(self.margin, self.reduction) def forward(self, input1, input2, target): return self.margin_ranking_loss(input1, input2, target) class HingeEmbeddingLoss(_Loss): def __init__(self, margin=1.0, size_average=None, reduce=None, reduction='mean'): super(HingeEmbeddingLoss, self).__init__(size_average, reduce, reduction) self.margin = margin self.hinge_embedding_loss = nn.HingeEmbeddingLoss(margin=self.margin, reduction=self.reduction) def forward(self, input, target): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) out = self.hinge_embedding_loss(input, target) return cast_to_adapter_tensor(out) class MultiLabelMarginLoss(_Loss): def __init__(self, size_average=None, reduce=None, reduction='mean'): super(MultiLabelMarginLoss, self).__init__(size_average, reduce, reduction) self.multilabel_margin_loss = ms.ops.MultilabelMarginLoss(reduction=self.reduction) def forward(self, input, target): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) if target.dtype == ms.int64: target = target.astype(ms.int32) out = self.multilabel_margin_loss(input, target) return cast_to_adapter_tensor(out) class MultiLabelSoftMarginLoss(_WeightedLoss): def __init__(self, weight=None, size_average=None, reduce=None, reduction='mean'): super(MultiLabelSoftMarginLoss, self).__init__(weight, size_average, reduce, reduction) self.multilabel_soft_margin_loss = nn.MultiLabelSoftMarginLoss(weight=self.weight, reduction=self.reduction) def forward(self, input, target): input = cast_to_ms_tensor(input) target = cast_to_ms_tensor(target) out = self.multilabel_soft_margin_loss(input, target) return cast_to_adapter_tensor(out) class TripletMarginWithDistanceLoss(_Loss): def __init__(self, *, distance_function=None, margin: float = 1.0, swap: bool = False, reduction: str = 'mean'): super(TripletMarginWithDistanceLoss, self).__init__(size_average=None, reduce=None, reduction=reduction) self.distance_function = distance_function self.margin = margin self.swap = swap def forward(self, anchor, positive, negative): return F.triplet_margin_with_distance_loss(anchor, positive, negative, distance_function=self.distance_function, margin=self.margin, swap=self.swap, reduction=self.reduction) class CTCLoss(_Loss): def __init__(self, blank=0, reduction='mean', zero_infinity=False): super(CTCLoss, self).__init__(reduction=reduction) self.blank = blank self.zero_infinity = zero_infinity def forward(self, log_probs, targets, input_lengths, target_lengths): return F.ctc_loss(log_probs, targets, input_lengths, target_lengths, self.blank, self.reduction, self.zero_infinity) msadapter/ms_adapter/pytorch/nn/modules/module.py000066400000000000000000000266131436713554100226100ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from collections import OrderedDict from mindspore.nn import Cell from mindspore import Tensor as ms_Tensor from ms_adapter.pytorch.tensor import tensor from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.common.device import Device __all__ = ['Module'] class Module(Cell): def __init__(self, auto_prefix=True, flags=None): super(Module, self).__init__(auto_prefix, flags) self.training = True def __del__(self): pass def __repr__(self): extra_str = self.extra_repr() info_str = self.__class__.__name__ + '(' if self._cells: sub_str = '\n' if extra_str: sub_str += '{}\n'.format(self.extra_repr()) for key, value in self._cells.items(): sub_str += ' ({}): {}\n'.format(key, repr(value)) sub_str = sub_str.replace('\n', '\n') + ')' info_str += sub_str else: info_str += extra_str + ')' return info_str def extra_repr(self): r"""Set the extra representation of the module""" return '' def construct(self, *inputs, **kwargs): return self.forward(*inputs, **kwargs) def _run_construct(self, cast_inputs, kwargs): """Run the construct function""" if self._enable_forward_pre_hook: cast_inputs = self._run_forward_pre_hook(cast_inputs) if self._enable_backward_hook: output = self._backward_hook_construct(*cast_inputs) elif hasattr(self, "_shard_fn"): output = self._shard_fn(*cast_inputs, **kwargs) else: output = self.construct(*cast_inputs, **kwargs) if self._enable_forward_hook: output = self._run_forward_hook(cast_inputs, output) return output def forward(self, *inputs, **kwargs): raise NotImplementedError("The forward method must be implemented by inherited class") def train(self, mode=True): self.set_train(mode) return self def eval(self): self.set_train(False) return self def modules(self): result = [] cells_names = self.cells_and_names() for m in cells_names: result.append(m[1]) return iter(result) def named_modules(self, memo=None, prefix='', remove_duplicate=None): if memo: raise NotImplementedError("For named_modules, `memo` not implemented.") unsupported_attr(remove_duplicate) output = self.cells_and_names(name_prefix=prefix) return output def _parameters_and_names(self, name_prefix='', expand=True): cells = [] if expand: cells = self.cells_and_names(name_prefix=name_prefix) else: cells.append((name_prefix, self)) params_set = set() for cell_name, cell in cells: params = cell._params.items() for par_name, par in params: if par.inited_param is not None: par = par.inited_param if par is not None and id(par) not in params_set: params_set.add(id(par)) par_new_name = par_name if cell_name: par_new_name = cell_name + '.' + par_new_name # TODO Update parameter names to avoid duplicates par.name = par_new_name yield par_new_name, par def add_module(self, name, module): if not isinstance(module, Module) and module is not None: raise TypeError("{} is not a Module subclass".format( module.__name__)) elif hasattr(self, name) and name not in self._cells: raise KeyError("attribute '{}' already exists".format(name)) elif '.' in name: raise KeyError("module name can't contain \".\", got: {}".format(name)) elif name == '': raise KeyError("module name can't be empty string \"\"") self._cells[name] = module def register_module(self, name, module): """Alias for :func:`add_module`.""" self.add_module(name, module) def named_parameters(self, prefix='', recurse=True): return self._parameters_and_names(prefix, recurse) def parameters_and_names(self, name_prefix='', expand=True): return self._parameters_and_names(name_prefix=name_prefix, expand=expand) def named_children(self): r"""Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself. Yields: (string, Module): Tuple containing a name and child module Example:: >>> for name, module in model.named_children(): >>> if name in ['conv4', 'conv5']: >>> print(module) """ memo = set() for name, module in self._cells.items(): if module is not None and module not in memo: memo.add(module) yield name, module def children(self): r"""Returns an iterator over immediate children modules. Yields: Module: a child module """ for _, module in self.named_children(): yield module def apply(self, fn=None): r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``) as well as self. Typical use includes initializing the parameters of a model (see also :ref:`nn-init-doc`). Args: fn (:class:`Module` -> None): function to be applied to each submodule Returns: Module: self Example:: >>> def init_weights(m): >>> print(m) >>> if type(m) == nn.Linear: >>> m.weight.fill_(1.0) >>> print(m.weight) >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)) >>> net.apply(init_weights) """ for module in self.children(): module.apply(fn) fn(self) return self def parameters(self, recurse = True): for _, param in self.named_parameters(recurse=recurse): yield param def state_dict(self, destination=None, prefix='', keep_vars=False): unsupported_attr(keep_vars) unsupported_attr(prefix) if destination is None: destination = OrderedDict() for n, v in self.named_parameters(): destination[n] = tensor(v) return destination def register_buffer(self, name, tensor, persistent=True): r"""Adds a buffer to the module. This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's ``running_mean`` is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:`persistent` to ``False``. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:`state_dict`. Buffers can be accessed as attributes using given names. Args: name (string): name of the buffer. The buffer can be accessed from this module using the given name tensor (Tensor or None): buffer to be registered. If ``None``, then operations that run on buffers, such as :attr:`cuda`, are ignored. If ``None``, the buffer is **not** included in the module's :attr:`state_dict`. persistent (bool): whether the buffer is part of this module's :attr:`state_dict`. """ unsupported_attr(persistent) if '_params' not in self.__dict__: raise AttributeError("cannot assign buffer before Module.__init__() call.") elif not isinstance(name, str): raise TypeError("buffer name should be a string. " "Got {}".format(type(name))) elif '.' in name: raise KeyError("buffer name can't contain \".\"") elif name == '': raise KeyError("buffer name can't be empty string \"\"") elif hasattr(self, name) and name not in self._params: raise KeyError("attribute '{}' already exists".format(name)) elif tensor is not None and not isinstance(tensor, ms_Tensor): raise TypeError("cannot assign '{}' object to buffer '{}' " "(Tensor or None required)" .format(type(tensor), name)) else: self._params[name] = Parameter(tensor, name=name, requires_grad=False) def to(self, *args, **kwargs): # TODO: # Note that this API requires the user to ensure the correctness of the input currently, # and only the function of modifying device is available. args_len = len(args) kwargs_len = len(kwargs) if args_len == 0 and kwargs_len == 0: raise ValueError("Module.to is missing inputs, please check.") elif (args_len + kwargs_len > 1) or (kwargs_len > 0 and "device" not in kwargs): raise ValueError("Currently only the function of modifying device is available.") elif (args_len > 0 and not isinstance(args[0], (str, Device))) or \ (kwargs_len > 0 and not isinstance(kwargs.get("device"), (str, Device))): raise ValueError("Currently only the function of modifying device is available, " "which via a string or torch.device.") return self def register_parameter(self, name, param): """Adds a parameter to the module. The parameter can be accessed as an attribute using given name. Args: name (string): name of the parameter. The parameter can be accessed from this module using the given name param (Parameter or None): parameter to be added to the module. If ``None``, then operations that run on parameters, such as :attr:`cuda`, are ignored. If ``None``, the parameter is **not** included in the module's :attr:`state_dict`. """ if '_params' not in self.__dict__: raise AttributeError("cannot assign parameter before Module.__init__() call") elif not isinstance(name, str): raise TypeError("parameter name should be a string. Got {}".format(type(name))) elif '.' in name: raise KeyError("parameter name can't contain \".\"") elif name == '': raise KeyError("parameter name can't be empty string \"\"") elif hasattr(self, name) and name not in self._params: raise KeyError("attribute '{}' already exists".format(name)) if param is None: self._params[name] = None elif not isinstance(param, Parameter): raise TypeError("cannot assign '{}' object to parameter '{}' " "(nn.Parameter or None required)" .format(type(param), name)) else: self._params[name] = param def cuda(self, device): unsupported_attr(device) return self msadapter/ms_adapter/pytorch/nn/modules/normalization.py000066400000000000000000000117151436713554100242060ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import mindspore.ops.functional as F import mindspore.ops.operations as P from mindspore.ops._primitive_cache import _get_cache_prim from ms_adapter.pytorch.tensor import cast_to_adapter_tensor, cast_to_ms_tensor from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.nn import init import ms_adapter.pytorch.functional as torch_func import ms_adapter.pytorch.nn.functional as torch_nn_func from .module import Module __all__ = ['LayerNorm', 'GroupNorm', 'LocalResponseNorm'] class LayerNorm(Module): def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True, device=None, dtype=None): unsupported_attr(device) unsupported_attr(dtype) if not isinstance(normalized_shape, (list, tuple, int)): raise TypeError("`normalized_shape` should be in type of `list`, `tuple`, `int`" "but got {}".format(type(normalized_shape))) super(LayerNorm,self).__init__() self.eps = eps self.normalized_shape = normalized_shape self.elementwise_affine = elementwise_affine self.weight = Parameter(torch_func.empty(self.normalized_shape)) self.bias = Parameter(torch_func.empty(self.normalized_shape)) if isinstance(self.normalized_shape, int): self.normalized_shape_rank = 1 else: self.normalized_shape_rank = len(self.normalized_shape) if self.elementwise_affine: self.weight.requires_grad = True self.bias.requires_grad = True else: self.weight.requires_grad = False self.bias.requires_grad = False self.reset_parameters() def reset_parameters(self): init.ones_(self.weight) init.zeros_(self.bias) def forward(self, input): input = cast_to_ms_tensor(input) begin_axis = ms.ops.rank(input) - self.normalized_shape_rank layer_norm_ops = _get_cache_prim(ms.ops.LayerNorm)(begin_norm_axis=begin_axis, begin_params_axis=begin_axis, epsilon=self.eps) output_x, _, _ = layer_norm_ops(input, self.weight, self.bias) return cast_to_adapter_tensor(output_x) def extra_repr(self): return "{normalized_shape}, eps={eps}, "\ "elementwise_affine={elementwise_affine}".format(**self.__dict__) class GroupNorm(Module): def __init__(self, num_groups, num_channels, eps=1e-5, affine=True, device=None, dtype=None): unsupported_attr(device) unsupported_attr(dtype) super(GroupNorm, self).__init__() self.num_groups = num_groups self.num_channels = num_channels self.eps = eps self.affine = affine self.weight = Parameter(torch_func.empty(num_channels), requires_grad=self.affine) self.bias = Parameter(torch_func.empty(num_channels), requires_grad=self.affine) self.shape = F.shape self.size = F.size self.reshape = F.reshape self.reduce_mean = P.ReduceMean(keep_dims=True) self.square = F.square self.reduce_sum = P.ReduceSum(keep_dims=True) self.sqrt = P.Sqrt() self.reset_parameters() def reset_parameters(self): init.ones_(self.weight) init.zeros_(self.bias) def _cal_output(self, x): """calculate groupnorm output""" x = cast_to_adapter_tensor(x) shape = self.shape(x) batch = shape[0] x = self.reshape(x, (batch, self.num_groups, -1)) mean = self.reduce_mean(x, 2) tmp = self.size(x) // batch var = self.reduce_sum(self.square(x - mean), 2) / (tmp / self.num_groups) std = self.sqrt(var + self.eps) x = (x - mean) / std x = self.reshape(x, shape) ndim = len(shape) param_shape = [1] * (ndim - 1) param_shape[0] = -1 param_shape = tuple(param_shape) output = x * self.reshape(self.weight, param_shape) + self.reshape(self.bias, param_shape) output = cast_to_adapter_tensor(output) return output def forward(self, input): output = self._cal_output(input) return output def extra_repr(self): return '{num_groups}, {num_channels}, eps={eps}, ' \ 'affine={affine}'.format(**self.__dict__) class LocalResponseNorm(Module): def __init__(self, size, alpha=1e-4, beta=0.75, k=1.): super(LocalResponseNorm, self).__init__() self.size = size self.alpha = alpha self.beta = beta self.k = k def forward(self, input): return torch_nn_func.local_response_norm(input, self.size, self.alpha, self.beta, self.k) def extra_repr(self): return '{size}, alpha={alpha}, beta={beta}, k={k}'.format(**self.__dict__) msadapter/ms_adapter/pytorch/nn/modules/padding.py000066400000000000000000000241441436713554100227260ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from mindspore import nn from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from .module import Module __all__ = ['ConstantPad1d', 'ConstantPad2d', 'ConstantPad3d', 'ReflectionPad1d', 'ReflectionPad2d', 'ReflectionPad3d', 'ZeroPad2d', 'ReplicationPad1d', 'ReplicationPad2d', 'ReplicationPad3d'] class _ConstantPadNd(Module): def __init__(self, padding, value): super(_ConstantPadNd, self).__init__() self.padding = padding self.value = value self.pad_fun = None def forward(self, input): input = cast_to_ms_tensor(input) output = self.pad_fun(input) return cast_to_adapter_tensor(output) def extra_repr(self) -> str: return 'padding={}, value={}'.format(self.padding, self.value) class ConstantPad1d(_ConstantPadNd): r"""Pads the input tensor boundaries with a constant value. For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. Args: padding (int, tuple): the size of the padding. If is `int`, uses the same padding in both boundaries. If a 2-`tuple`, uses (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`) Shape: - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`. - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` Examples:: >>> m = nn.ConstantPad1d(2, 3.5) >>> input = ms_adapter.pytorch.ones(1, 2, 4) >>> m(input) """ def __init__(self, padding, value: float): super(ConstantPad1d, self).__init__(padding, value) self.pad_fun = nn.ConstantPad1d(self.padding, self.value) class ConstantPad2d(_ConstantPadNd): r"""Pads the input tensor boundaries with a constant value. For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. Args: padding (int, tuple): the size of the padding. If is `int`, uses the same padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`) Shape: - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`. - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}` :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` Examples:: >>> m = nn.ConstantPad2d(2, 3.5) >>> input = ms_adapter.pytorch.ones(1, 2, 2) >>> m(input) """ def __init__(self, padding, value): super(ConstantPad2d, self).__init__(padding, value) self.pad_fun = nn.ConstantPad2d(self.padding, self.value) class ConstantPad3d(_ConstantPadNd): r"""Pads the input tensor boundaries with a constant value. For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. Args: padding (int, tuple): the size of the padding. If is `int`, uses the same padding in all boundaries. If a 6-`tuple`, uses (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`, :math:`\text{padding\_front}`, :math:`\text{padding\_back}`) Shape: - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`. - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`, where :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}` :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}` :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` Examples:: >>> m = nn.ConstantPad3d(3, 3.5) >>> input = ms_adapter.pytorch.ones(16, 3, 10, 20, 30) >>> output = m(input) """ def __init__(self, padding, value): super(ConstantPad3d, self).__init__(padding, value) self.pad_fun = nn.ConstantPad3d(self.padding, self.value) class _ReflectionPadNd(Module): def __init__(self, padding): super(_ReflectionPadNd, self).__init__() self.padding = padding self.pad_fun = None def forward(self, input): input = cast_to_ms_tensor(input) output = self.pad_fun(input) return cast_to_adapter_tensor(output) def extra_repr(self) -> str: return '{}'.format(self.padding) class ReflectionPad1d(_ReflectionPadNd): r"""Pads the input tensor using the reflection of the input boundary. For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. Args: padding (int, tuple): the size of the padding. If is `int`, uses the same padding in all boundaries. If a 2-`tuple`, uses (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`) Shape: - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`. - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` Examples:: >>> m = nn.ReflectionPad1d(2) >>> input = ms_adapter.pytorch.ones(1, 2, 4) >>> m(input) """ def __init__(self, padding): super(ReflectionPad1d, self).__init__(padding) self.pad_fun = nn.ReflectionPad1d(self.padding) class ReflectionPad2d(_ReflectionPadNd): r"""Pads the input tensor using the reflection of the input boundary. For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. Args: padding (int, tuple): the size of the padding. If is `int`, uses the same padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`) Shape: - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`. - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})` where :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}` :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` Examples:: >>> m = nn.ReflectionPad2d(2) >>> input = ms_adapter.pytorch.ones(1, 1, 3, 3) >>> m(input) """ def __init__(self, padding): super(ReflectionPad2d, self).__init__(padding) self.pad_fun = nn.ReflectionPad2d(self.padding) def forward(self, input): input = cast_to_ms_tensor(input) output = self.pad_fun(input) return cast_to_adapter_tensor(output) class ReflectionPad3d(_ReflectionPadNd): r"""Pads the input tensor using the reflection of the input boundary. For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. Args: padding (int, tuple): the size of the padding. If is `int`, uses the same padding in all boundaries. If a 6-`tuple`, uses (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`, :math:`\text{padding\_front}`, :math:`\text{padding\_back}`) Shape: - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`. - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`, where :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}` :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}` :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` Examples:: >>> m = nn.ReflectionPad3d(1) >>> input = ms_adapter.pytorch.ones(1, 1, 2, 2, 2) >>> m(input) """ def __init__(self, padding): super(ReflectionPad3d, self).__init__(padding) self.pad_fun = nn.ReflectionPad3d(self.padding) def forward(self, input): input = cast_to_ms_tensor(input) output = self.pad_fun(input) return cast_to_adapter_tensor(output) class ZeroPad2d(ConstantPad2d): r"""Pads the input tensor boundaries with zero. For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. Args: padding (int, tuple): the size of the padding. If is `int`, uses the same padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`) Shape: - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`. - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}` :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` Examples:: >>> m = nn.ZeroPad2d(2) >>> input = ms_adapter.pytorch.ones(1, 1, 3, 3) >>> m(input) """ def __init__(self, padding) -> None: super(ZeroPad2d, self).__init__(padding, 0.) class _ReplicationPadNd(Module): def __init__(self, padding): super(_ReplicationPadNd, self).__init__() self.padding = padding self.pad_fun = None def forward(self, input): input = cast_to_ms_tensor(input) output = self.pad_fun(input) return cast_to_adapter_tensor(output) def extra_repr(self) -> str: return '{}'.format(self.padding) class ReplicationPad1d(_ReplicationPadNd): def __init__(self, padding): super(ReplicationPad1d, self).__init__(padding) self.pad_fun = nn.ReplicationPad1d(padding) class ReplicationPad2d(_ReplicationPadNd): def __init__(self, padding): super(ReplicationPad2d, self).__init__(padding) self.pad_fun = nn.ReplicationPad2d(padding) class ReplicationPad3d(_ReplicationPadNd): def __init__(self, padding): super(ReplicationPad3d, self).__init__(padding) self.pad_fun = nn.ReplicationPad3d(padding) msadapter/ms_adapter/pytorch/nn/modules/pixel_shuffle.py000066400000000000000000000015641436713554100241560ustar00rootroot00000000000000from ms_adapter.pytorch.nn.modules.module import Module from ms_adapter.pytorch.nn.functional import pixel_shuffle, pixel_unshuffle __all__ = ['PixelShuffle', 'PixelUnshuffle'] class PixelShuffle(Module): def __init__(self, upscale_factor): super(PixelShuffle, self).__init__() self.upscale_factor = upscale_factor def forward(self, input): return pixel_shuffle(input, self.upscale_factor) def extra_repr(self) -> str: return 'upscale_factor={}'.format(self.upscale_factor) class PixelUnshuffle(Module): def __init__(self, downscale_factor): super(PixelUnshuffle, self).__init__() self.downscale_factor = downscale_factor def forward(self, input): return pixel_unshuffle(input, self.downscale_factor) def extra_repr(self) -> str: return 'downscale_factor={}'.format(self.downscale_factor) msadapter/ms_adapter/pytorch/nn/modules/pooling.py000066400000000000000000000406361436713554100227730ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from typing import Iterable import mindspore as ms from mindspore.ops import operations as P from mindspore.ops import functional as F from mindspore.ops._primitive_cache import _get_cache_prim import ms_adapter.pytorch.nn.functional as Adapter_F from ms_adapter.utils import unsupported_attr, is_under_ascend_context from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from .module import Module __all__ = ['MaxPool1d', 'MaxPool2d', 'MaxPool3d', 'AvgPool1d', 'AvgPool2d', 'AvgPool3d', 'AdaptiveAvgPool1d', 'AdaptiveAvgPool2d', 'AdaptiveAvgPool3d', 'AdaptiveMaxPool1d', 'AdaptiveMaxPool2d', 'AdaptiveMaxPool3d', 'LPPool1d', 'LPPool2d', 'FractionalMaxPool2d', 'FractionalMaxPool3d'] class _MaxPoolNd(Module): def __init__(self, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False): super(_MaxPoolNd, self).__init__() self.kernel_size = kernel_size self.stride = stride if (stride is not None) else kernel_size self.padding = padding self.dilation = dilation self.return_indices = return_indices self.ceil_mode = ceil_mode def extra_repr(self): return 'kernel_size={kernel_size}, stride={stride}, padding={padding}' \ ', dilation={dilation}, ceil_mode={ceil_mode}'.format(**self.__dict__) class MaxPool1d(_MaxPoolNd): def forward(self, input): input = cast_to_ms_tensor(input) output = Adapter_F.max_pool1d(input, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, self.return_indices) return cast_to_adapter_tensor(output) class MaxPool2d(_MaxPoolNd): def __init__(self, kernel_size=1, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False): super(MaxPool2d, self).__init__(kernel_size, stride, padding, dilation, return_indices, ceil_mode) def forward(self, input): return Adapter_F.max_pool2d(input, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, self.return_indices) class MaxPool3d(_MaxPoolNd): def __init__(self, kernel_size=1, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False): super(MaxPool3d, self).__init__(kernel_size, stride, padding, dilation, return_indices, ceil_mode) def forward(self, input): return Adapter_F.max_pool3d(input, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, self.return_indices) class _AvgPoolNd(Module): def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override = None): super(_AvgPoolNd, self).__init__() self.kernel_size = kernel_size self.stride = stride if (stride is not None) else kernel_size self.padding = padding self.ceil_mode = ceil_mode self.count_include_pad = count_include_pad self.divisor_override = divisor_override unsupported_attr(ceil_mode) unsupported_attr(count_include_pad) unsupported_attr(divisor_override) if ceil_mode is True or count_include_pad is False or divisor_override is not None: raise NotImplementedError("These parameters cannot be set now.") def extra_repr(self): return 'kernel_size={}, stride={}, padding={}'.format( self.kernel_size, self.stride, self.padding ) class AvgPool1d(_AvgPoolNd): def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True): super(AvgPool1d, self).__init__(kernel_size, stride, padding, ceil_mode, count_include_pad) self.padding = padding self.kernel_size = kernel_size def forward(self, input): return Adapter_F.avg_pool1d(input, kernel_size=self.kernel_size, stride=self.stride, padding=self.padding) class AvgPool2d(_AvgPoolNd): def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None): super(AvgPool2d, self).__init__(kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) self.padding = padding self.kernel_size = kernel_size def forward(self, input): return Adapter_F.avg_pool2d(input, kernel_size=self.kernel_size, stride=self.stride, padding=self.padding) class AvgPool3d(_AvgPoolNd): def forward(self, input): return Adapter_F.avg_pool3d(input, kernel_size=self.kernel_size, stride=self.stride, padding=self.padding, ceil_mode=self.ceil_mode) class _AdaptiveAvgPoolNd(Module): def __init__(self, output_size): super(_AdaptiveAvgPoolNd, self).__init__() self.output_size = output_size def extra_repr(self): return 'output_size={}'.format(self.output_size) class AdaptiveAvgPool1d(_AdaptiveAvgPoolNd): def __init__(self, output_size): """Initialize AdaptiveMaxPool1d.""" super(AdaptiveAvgPool1d, self).__init__(output_size) self.expand = P.ExpandDims() self.squeeze = P.Squeeze(2) self.output_size = output_size self.shape = F.shape def construct(self, input): input = cast_to_ms_tensor(input) _, _, width = self.shape(input) stride = width // self.output_size kernel_size = width - (self.output_size - 1) * stride stride = (1, width // self.output_size) kernel_size = (1, kernel_size) max_pool = _get_cache_prim(P.AvgPool)(kernel_size=kernel_size, strides=stride, pad_mode="valid", data_format="NCHW") input = self.expand(input, 2) x = max_pool(input) x = self.squeeze(x) return cast_to_adapter_tensor(x) class AdaptiveAvgPool2d(_AdaptiveAvgPoolNd): def __init__(self, output_size): super(AdaptiveAvgPool2d, self).__init__(output_size) self.output_size = output_size self.shape = P.Shape() if not isinstance(self.output_size, Iterable): self.output_size = [self.output_size, ] * 2 self.condition = [0,] * 2 if None in self.output_size: self.output_size = list(self.output_size) if self.output_size[0] is None: self.condition [0] = 1 self.output_size[0] = 0 if self.output_size[1] is None: self.condition [1] = 1 self.output_size[1] = 0 def forward(self, input): input = cast_to_ms_tensor(input) _, _, h, w = self.shape(input) out_h = self.output_size[0] + self.condition[0] * h out_w = self.output_size[1] + self.condition[1] * w stride_h = h // out_h kernel_h = h - (out_h - 1) * stride_h stride_w = w // out_w kernel_w = w - (out_w - 1) * stride_w avg_pool = _get_cache_prim(P.AvgPool)( kernel_size=(kernel_h, kernel_w), strides=(stride_h, stride_w), pad_mode="valid", data_format="NCHW" ) outputs = avg_pool(input) return cast_to_adapter_tensor(outputs) class AdaptiveAvgPool3d(_AdaptiveAvgPoolNd): def __init__(self, output_size): super(AdaptiveAvgPool3d, self).__init__(output_size) self.output_size = output_size self.shape = P.Shape() if not isinstance(self.output_size, Iterable): self.output_size = [self.output_size, ] * 3 self.condition = [0,] * 3 if None in self.output_size: self.output_size = list(self.output_size) if self.output_size[0] is None: self.condition [0] = 1 self.output_size[0] = 0 if self.output_size[1] is None: self.condition [1] = 1 self.output_size[1] = 0 if self.output_size[2] is None: self.condition[2] = 1 self.output_size[2] = 0 def forward(self, input): input = cast_to_ms_tensor(input) _, _, d, h, w = self.shape(input) out_d = self.output_size[0] + self.condition[0] * d out_h = self.output_size[1] + self.condition[1] * h out_w = self.output_size[2] + self.condition[2] * w stride_d = d // out_d kernel_d = d - (out_d - 1) * stride_d stride_h = h // out_h kernel_h = h - (out_h - 1) * stride_h stride_w = w // out_w kernel_w = w - (out_w - 1) * stride_w avg_pool = _get_cache_prim(P.AvgPool3D)(kernel_size=(kernel_d, kernel_h, kernel_w), strides=(stride_d, stride_h, stride_w), pad_mode="valid", data_format="NCDHW") outputs = avg_pool(input) return cast_to_adapter_tensor(outputs) class _AdaptiveMaxPoolNd(Module): def __init__(self, output_size, return_indices = False): super(_AdaptiveMaxPoolNd, self).__init__() self.output_size = output_size self.return_indices = return_indices def extra_repr(self) -> str: return 'output_size={}'.format(self.output_size) class AdaptiveMaxPool1d(_AdaptiveMaxPoolNd): def __init__(self, output_size, return_indices = False): """Initialize AdaptiveMaxPool1d.""" super(AdaptiveMaxPool1d, self).__init__(output_size, return_indices) self.expand = P.ExpandDims() self.squeeze = P.Squeeze(2) self.output_size = output_size self.shape = F.shape self.return_indices = return_indices def construct(self, input): input = cast_to_ms_tensor(input) _, _, width = self.shape(input) stride = width // self.output_size kernel_size = width - (self.output_size - 1) * stride stride = (1, width // self.output_size) kernel_size = (1, kernel_size) if self.return_indices: max_pool = P.MaxPoolWithArgmax(kernel_size=kernel_size, strides=stride, pad_mode='valid', data_format="NCHW") x = self.expand(input, 2) x, idx = max_pool(x) x = self.squeeze(x) # TODO: to avoid ascend not return ms.int32 but ms.uint16 idx = idx.astype(ms.int32) idx = self.squeeze(idx) return cast_to_adapter_tensor((x, idx)) else: max_pool = P.MaxPool(kernel_size=kernel_size, strides=stride, pad_mode="valid", data_format="NCHW") x = self.expand(input, 2) x = max_pool(x) x = self.squeeze(x) return cast_to_adapter_tensor(x) class AdaptiveMaxPool2d(_AdaptiveMaxPoolNd): def forward(self, input): input = cast_to_ms_tensor(input) if is_under_ascend_context() and len(input.shape) == 3: input = ms.ops.expand_dims(input, 0) output = ms.ops.adaptive_max_pool2d(input, self.output_size, self.return_indices) if self.return_indices: output[0] = ms.ops.squeeze(output[0], 0) output[1] = ms.ops.squeeze(output[1], 1) else: output = ms.ops.squeeze(output, 0) else: output = ms.ops.adaptive_max_pool2d(input, self.output_size, self.return_indices) return cast_to_adapter_tensor(output) class AdaptiveMaxPool3d(_AdaptiveMaxPoolNd): def __init__(self, output_size, return_indices = False): super(AdaptiveMaxPool3d, self).__init__(output_size, return_indices) self.output_size = output_size self.shape = P.Shape() if not isinstance(self.output_size, Iterable): self.output_size = [self.output_size, ] * 3 self.condition = [0,] * 3 if None in self.output_size: self.output_size = list(self.output_size) if self.output_size[0] is None: self.condition [0] = 1 self.output_size[0] = 0 if self.output_size[1] is None: self.condition [1] = 1 self.output_size[1] = 0 if self.output_size[2] is None: self.condition[2] = 1 self.output_size[2] = 0 if return_indices: raise NotImplementedError('AdaptiveMaxPool3d doesn\'t support return_indices now.') def forward(self, input): input = cast_to_ms_tensor(input) _, _, d, h, w = self.shape(input) out_d = self.output_size[0] + self.condition[0] * d out_h = self.output_size[1] + self.condition[1] * h out_w = self.output_size[2] + self.condition[2] * w stride_d = d // out_d kernel_d = d - (out_d - 1) * stride_d stride_h = h // out_h kernel_h = h - (out_h - 1) * stride_h stride_w = w // out_w kernel_w = w - (out_w - 1) * stride_w avg_pool = P.MaxPool3D(kernel_size=(kernel_d, kernel_h, kernel_w), strides=(stride_d, stride_h, stride_w), pad_mode="valid", data_format="NCDHW") outputs = avg_pool(input) return cast_to_adapter_tensor(outputs) class _LPPoolNd(Module): def __init__(self, norm_type, kernel_size, stride = None, ceil_mode = False): super(_LPPoolNd, self).__init__() self.norm_type = norm_type self.kernel_size = kernel_size self.stride = stride self.ceil_mode = ceil_mode def extra_repr(self): return 'norm_type={norm_type}, kernel_size={kernel_size}, stride={stride}, ' \ 'ceil_mode={ceil_mode}'.format(**self.__dict__) class LPPool1d(_LPPoolNd): def forward(self, input): return Adapter_F.lp_pool1d(input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode) class LPPool2d(_LPPoolNd): def forward(self, input): return Adapter_F.lp_pool2d(input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode) class FractionalMaxPool2d(Module): def __init__(self, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None): super(FractionalMaxPool2d, self).__init__() self.kernel_size = kernel_size self.return_indices = return_indices self.output_size = output_size self.output_ratio = output_ratio self._random_samples = _random_samples if output_size is None and output_ratio is None: raise ValueError("FractionalMaxPool2d requires specifying either " "an output size, or a pooling ratio") if output_size is not None and output_ratio is not None: raise ValueError("only one of output_size and output_ratio may be specified") if self.output_ratio is not None: if not (0 < self.output_ratio[0] < 1 and 0 < self.output_ratio[1] < 1): raise ValueError("output_ratio must be between 0 and 1 (got {})" .format(output_ratio)) def forward(self, input): return Adapter_F.fractional_max_pool2d(input, self.kernel_size, self.output_size, self.output_ratio, self.return_indices, self._random_samples) class FractionalMaxPool3d(Module): def __init__(self, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None): super(FractionalMaxPool3d, self).__init__() self.kernel_size = kernel_size self.return_indices = return_indices self.output_size = output_size self.output_ratio = output_ratio self._random_samples = _random_samples if output_size is None and output_ratio is None: raise ValueError("FractionalMaxPool3d requires specifying either " "an output size, or a pooling ratio") if output_size is not None and output_ratio is not None: raise ValueError("only one of output_size and output_ratio may be specified") if self.output_ratio is not None: if not (0 < self.output_ratio[0] < 1 and 0 < self.output_ratio[1] < 1): raise ValueError("output_ratio must be between 0 and 1 (got {})" .format(output_ratio)) def forward(self, input): return Adapter_F.fractional_max_pool3d(input, self.kernel_size, self.output_size, self.output_ratio, self.return_indices, self._random_samples) msadapter/ms_adapter/pytorch/nn/modules/rnn.py000066400000000000000000000522231436713554100221140ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numbers import warnings import math import mindspore as ms from mindspore.nn.layer.rnns import _DynamicRNNRelu, _DynamicRNNTanh, _DynamicLSTMCPUGPU, _DynamicLSTMAscend, \ _DynamicGRUAscend, _DynamicGRUCPUGPU from mindspore.nn.layer.rnn_cells import _rnn_tanh_cell, _rnn_relu_cell, _lstm_cell, _gru_cell from ms_adapter.pytorch.nn.modules.module import Module from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.pytorch.functional import empty, zeros from ms_adapter.pytorch.nn import init from ms_adapter.utils import unsupported_attr, is_under_ascend_context class RNNBase(Module): def __init__(self, mode, input_size, hidden_size, num_layers=1, bias=True, batch_first=False, dropout=0., bidirectional=False, proj_size=0, device=None, dtype=None): unsupported_attr(device) super(RNNBase, self).__init__() self.mode = mode self.input_size = input_size self.hidden_size = hidden_size self.num_layers = num_layers self.bias = bias self.batch_first = batch_first self.dropout = float(dropout) self.bidirectional = bidirectional self.proj_size = proj_size self.num_directions = 2 if bidirectional else 1 if not isinstance(dropout, numbers.Number) or not 0 <= dropout <= 1 or \ isinstance(dropout, bool): raise ValueError("dropout should be a number in range [0, 1] " "representing the probability of an element being " "zeroed") if dropout > 0 and num_layers == 1: warnings.warn("dropout option adds dropout after all but last " "recurrent layer, so non-zero dropout expects " "num_layers greater than 1, but got dropout={} and " "num_layers={}".format(dropout, num_layers)) if proj_size < 0: raise ValueError("proj_size should be a positive integer or zero to disable projections") if proj_size >= hidden_size: raise ValueError("proj_size has to be smaller than hidden_size") if mode == 'LSTM': gate_size = 4 * hidden_size elif mode == 'GRU': if is_under_ascend_context() and hidden_size % 16 != 0: raise ValueError(f"GRU on ascend do not support hidden size that is not divisible by 16, " f"but get hidden size {hidden_size}, please reset the argument.") gate_size = 3 * hidden_size elif mode == 'RNN_TANH': gate_size = hidden_size elif mode == 'RNN_RELU': gate_size = hidden_size else: raise ValueError("Unrecognized RNN mode: " + mode) self._flat_weights_names = [] self._all_weights = [] for layer in range(num_layers): for direction in range(self.num_directions): real_hidden_size = proj_size if proj_size > 0 else hidden_size layer_input_size = input_size if layer == 0 else real_hidden_size * self.num_directions w_ih = Parameter(empty((gate_size, layer_input_size), dtype=dtype)) w_hh = Parameter(empty((gate_size, real_hidden_size), dtype=dtype)) b_ih = Parameter(empty(gate_size, dtype=dtype)) b_hh = Parameter(empty(gate_size, dtype=dtype)) layer_params = () if self.proj_size == 0: if bias: layer_params = (w_ih, w_hh, b_ih, b_hh) else: layer_params = (w_ih, w_hh) else: w_hr = Parameter(empty((proj_size, hidden_size), dtype=dtype)) if bias: layer_params = (w_ih, w_hh, b_ih, b_hh, w_hr) else: layer_params = (w_ih, w_hh, w_hr) suffix = '_reverse' if direction == 1 else '' param_names = ['weight_ih_l{}{}', 'weight_hh_l{}{}'] if bias: param_names += ['bias_ih_l{}{}', 'bias_hh_l{}{}'] if self.proj_size > 0: param_names += ['weight_hr_l{}{}'] param_names = [x.format(layer, suffix) for x in param_names] for name, param in zip(param_names, layer_params): setattr(self, name, param) self._flat_weights_names.extend(param_names) self._all_weights.append(param_names) self._flat_weights = \ [(lambda wn: getattr(self, wn) if hasattr(self, wn) else None)(wn) for wn in self._flat_weights_names] self.reset_parameters() def __setattr__(self, attr, value): if hasattr(self, "_flat_weights_names") and attr in self._flat_weights_names: # keep self._flat_weights up to date if you do self.weight = ... idx = self._flat_weights_names.index(attr) self._flat_weights[idx] = value super(RNNBase, self).__setattr__(attr, value) def reset_parameters(self) -> None: stdv = 1.0 / math.sqrt(self.hidden_size) if self.hidden_size > 0 else 0 for weight in self.parameters(): init.uniform_(weight, -stdv, stdv) def extra_repr(self): s = '{input_size}, {hidden_size}' if self.proj_size != 0: s += ', proj_size={proj_size}' if self.num_layers != 1: s += ', num_layers={num_layers}' if self.bias is not True: s += ', bias={bias}' if self.batch_first is not False: s += ', batch_first={batch_first}' if self.dropout != 0: s += ', dropout={dropout}' if self.bidirectional is not False: s += ', bidirectional={bidirectional}' return s.format(**self.__dict__) @property def all_weights(self): return [[getattr(self, weight) for weight in weights] for weights in self._all_weights] def __setstate__(self, d): super(RNNBase, self).__setstate__(d) if 'all_weights' in d: self._all_weights = d['all_weights'] # In PyTorch 1.8 we added a proj_size member variable to LSTM. # LSTMs that were serialized via torch.save(module) before PyTorch 1.8 # don't have it, so to preserve compatibility we set proj_size here. if 'proj_size' not in d: self.proj_size = 0 if isinstance(self._all_weights[0][0], str): return num_layers = self.num_layers num_directions = 2 if self.bidirectional else 1 self._flat_weights_names = [] self._all_weights = [] for layer in range(num_layers): for direction in range(num_directions): suffix = '_reverse' if direction == 1 else '' weights = ['weight_ih_l{}{}', 'weight_hh_l{}{}', 'bias_ih_l{}{}', 'bias_hh_l{}{}', 'weight_hr_l{}{}'] weights = [x.format(layer, suffix) for x in weights] if self.bias: if self.proj_size > 0: self._all_weights += [weights] self._flat_weights_names.extend(weights) else: self._all_weights += [weights[:4]] self._flat_weights_names.extend(weights[:4]) else: if self.proj_size > 0: self._all_weights += [weights[:2]] + [weights[-1:]] self._flat_weights_names.extend(weights[:2] + [weights[-1:]]) else: self._all_weights += [weights[:2]] self._flat_weights_names.extend(weights[:2]) self._flat_weights = \ [(lambda wn: getattr(self, wn) if hasattr(self, wn) else None)(wn) for wn in self._flat_weights_names] def _get_weight_and_bias(self, num_directions, layer, bias): _param_nums_per_directions = 4 if bias else 2 _param_nums_per_layer = num_directions * _param_nums_per_directions offset = _param_nums_per_layer * layer param = () for _ in range(num_directions): if bias: param += tuple(self._flat_weights[offset:offset + _param_nums_per_directions]) else: param += tuple(self._flat_weights[offset:offset + _param_nums_per_directions]) param += (None, None) offset = offset + _param_nums_per_directions return param def forward(self, input, hx=None): if len(input.shape) not in (2, 3): raise RuntimeError(f"For RNN, input should be 2D or 3D, but got {len(input.shape)}D.") is_batched = len(input.shape) == 3 input = cast_to_ms_tensor(input) if not is_batched: input = ms.ops.unsqueeze(input, 1) else: if self.batch_first: input = ms.ops.transpose(input, (1, 0, 2)) x_dtype = input.dtype max_batch_size = input.shape[1] num_directions = 2 if self.bidirectional else 1 if hx is None: hx = zeros(self.num_layers * num_directions, max_batch_size, self.hidden_size, dtype=x_dtype) hx = cast_to_ms_tensor(hx) else: hx = cast_to_ms_tensor(hx) if len(hx.shape) not in (2, 3): raise RuntimeError(f"For RNN, hx should be 2D or 3D, but got {len(hx.shape)}D.") if not is_batched: if len(hx.shape) != 2: raise RuntimeError("For RNN, hx ndim should be equal to input") hx = ms.ops.unsqueeze(hx, 1) pre_layer = input h_n = () # For jit output = None if num_directions == 1: for i in range(self.num_layers): w_ih, w_hh, b_ih, b_hh = self._get_weight_and_bias(num_directions, i, self.bias) output, h_t = self.rnn_cell(pre_layer, hx[i], None, w_ih, w_hh, b_ih, b_hh) h_n += (h_t,) pre_layer = ms.ops.dropout(output, 1 - self.dropout)[0] \ if (self.dropout != 0 and i < self.num_layers - 1) else output else: for i in range(self.num_layers): w_ih, w_hh, b_ih, b_hh, w_ih_b, w_hh_b, b_ih_b, b_hh_b = \ self._get_weight_and_bias(num_directions, i, self.bias) x_b = ms.ops.reverse(pre_layer, [0]) output, h_t = self.rnn_cell(pre_layer, hx[2 * i], None, w_ih, w_hh, b_ih, b_hh) output_b, h_t_b = self.rnn_cell(x_b, hx[2 * i + 1], None, w_ih_b, w_hh_b, b_ih_b, b_hh_b) output_b = ms.ops.reverse(output_b, [0]) output = ms.ops.concat((output, output_b), 2) h_n += (h_t,) h_n += (h_t_b,) pre_layer = ms.ops.dropout(output, 1 - self.dropout)[0] \ if (self.dropout != 0 and i < self.num_layers - 1) else output h_n = ms.ops.concat(h_n, 0) h_n = h_n.view(hx.shape) if not is_batched: output = ms.ops.squeeze(output, 1) h_n = ms.ops.squeeze(h_n, 1) else: if self.batch_first: output = ms.ops.transpose(output, (1, 0, 2)) return cast_to_adapter_tensor(output.astype(x_dtype)), cast_to_adapter_tensor(h_n.astype(x_dtype)) class RNN(RNNBase): def __init__(self, *args, **kwargs): if 'proj_size' in kwargs: raise ValueError("proj_size argument is only supported for LSTM, not RNN or GRU") self.nonlinearity = kwargs.pop('nonlinearity', 'tanh') if self.nonlinearity == 'tanh': mode = 'RNN_TANH' elif self.nonlinearity == 'relu': mode = 'RNN_RELU' else: raise ValueError("Unknown nonlinearity '{}'".format(self.nonlinearity)) super(RNN, self).__init__(mode, *args, **kwargs) if mode == 'RNN_TANH': self.rnn_cell = _DynamicRNNRelu() elif mode == 'RNN_RELU': self.rnn_cell = _DynamicRNNTanh() class GRU(RNNBase): def __init__(self, *args, **kwargs): if 'proj_size' in kwargs: raise ValueError("proj_size argument is only supported for LSTM, not RNN or GRU") super(GRU, self).__init__('GRU', *args, **kwargs) if is_under_ascend_context(): self.rnn_cell = _DynamicGRUAscend() else: self.rnn_cell = _DynamicGRUCPUGPU() class LSTM(RNNBase): def __init__(self, *args, **kwargs): super(LSTM, self).__init__('LSTM', *args, **kwargs) if self.proj_size > 0: raise NotImplementedError("For LSTM, proj_size > 0 is not supported yet.") if is_under_ascend_context(): self.lstm_cell = _DynamicLSTMAscend() else: self.lstm_cell = _DynamicLSTMCPUGPU() def forward(self, input, hx=None): if len(input.shape) not in (2, 3): raise RuntimeError(f"For LSTM, input should be 2D or 3D, but got {len(input.shape)}D.") is_batched = len(input.shape) == 3 input = cast_to_ms_tensor(input) if not is_batched: input = ms.ops.unsqueeze(input, 1) else: if self.batch_first: input = ms.ops.transpose(input, (1, 0, 2)) x_dtype = input.dtype max_batch_size = input.shape[1] num_directions = 2 if self.bidirectional else 1 real_hidden_size = self.proj_size if self.proj_size > 0 else self.hidden_size if hx is None: h_zeros = zeros(self.num_layers * num_directions, max_batch_size, real_hidden_size, dtype=x_dtype) c_zeros = zeros(self.num_layers * num_directions, max_batch_size, self.hidden_size, dtype=x_dtype) hx = (h_zeros, c_zeros) hx = cast_to_ms_tensor(hx) else: hx = cast_to_ms_tensor(hx) if is_batched: if (len(hx[0].shape) != 3 or len(hx[1].shape) != 3): msg = ("For batched 3-D input, hx and cx should " f"also be 3-D but got ({len(hx[0].shape)}-D, {len(hx[1].shape)}-D) tensors") raise RuntimeError(msg) else: if len(hx[0].shape) != 2 or len(hx[1].shape) != 2: msg = ("For unbatched 2-D input, hx and cx should " f"also be 2-D but got ({len(hx[0].shape)}-D, {len(hx[1].shape)}-D) tensors") raise RuntimeError(msg) hx = (ms.ops.unsqueeze(hx[0], 1), ms.ops.unsqueeze(hx[1], 1)) pre_layer = input h_n = () c_n = () # For jit output = None if num_directions == 1: for i in range(self.num_layers): w_ih, w_hh, b_ih, b_hh = self._get_weight_and_bias(num_directions, i, self.bias) h_i = (hx[0][i], hx[1][i]) output, hc_t = self.lstm_cell(pre_layer, h_i, None, w_ih, w_hh, b_ih, b_hh) h_t, c_t = hc_t h_n += (h_t,) c_n += (c_t,) pre_layer = ms.ops.dropout(output, 1 - self.dropout)[0] \ if (self.dropout != 0 and i < self.num_layers - 1) else output else: for i in range(self.num_layers): w_ih, w_hh, b_ih, b_hh, w_ih_b, w_hh_b, b_ih_b, b_hh_b = \ self._get_weight_and_bias(num_directions, i, self.bias) x_b = ms.ops.reverse(pre_layer, [0]) h_i = (hx[0][2 * i], hx[1][2 * i]) h_b_i = (hx[0][2 * i + 1], hx[1][2 * i + 1]) output, hc_t = self.lstm_cell(pre_layer, h_i, None, w_ih, w_hh, b_ih, b_hh) output_b, hc_t_b = self.lstm_cell(x_b, h_b_i, None, w_ih_b, w_hh_b, b_ih_b, b_hh_b) output_b = ms.ops.reverse(output_b, [0]) output = ms.ops.concat((output, output_b), 2) h_t, c_t = hc_t h_t_b, c_t_b = hc_t_b h_n += (h_t,) h_n += (h_t_b,) c_n += (c_t,) c_n += (c_t_b,) pre_layer = ms.ops.dropout(output, 1 - self.dropout)[0] \ if (self.dropout != 0 and i < self.num_layers - 1) else output h_n = ms.ops.concat(h_n, 0) h_n = h_n.view(hx[0].shape) c_n = ms.ops.concat(c_n, 0) c_n = c_n.view(hx[1].shape) if not is_batched: output = ms.ops.squeeze(output, 1) h_n = ms.ops.squeeze(h_n, 1) c_n = ms.ops.squeeze(c_n, 1) else: if self.batch_first: output = ms.ops.transpose(output, (1, 0, 2)) return cast_to_adapter_tensor(output.astype(x_dtype)), \ cast_to_adapter_tensor((h_n.astype(x_dtype), c_n.astype(x_dtype))) class RNNCellBase(Module): def __init__(self, input_size, hidden_size, bias, num_chunks, device=None, dtype=None): unsupported_attr(device) super(RNNCellBase, self).__init__() self.input_size = input_size self.hidden_size = hidden_size self.bias = bias self.weight_ih = Parameter(empty((num_chunks * hidden_size, input_size), dtype=dtype)) self.weight_hh = Parameter(empty((num_chunks * hidden_size, hidden_size), dtype=dtype)) if bias: self.bias_ih = Parameter(empty(num_chunks * hidden_size, dtype=dtype)) self.bias_hh = Parameter(empty(num_chunks * hidden_size, dtype=dtype)) else: self.register_parameter('bias_ih', None) self.register_parameter('bias_hh', None) self._rnn_cell = None self.reset_parameters() def extra_repr(self) -> str: s = '{input_size}, {hidden_size}' if 'bias' in self.__dict__ and self.bias is not True: s += ', bias={bias}' if 'nonlinearity' in self.__dict__ and self.nonlinearity != "tanh": s += ', nonlinearity={nonlinearity}' return s.format(**self.__dict__) def reset_parameters(self) -> None: stdv = 1.0 / math.sqrt(self.hidden_size) if self.hidden_size > 0 else 0 for weight in self.parameters(): init.uniform_(weight, -stdv, stdv) def forward(self, input, hx=None): input = cast_to_ms_tensor(input) if len(input.shape) not in (1, 2): raise RuntimeError(f"RNNCell: Expected input to be 1-D or 2-D but received {len(input.shape)}-D tensor") is_batched = len(input.shape) == 2 if not is_batched: input = ms.ops.unsqueeze(input, 0) if hx is None: hx = zeros(input.shape[0], self.hidden_size, dtype=input.dtype) hx = cast_to_ms_tensor(hx) else: hx = cast_to_ms_tensor(hx) hx = ms.ops.unsqueeze(hx, 0) if not is_batched else hx ret = self._rnn_cell(input, hx, self.weight_ih, self.weight_hh, self.bias_ih, self.bias_hh) if not is_batched: ret = ms.ops.squeeze(ret, 0) return cast_to_adapter_tensor(ret) class RNNCell(RNNCellBase): def __init__(self, input_size, hidden_size, bias=True, nonlinearity="tanh", device=None, dtype=None): super(RNNCell, self).__init__(input_size, hidden_size, bias, num_chunks=1, device=device, dtype=dtype) self.nonlinearity = nonlinearity if self.nonlinearity == "tanh": self._rnn_cell = _rnn_tanh_cell elif self.nonlinearity == "relu": self._rnn_cell = _rnn_relu_cell else: raise RuntimeError( "Unknown nonlinearity: {}".format(self.nonlinearity)) class LSTMCell(RNNCellBase): def __init__(self, input_size, hidden_size, bias=True, device=None, dtype=None): super(LSTMCell, self).__init__(input_size, hidden_size, bias, num_chunks=4, device=device, dtype=dtype) def forward(self, input, hx=None): input = cast_to_ms_tensor(input) if len(input.shape) not in (1, 2): raise RuntimeError(f"LSTMCell: Expected input to be 1-D or 2-D but received {len(input.shape)}-D tensor") is_batched = len(input.shape) == 2 if not is_batched: input = ms.ops.unsqueeze(input, 0) if hx is None: _zeros = zeros(input.shape[0], self.hidden_size, dtype=input.dtype) hx = (_zeros, _zeros) hx = cast_to_ms_tensor(hx) else: hx = cast_to_ms_tensor(hx) hx = (ms.ops.unsqueeze(hx[0], 0), ms.ops.unsqueeze(hx[1], 0)) if not is_batched else hx hx = cast_to_ms_tensor(hx) ret = _lstm_cell(input, hx, self.weight_ih, self.weight_hh, self.bias_ih, self.bias_hh) if not is_batched: ret = (ms.ops.squeeze(ret[0], 0), ms.ops.squeeze(ret[1], 0)) return cast_to_adapter_tensor(ret) class GRUCell(RNNCellBase): def __init__(self, input_size, hidden_size, bias=True, device=None, dtype=None): super(GRUCell, self).__init__(input_size, hidden_size, bias, num_chunks=3, device=device, dtype=dtype) self._rnn_cell = _gru_cell msadapter/ms_adapter/pytorch/nn/modules/sparse.py000066400000000000000000000073461436713554100226220ustar00rootroot00000000000000import ms_adapter.pytorch.nn.functional as Adapter_F from ms_adapter.pytorch.functional import empty from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.utils import unsupported_attr from ms_adapter.pytorch.nn.modules.module import Module from ms_adapter.pytorch.nn.init import normal_ __all__ = ['Embedding'] class Embedding(Module): def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2., scale_grad_by_freq=False, sparse=False, _weight=None, device=None, dtype=None): unsupported_attr(scale_grad_by_freq) unsupported_attr(sparse) unsupported_attr(device) #TODO: padding_idx is not supported, because can not not updating the gradient of # weight[padding_idx] if padding_idx: raise NotImplementedError("nn.Embedding: `padding_idx` is not supported until now.") super(Embedding, self).__init__() self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim if padding_idx is not None: if padding_idx > 0: if padding_idx >= self.num_embeddings: raise ValueError('Padding_idx must be within num_embeddings') elif padding_idx < 0: if padding_idx < -self.num_embeddings: raise ValueError('Padding_idx must be within num_embeddings') padding_idx = self.num_embeddings + padding_idx self.padding_idx = padding_idx self.max_norm = max_norm self.norm_type = norm_type self.scale_grad_by_freq = scale_grad_by_freq if _weight is None: self.weight = Parameter(empty((num_embeddings, embedding_dim), dtype=dtype)) self.reset_parameters() else: if list(_weight.shape) != [num_embeddings, embedding_dim]: raise ValueError('Shape of weight does not match num_embeddings and embedding_dim') self.weight = Parameter(_weight) self.sparse = sparse def reset_parameters(self): normal_(self.weight) self._fill_padding_idx_with_zero() def _fill_padding_idx_with_zero(self): # TODO: to support `padding_idx` in the future return def forward(self, input): return Adapter_F.embedding( input, self.weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse) def extra_repr(self): s = '{num_embeddings}, {embedding_dim}' if self.padding_idx is not None: s += ', padding_idx={padding_idx}' if self.max_norm is not None: s += ', max_norm={max_norm}' if self.norm_type != 2: s += ', norm_type={norm_type}' if self.scale_grad_by_freq is not False: s += ', scale_grad_by_freq={scale_grad_by_freq}' if self.sparse is not False: s += ', sparse=True' return s.format(**self.__dict__) @classmethod def from_pretrained(cls, embeddings, freeze=True, padding_idx=None, max_norm=None, norm_type=2., scale_grad_by_freq=False, sparse=False): embeddings_rank = len(embeddings.shape) if embeddings_rank != 2: raise ValueError(f"dim of embedding must be 2, but got {embeddings_rank}") rows, cols = embeddings.shape embedding = cls( num_embeddings=rows, embedding_dim=cols, _weight=embeddings, padding_idx=padding_idx, max_norm=max_norm, norm_type=norm_type, scale_grad_by_freq=scale_grad_by_freq, sparse=sparse) embedding.weight.requires_grad = not freeze return embedding msadapter/ms_adapter/pytorch/nn/modules/transformer.py000066400000000000000000000000001436713554100236430ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/nn/modules/unpooling.py000066400000000000000000000025071436713554100233310ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import ms_adapter.pytorch.nn.functional as Adapter_F from .module import Module __all__ = ['MaxUnpool1d', 'MaxUnpool2d', 'MaxUnpool3d'] class _MaxUnpoolNd(Module): def __init__(self, kernel_size, stride=None, padding=0): super(_MaxUnpoolNd, self).__init__() self.kernel_size = kernel_size self.stride = stride if (stride is not None) else kernel_size self.padding = padding def extra_repr(self) -> str: return 'kernel_size={}, stride={}, padding={}'.format( self.kernel_size, self.stride, self.padding ) class MaxUnpool1d(_MaxUnpoolNd): def forward(self, input, indices, output_size = None): return Adapter_F.max_unpool1d(input, indices, self.kernel_size, self.stride, self.padding, output_size) class MaxUnpool2d(_MaxUnpoolNd): def forward(self, input, indices, output_size = None): return Adapter_F.max_unpool2d(input, indices, self.kernel_size, self.stride, self.padding, output_size) class MaxUnpool3d(_MaxUnpoolNd): def forward(self, input, indices, output_size = None): return Adapter_F.max_unpool3d(input, indices, self.kernel_size, self.stride, self.padding, output_size) msadapter/ms_adapter/pytorch/nn/modules/upsampling.py000066400000000000000000000062021436713554100234720ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch.nn import functional as adapter_F from .module import Module __all__ = ['Upsample', 'UpsamplingNearest2d', 'UpsamplingBilinear2d'] class Upsample(Module): r"""Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data. The input data is assumed to be of the form `minibatch x channels x [optional depth] x [optional height] x width`. Hence, for spatial inputs, we expect a 4D Tensor and for volumetric inputs, we expect a 5D Tensor. The algorithms available for upsampling are nearest neighbor and linear, bilinear, bicubic and trilinear for 3D, 4D and 5D input Tensor, respectively. One can either give a :attr:`scale_factor` or the target output :attr:`size` to calculate the output size. (You cannot give both, as it is ambiguous) Args: size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int], optional): output spatial sizes scale_factor (float or Tuple[float] or Tuple[float, float] or Tuple[float, float, float], optional): multiplier for spatial size. Has to match input size if it is a tuple. mode (str, optional): the upsampling algorithm: one of ``'nearest'``, ``'linear'``, ``'bilinear'``, ``'bicubic'`` and ``'trilinear'``. Default: ``'nearest'`` align_corners (bool, optional): if ``True``, the corner pixels of the input and output tensors are aligned, and thus preserving the values at those pixels. This only has effect when :attr:`mode` is ``'linear'``, ``'bilinear'``, or ``'trilinear'``. Default: ``False`` """ def __init__(self, size=None, scale_factor=None, mode='bilinear', align_corners=None, recompute_scale_factor=None): super(Upsample, self).__init__() self.name = type(self).__name__ self.size = size self.scale_factor = scale_factor self.mode = mode self.align_corners = align_corners if recompute_scale_factor is not None: raise ValueError("recompute_scale_factor is not supported") def forward(self, input): return adapter_F.interpolate( input, size=self.size, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners ) def extra_repr(self) -> str: if self.scale_factor is not None: info = 'scale_factor=' + str(self.scale_factor) else: info = 'size=' + str(self.size) info += ', mode=' + self.mode return info class UpsamplingNearest2d(Upsample): def __init__(self, size=None, scale_factor=None): super(UpsamplingNearest2d, self).__init__(size, scale_factor, mode='nearest') def forward(self, input): return adapter_F.upsample_nearest(input, self.size, self.scale_factor) class UpsamplingBilinear2d(Upsample): def __init__(self, size=None, scale_factor=None): super(UpsamplingBilinear2d, self).__init__(size, scale_factor, mode='bilinear') def forward(self, input): return adapter_F.upsample_bilinear(input, self.size, self.scale_factor) msadapter/ms_adapter/pytorch/nn/modules/utils.py000066400000000000000000000014271436713554100224570ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import collections from itertools import repeat def _ntuple(n, name="parse"): def parse(x): if isinstance(x, list) and len(x) == 1: x = x[0] if isinstance(x, collections.abc.Iterable): return tuple(x) return tuple(repeat(x, n)) parse.__name__ = name return parse _single = _ntuple(1, "_single") _pair = _ntuple(2, "_pair") _triple = _ntuple(3, "_triple") _quadruple = _ntuple(4, "_quadruple") def _reverse_repeat_tuple(t, n): r"""Reverse the order of `t` and repeat each element for `n` times. This can be used to translate padding arg used by Conv and Pooling modules to the ones used by `F.pad`. """ return tuple(x for x in reversed(t) for _ in range(n)) msadapter/ms_adapter/pytorch/nn/parameter.py000066400000000000000000000352311436713554100216270ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- """Parameter interface""" import sys import numbers from copy import copy import mindspore as ms import mindspore.common.dtype as mstype from mindspore.common.initializer import initializer from mindspore._checkparam import Validator from mindspore._c_expression import Tensor as Tensor_ from mindspore.parallel._tensor import _get_slice_index from mindspore.parallel._auto_parallel_context import auto_parallel_context from mindspore.parallel._ps_context import _is_role_worker, _is_role_sched, _clone_hash_table from mindspore.parallel._ps_context import _insert_accumu_init_info from mindspore.ops import functional as F from ms_adapter.pytorch.tensor import Tensor, cast_to_adapter_tensor __all__ = ['Parameter', 'ParameterTuple'] PARAMETER_NAME_DEFAULT = "Parameter" PARAMETER_NAME_PREFIX_MAX_LEN = 1024 # Global variable for parameter unique key. _GLOBAL_PARAMETER_KEY = -1 def _is_in_parallel_mode(): """Get parallel mode.""" return auto_parallel_context().get_parallel_mode() in ["semi_auto_parallel", "auto_parallel"] def init_to_value(init): """ Get value of initializer. Returns: Value of the initializer. Raises: ValueError: The value of the argument 'init' is not correct. """ if isinstance(init, str): if init == 'zeros': return 0.0 if init == 'ones': return 1.0 raise ValueError("The argument 'init' should be one of values in ['zeros', 'ones'].") if isinstance(init, numbers.Number): return float(init) raise ValueError("The argument 'init' should be number or string, but got {}.".format(type(init))) def _get_unique_parameter_key(): """ Get parameter unique key. Used to identify the same Parameter for Worker and Server in the embedding cache scenario. Returns: Integer. The unique parameter key. """ global _GLOBAL_PARAMETER_KEY _GLOBAL_PARAMETER_KEY += 1 return _GLOBAL_PARAMETER_KEY class Parameter(ms.Parameter): # Parameter is an subclass of ms.Parameter and adapter.Tensor, only 'Parameter' in methods need be overload. _base_type = {} def __new__(cls, data, *args, **kwargs): init_data_flag = bool(isinstance(data, ms.Tensor) and data.has_init) rc = sys.getrefcount(data) _, *class_init_args = Parameter._get_parameter_new_args(data, rc) new_type = Parameter._get_base_class(Tensor) obj = Tensor.__new__(new_type) Tensor.__init__(obj, *class_init_args, inner=True) # it's better to make the Initializer a kind of tensor. obj.init_mode = None obj.is_default_input_init = init_data_flag if obj.has_init: obj.init_mode = data return obj def __reduce_ex__(self, _): data = self if self.init_mode is not None: data = self.init_mode else: # cast to break deep infinite loop while deepcopy data = Tensor(self) return ( Parameter, (data, self.name, self.requires_grad, self.layerwise_parallel)) def __init__(self, data, name=None, requires_grad=True, layerwise_parallel=False, parallel_optimizer=True): super().__init__(default_input=data, name=name, requires_grad=requires_grad, layerwise_parallel=layerwise_parallel, parallel_optimizer=parallel_optimizer) def __deepcopy__(self, memodict): new_obj = Parameter(self) new_obj.name = self.name new_obj._inited_param = self._inited_param return new_obj def __str__(self): return f'Parameter containing: {Tensor_.__repr__(self.data)}, requires_grad={self.requires_grad})' def __parameter__(self): """For parse check.""" @staticmethod def _get_base_class(input_class): input_class_name = Parameter.__name__ if input_class_name in Parameter._base_type: new_type = Parameter._base_type.get(input_class_name) else: new_type = type(input_class_name, (Parameter, input_class), {}) Parameter._base_type[input_class_name] = new_type return new_type def copy(self): """ Copy the parameter. Returns: Parameter, a new parameter. """ return self.clone(init='same') def clone(self, init='same'): """ Clone the parameter. Args: init (Union[Tensor, str, numbers.Number]): Initialize the shape and dtype of the parameter. If `init` is a `Tensor` or `numbers.Number`, clone a new parameter with the same shape and dtype, and the data of the new parameter will be set according to `init`. If `init` is a `str`, the `init` should be the alias of the class inheriting from `Initializer`. For example, if `init` is 'same', clone a new parameter with the same data, shape, and dtype. Default: 'same'. Returns: Parameter, a new parameter. """ x = copy(self) param_info_clone = self.param_info.clone() info = self.param_info if hasattr(info, "cloned_obj"): info.cloned_obj.append(x) else: info.cloned_obj = [x] self.param_info = info param_info_clone.obj = x x.param_info = param_info_clone x.is_init = False x.init = self.init x.is_param_ps = self.is_param_ps x.init_in_server = self.init_in_server x.cache_enable = self.cache_enable if x.cache_enable: x.key = _get_unique_parameter_key() x.requires_aggr = self.requires_aggr if self.cache_shape: x.cache_shape = self.cache_shape if init != 'same': shape = self.shape dtype = self.dtype init_data = initializer(init, shape=shape, dtype=dtype) x.set_data(cast_to_adapter_tensor(init_data)) return x @property def data(self): """Return the parameter object.""" return self def _update_tensor_data(self, data): """Update the parameter by a Tensor.""" if isinstance(self, Tensor): self.init_flag = False self.init = None return self.assign_value(data) new_param = Parameter(data, self.name, self.requires_grad) new_param.param_info = self.param_info return new_param @staticmethod def _from_tensor(tensor, *args, **kwargs): """Create a `Parameter` that data is shared from a `Tensor`.""" if not isinstance(tensor, Tensor_): raise TypeError(f"The type of input must be Tensor, but got {type(tensor)}.") param = Tensor_.__new__(Parameter) Tensor_.__init__(param, tensor) param.init = None param.init_mode = None param.is_default_input_init = False Parameter.__init__(param, tensor, *args, **kwargs) return param def set_data(self, data, slice_shape=False): """ Set Parameter's data. Args: data (Union[Tensor, int, float]): New data. slice_shape (bool): If slice the parameter is set to true, the shape is not checked for consistency. Default: False. Returns: Parameter, the parameter after set data. """ if not isinstance(data, (ms.Tensor, int, float)): raise TypeError(f"Parameter data must be [`Tensor`, `int`, `float`] or a kind of `Tensor` " f"(like `Tensor`). But with type {type(data)}.") if isinstance(data, (int, float)): if self.dtype in mstype.int_type and isinstance(data, float): self._raise_type_error(mstype.float_) data = Tensor(data, self.dtype) # both not init. incoming_tensor_is_init = isinstance(data, ms.Tensor) and not data.has_init current_tensor_is_init = isinstance(self, Tensor) and not self.has_init Parameter._set_data_check_input_valid(self.shape, data.shape, current_tensor_is_init, incoming_tensor_is_init, slice_shape) if self.dtype != data.dtype: if mstype.implicit_conversion_seq[self.dtype] < mstype.implicit_conversion_seq[data.dtype]: self._raise_type_error(data.dtype) else: if isinstance(data, ms.Tensor) and data.init is not None: data.init_data() data = F.cast(data, self.dtype) if isinstance(data, ms.Tensor) and data.has_init: # The parameter has been initialized, directly update by the data if current_tensor_is_init: self._update_tensor_data(data.init_data()) else: # also update the related inited parameter data if self.inited_param is not None: self.inited_param.set_data(data) self.init_mode = data elif incoming_tensor_is_init or current_tensor_is_init: self._update_tensor_data(data) self.sliced = slice_shape return self @staticmethod def _get_init_data_args(layout=None): """Get the data layout args.""" init_data_args = () if layout: if not isinstance(layout, tuple): raise TypeError("The argument 'layout' should be tuple, but got {}.".format(type(layout))) if len(layout) < 6: raise ValueError("The length of 'layout' must be larger than 5, but got {}.".format(len(layout))) slice_index = int(_get_slice_index(layout[0], layout[1])) init_data_args += (slice_index, layout[2], layout[5]) return init_data_args def init_data(self, layout=None, set_sliced=False): """ Initialize the parameter's data. Args: layout (Union[None, tuple]): The parameter's layout info. layout [dev_mat, tensor_map, slice_shape, filed_size, uniform_split, opt_shard_group]. Default: None. It's not None only in 'SEMI_AUTO_PARALLEL' or 'AUTO_PARALLEL' mode. - dev_mat (list(int)): The parameter's device matrix. - tensor_map (list(int)): The parameter's tensor map. - slice_shape (list(int)): The parameter's slice shape. - filed_size (int): The parameter's filed size. - uniform_split (bool): Whether the parameter is split evenly. - opt_shard_group (str): The group of the parameter while running optimizer parallel. set_sliced (bool): True if the parameter is set sliced after initializing the data. Default: False. Raises: RuntimeError: If it is from Initializer, and parallel mode has changed after the Initializer created. ValueError: If the length of the layout is less than 6. TypeError: If `layout` is not tuple. Returns: Parameter, the `Parameter` after initializing data. If current `Parameter` was already initialized before, returns the same initialized `Parameter`. """ if self.is_default_input_init and self.is_in_parallel != _is_in_parallel_mode(): raise RuntimeError("Must set or change parallel mode before any Tensor created.") if self.init_mode is None: return self if self.inited_param is not None: return self.inited_param init_data_args = self._get_init_data_args(layout) if self.init_in_server and self.is_param_ps and isinstance(self.init_mode, Tensor) and \ self.init_mode.init is not None and (_is_role_worker() or _is_role_sched()): if self.cache_enable: data = self.init_mode.init_data(*init_data_args) else: data = self.init_mode.init_data(0, [1]) else: data = self.init_mode.init_data(*init_data_args) obj = self._update_tensor_data(data) if id(obj) != id(self): self._inited_param = obj obj.init_mode = None obj.sliced = set_sliced return obj def requires_grad_(self, requires_grad=True): self.requires_grad = requires_grad class ParameterTuple(tuple): """ Inherited from tuple, ParameterTuple is used to save multiple parameter. Note: It is used to store the parameters of the network into the parameter tuple collection. """ def __new__(cls, iterable): """Create instance object of ParameterTuple.""" data = tuple(iterable) ids = set() names = set() for x in data: if not isinstance(x, Parameter): raise TypeError(f"For ParameterTuple initialization, " f"ParameterTuple input should be 'Parameter' collection, " f"but got a {type(iterable)}. ") if id(x) not in ids: if x.name in names: raise ValueError("The value {} , its name '{}' already exists. " "Please set a unique name for the parameter.".format(x, x.name)) names.add(x.name) ids.add(id(x)) return tuple.__new__(ParameterTuple, tuple(data)) def clone(self, prefix, init='same'): """ Clone the parameters in ParameterTuple element-wisely to generate a new ParameterTuple. Args: prefix (str): Namespace of parameter, the prefix string will be added to the names of parameters in parametertuple. init (Union[Tensor, str, numbers.Number]): Clone the shape and dtype of Parameters in ParameterTuple and set data according to `init`. Default: 'same'. If `init` is a `Tensor` , set the new Parameter data to the input Tensor. If `init` is `numbers.Number` , set the new Parameter data to the input number. If `init` is a `str`, data will be seted according to the initialization method of the same name in the `Initializer`. If `init` is 'same', the new Parameter has the same value with the original Parameter. Returns: Tuple, the new Parameter tuple. """ Validator.check_str_by_regular(prefix) new = [] for x in self: x1 = x.clone(init) x1.name = prefix + "." + x1.name new.append(x1) if not x1.cache_enable: continue if _is_role_worker(): _clone_hash_table(x.name, x.key, x1.name, x1.key) _insert_accumu_init_info(x1.name, init_to_value(init)) return ParameterTuple(new) def __parameter_tuple__(self): """For parse check.""" msadapter/ms_adapter/pytorch/optim/000077500000000000000000000000001436713554100200065ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/optim/__init__.py000066400000000000000000000021731436713554100221220ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from mindspore.nn.optim import SGD from mindspore.nn.optim import Adam from mindspore.nn.optim import Momentum from mindspore.nn.optim import LARS from mindspore.nn.optim import AdamWeightDecay from mindspore.nn.optim import LazyAdam from mindspore.nn.optim import AdamOffload from mindspore.nn.optim import Lamb from mindspore.nn.optim import ASGD from mindspore.nn.optim import RMSProp from mindspore.nn.optim import Rprop from mindspore.nn.optim import FTRL from mindspore.nn.optim import ProximalAdagrad from mindspore.nn.optim import Adagrad from mindspore.nn.optim import thor from mindspore.nn.optim import AdaFactor from mindspore.nn.optim import AdaSumByDeltaWeightWrapCell from mindspore.nn.optim import AdaSumByGradWrapCell from mindspore.nn.optim import AdaMax from mindspore.nn.optim import Adadelta __all__ = ['Momentum', 'LARS', 'Adam', 'AdamWeightDecay', 'LazyAdam', 'AdamOffload', 'Lamb', 'SGD', 'ASGD', 'Rprop', 'FTRL', 'RMSProp', 'ProximalAdagrad', 'Adagrad', 'thor', 'AdaFactor', 'AdaSumByDeltaWeightWrapCell', 'AdaSumByGradWrapCell', 'AdaMax', 'Adadelta'] msadapter/ms_adapter/pytorch/package_info.py000066400000000000000000000004351436713554100216400ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- MAJOR = 0 MINOR = 1 PATCH = 0 PRE_RELEASE = '' # Use the following formatting: (major, minor, patch, prerelease) VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE) __version__ = version = '.'.join(map(str, VERSION[:3])) + ''.join(VERSION[3:]) msadapter/ms_adapter/pytorch/random.py000066400000000000000000000003661436713554100205150ustar00rootroot00000000000000import mindspore as ms import numpy as np def manual_seed(seed): return ms.set_seed(seed) def seed(): value = np.floor(np.random.random(1) * 2**32 - 1) return ms.set_seed(int(value)) def initial_seed(): return ms.get_seed() msadapter/ms_adapter/pytorch/storage.py000066400000000000000000000026751436713554100207060ustar00rootroot00000000000000import numpy as np import ms_adapter.pytorch.common.dtype as _dtype from ms_adapter.utils import unsupported_attr _default_dtype = np.float32 # TODO: to supported more methods class _TypedStorage(): def __new__(cls, *args, wrap_storage=None, dtype=None, device=None): if cls == _TypedStorage: return super().__new__(cls) if dtype is not None: raise RuntimeError("For TypedStorage, Keyword argument 'dtype' cannot be specified") _dtype = _cls_dtype_map.get(cls) return _TypedStorage(*args, wrap_storage=wrap_storage, dtype=_dtype, device=device) def __init__(self, *args, wrap_storage=None, dtype=None, device=None): unsupported_attr(device) if wrap_storage is not None: raise NotImplementedError("For TypeStorage, `wrap_storage` is not supported yet.") if len(args) != 1 or isinstance(args[0], int): raise NotImplementedError("For TypeStorage, `args` now only support sequence data.") if not dtype: dtype = _default_dtype self.data = np.frombuffer(args[0], dtype) self.dtype = dtype @classmethod def from_buffer(cls, bytestream): _dtype = _cls_dtype_map.get(cls) return _TypedStorage(bytestream, dtype=_dtype) class ByteStorage(_TypedStorage): @property def dtype(self): return _dtype.uint8 # TODO: To supported other type storage. _cls_dtype_map = {ByteStorage: np.uint8} msadapter/ms_adapter/pytorch/tensor.py000066400000000000000000002470441436713554100205550ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import warnings import numbers import numpy as np import mindspore as ms from mindspore.common import dtype as mstype from mindspore.common._register_for_tensor import tensor_operator_registry import mindspore.ops as P from mindspore.ops import constexpr from mindspore.ops._primitive_cache import _get_cache_prim from mindspore.ops.operations import _inner_ops as inner from mindspore.common.initializer import _init_random_normal, _init_random_uniform, Zero from mindspore._c_expression import Tensor as Tensor_ from ms_adapter.utils import unsupported_attr, is_under_gpu_context, get_backend, \ is_under_ascend_context, _infer_size, _ascend_tensor_general_cast import ms_adapter.pytorch.common.dtype as msdapter_dtype from ms_adapter.pytorch.common.device import Device from ms_adapter.pytorch.storage import _TypedStorage _dtypeDict = { 'float16': mstype.float16, 'float32': mstype.float32, 'float64': mstype.float64, 'int8': mstype.int8, 'int16': mstype.int16, 'int32': mstype.int32, 'int64': mstype.int64, 'uint8': mstype.uint8, 'uint16': mstype.uint16, 'uint32': mstype.uint32, 'uint64': mstype.uint64, 'bool': mstype.bool_, 'complex64': mstype.complex64, 'complex128': mstype.complex128, 'long': mstype.int64, 'half': mstype.float16, 'int': mstype.int32, 'double': mstype.float64, 'float': mstype.float32, 'char': mstype.int8, 'byte': mstype.uint8, 'short': mstype.int16 } kMaxInt8 = 2 ** 7 - 1 kMaxInt16 = 2 ** 15 - 1 kMaxInt32 = 2 ** 31 - 1 kMaxInt64 = 2 ** 63 - 1 kMaxUint8 = 2 ** 8 - 1 kMaxUint16 = 2 ** 16 - 1 kMaxUint32 = 2 ** 32 - 1 kMaxUint64 = 2 ** 64 - 1 kMantissaFloat16 = 2 ** 11 kMantissaFloat32 = 2 ** 24 kMantissaFloat64 = 2 ** 53 _dtype2typeDict = { 'float32': 'FloatTensor', 'float': 'FloatTensor', 'float64': 'DoubleTensor', 'double': 'DoubleTensor', 'complex64': 'ComplexFloatTensor', 'cfloat': 'ComplexFloatTensor', 'complex128': 'ComplexDoubleTensor', 'cdouble': 'ComplexDoubleTensor', 'float16': 'HalfTensor', 'half': 'HalfTensor', 'bfloat16': 'BFloat16Tensor', 'uint8': 'ByteTensor', 'int8': 'CharTensor', 'int16': 'ShortTensor', 'short': 'ShortTensor', 'int32': 'IntTensor', 'int': 'IntTensor', 'int64': 'LongTensor', 'long': 'LongTensor', 'bool': 'BoolTensor' } _type2dtypeDict = { 'FloatTensor': msdapter_dtype.float32, 'DoubleTensor': msdapter_dtype.float64, 'ComplexFloatTensor': msdapter_dtype.complex64, 'ComplexDoubleTensor': msdapter_dtype.complex128, 'HalfTensor': msdapter_dtype.float16, 'BFloat16Tensor': msdapter_dtype.bfloat16, 'ByteTensor': msdapter_dtype.uint8, 'CharTensor' : msdapter_dtype.int8, 'ShortTensor': msdapter_dtype.int16, 'IntTensor': msdapter_dtype.int32, 'LongTensor': msdapter_dtype.int64, 'BoolTensor': msdapter_dtype.bool } class Tensor(ms.Tensor): def __init__(self, *data, dtype=None, inner=False): def _process_data(data): _shape = None _input_data = None if len(data) == 1: if isinstance(data[0], int): _shape = data elif isinstance(data[0], (np.ndarray, ms.Tensor, list, Tensor_)): _input_data = data[0] elif isinstance(data[0], tuple): if len(data[0]) == 1: _shape = data[0] else: _input_data = data[0] elif isinstance(data[0], _TypedStorage): _input_data=data[0].data else: raise TypeError(f"For Tensor, data must be a sequence, got {type(data[0])}") elif len(data) > 1: if not isinstance(data[0], int): raise TypeError("For Tensor, elements of shape must be int.") _shape = data else: _input_data = () return _input_data, _shape if dtype is not None: dtype = _dtypeDict[str(dtype).split('.')[-1].lower()] if inner is True: super(Tensor, self).__init__(*data, dtype=dtype) else: _input_data, _shape = _process_data(data) if _shape: if dtype is None: dtype = mstype.float32 super(Tensor, self).__init__(shape=_shape, dtype=dtype, init=Zero()) self.init_data() else: if dtype is None: if not isinstance(_input_data, (ms.Tensor, Tensor_, _TypedStorage)): dtype=mstype.float32 super(Tensor, self).__init__(input_data=_input_data, dtype=dtype) def __neg__(self): tensor_ms = cast_to_ms_tensor(self) out = tensor_ms.__neg__() return cast_to_adapter_tensor(out) def __invert__(self): tensor_ms = cast_to_ms_tensor(self) if tensor_ms.dtype != ms.bool_: out = - 1 - tensor_ms else: out = tensor_ms.__invert__() return cast_to_adapter_tensor(out) def __round__(self): tensor_ms = cast_to_ms_tensor(self) out = tensor_ms.__round__() return cast_to_adapter_tensor(out) def __pos__(self): tensor_ms = cast_to_ms_tensor(self) out = tensor_ms.__pos__() return cast_to_adapter_tensor(out) def __abs__(self): tensor_ms = cast_to_ms_tensor(self) out = tensor_ms.__abs__() return cast_to_adapter_tensor(out) def __add__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__add__(other_ms) return cast_to_adapter_tensor(out) def __and__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) input_dtype = tensor_ms.dtype if input_dtype == mstype.bool_: # avoid BitwiseAnd op has not corresponding bprop. tensor_ms = tensor_ms.astype(mstype.int8) out = tensor_ms.mul(other_ms) out = out.astype(mstype.bool_) else: out = tensor_ms.__and__(other_ms) return cast_to_adapter_tensor(out) def __xor__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__xor__(other_ms) return cast_to_adapter_tensor(out) def __or__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) input_dtype = tensor_ms.dtype if input_dtype == mstype.bool_: # avoid BitwiseOr op has not corresponding bprop. tensor_ms = tensor_ms.astype(mstype.int8) out = tensor_ms.add(other_ms) out = out.astype(mstype.bool_) else: out = tensor_ms.__or__(other_ms) return cast_to_adapter_tensor(out) def __radd__(self, other): return self.__add__(other) def __iadd__(self, other): return self.__add__(other) def __sub__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__sub__(other_ms) return cast_to_adapter_tensor(out) def __rsub__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__rsub__(other_ms) return cast_to_adapter_tensor(out) def __isub__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__isub__(other_ms) return cast_to_adapter_tensor(out) def __mul__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__mul__(other_ms) return cast_to_adapter_tensor(out) def __rmul__(self, other): return self.__mul__(other) def __imul__(self, other): return self.__mul__(other) def __truediv__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) tensor_type = tensor_ms.dtype if 'Int' in str(tensor_type): tensor_ms = ms.ops.cast(tensor_ms, mstype.float32) out = tensor_ms.__truediv__(other_ms) return cast_to_adapter_tensor(out) def __rtruediv__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) tensor_type = tensor_ms.dtype if 'Int' in str(tensor_type): tensor_ms = ms.ops.cast(tensor_ms, mstype.float32) out = tensor_ms.__rtruediv__(other_ms) return cast_to_adapter_tensor(out) def __mod__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__mod__(other_ms) return cast_to_adapter_tensor(out) def __rmod__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__rmod__(other_ms) return cast_to_adapter_tensor(out) def __imod__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__imod__(other_ms) return cast_to_adapter_tensor(out) def __pow__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__pow__(other_ms) return cast_to_adapter_tensor(out) def __rpow__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__rpow__(other_ms) return cast_to_adapter_tensor(out) def __floordiv__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__floordiv__(other_ms) return cast_to_adapter_tensor(out) def __rfloordiv__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__rfloordiv__(other_ms) return cast_to_adapter_tensor(out) def __ifloordiv__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__ifloordiv__(other_ms) return cast_to_adapter_tensor(out) def __lt__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__lt__(other_ms) return cast_to_adapter_tensor(out) def __le__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__le__(other_ms) return cast_to_adapter_tensor(out) def __gt__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__gt__(other_ms) return cast_to_adapter_tensor(out) def __ge__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__ge__(other_ms) return cast_to_adapter_tensor(out) def __eq__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__eq__(other_ms) return cast_to_adapter_tensor(out) def __hash__(self): return hash(id(self)) def __ne__(self, other): tensor_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) out = tensor_ms.__ne__(other_ms) return cast_to_adapter_tensor(out) # __setitem__ no need to overload def _getitem_handler(self, index): tensor_ms = cast_to_ms_tensor(self) if isinstance(index, bool): if index: return tensor_ms.expand_dims(0) else: index = ms.Tensor(False) out = ms.ops.masked_select(tensor_ms, index) return out if isinstance(index, tuple) and isinstance(index[0], bool): if False in index: index = ms.Tensor(False) out = ms.ops.masked_select(tensor_ms, index) return out else: return tensor_ms.expand_dims(0) if isinstance(index, ms.Tensor) and index.dtype == ms.bool_: # TODO: delete here after support to reshape empty tensor # TODO: not support GRAPH_MODE if ms.ops.count_nonzero(index.astype(ms.int32)) == 0: out_shape = list(tensor_ms.shape) out_shape[0] = 0 out_shape = tuple(out_shape) out = ms.ops.zeros(out_shape, dtype=tensor_ms.dtype) return out ms_shape_len = len(tensor_ms.shape) index_shape_len = len(index.shape) out_shape = [-1] while index_shape_len < ms_shape_len: out_shape.append(tensor_ms.shape[index_shape_len]) index = index.expand_dims(-1) index_shape_len += 1 out = ms.ops.masked_select(tensor_ms, index) if len(out_shape) > 1: out = out.reshape(out_shape) else: if ms.ops.size(tensor_ms) == 0: #TODO: not support GRAPH_MODE input_shape = tensor_ms.shape scale = [1] * len(input_shape) tmp_shape = list(input_shape) for i in range(len(input_shape)): if input_shape[i] == 0: scale[i] = 0 tmp_shape[i] = 1 tmp_tensor = ms.ops.zeros(tuple(tmp_shape), dtype=tensor_ms.dtype) tmp_out = tmp_tensor.__getitem__(index) out_shape = list(tmp_out.shape) for i in range(len(out_shape)): out_shape[i] = out_shape[i] * scale[i] out = ms.ops.zeros(tuple(out_shape), dtype=tensor_ms.dtype) else: out = tensor_ms.__getitem__(index) return out def __getitem__(self, index): out = cast_to_adapter_tensor(self._getitem_handler(index)) if out is not self: out.parent_tensor_ = self out.index_of_parent_ = index return out def __getstate__(self): pickled = {"input_data": self.asnumpy(), "dtype": self.dtype, "const_arg": self.const_arg} return pickled def __setstate__(self, state): self.__init__(state["input_data"], dtype=state["dtype"]) def fill_adapter(self, val): if not isinstance(val, (int, float, bool)): raise TypeError("For 'Tensor.fill', the type of the argument 'value' must be int, float or bool, " "but got {}.".format(type(val))) output = tensor_operator_registry.get("fill")(self.dtype, self.shape, val) return cast_to_adapter_tensor(output) def fill_(self, val): output = self.fill_adapter(val) return _tensor_inplace_assign(self, output, "fill_", "fill_adapter") def normal_adapter(self, mean=0, std=1, *, generator=None): if generator is not None: raise ValueError("`generator` can not be supportted.") output = ms.Tensor(_init_random_normal(mean, std, self.shape), ms.float32) return cast_to_adapter_tensor(output) def normal_(self, mean=0, std=1, *, generator=None): output = self.normal_adapter(mean, std, generator=generator) return _tensor_inplace_assign(self, output, "normal_", "normal_adapter") def size(self, dim=None): """ tensor.size() has the same function as tensor.size() in PyTorch, but different from the tensor.size in MindSpore. """ if dim is None: return self.shape return self.shape[dim] def uniform_adpater(self, from_alias=0, to=1): #TODO: from_alias->from self_dtype = self.dtype output = ms.Tensor(_init_random_uniform(from_alias, to, self.shape), self_dtype) return cast_to_adapter_tensor(output) def uniform_(self, from_alias=0, to=1): output = self.uniform_adpater(from_alias, to) return _tensor_inplace_assign(self, output, "uniform_", "uniform_adpater") def random_adapter(self, from_alias=0, to=None, *, generator=None): #TODO: from_alias->from unsupported_attr(generator) if generator: raise NotImplementedError("generator is not supported.") self_dtype = self.dtype if not to: if self_dtype == ms.float64: return self.uniform_adpater(from_alias, kMantissaFloat64) elif self_dtype == ms.float32: return self.uniform_adpater(from_alias, kMantissaFloat32) elif self_dtype == ms.float16: return self.uniform_adpater(from_alias, kMantissaFloat16) elif self_dtype == ms.uint64: return self.uniform_adpater(from_alias, kMaxUint64) elif self_dtype == ms.uint32: return self.uniform_adpater(from_alias, kMaxUint32) elif self_dtype == ms.uint16: return self.uniform_adpater(from_alias, kMaxUint16) elif self_dtype == ms.uint8: return self.uniform_adpater(from_alias, kMaxUint8) elif self_dtype == ms.int64: return self.uniform_adpater(from_alias, kMaxInt64) elif self_dtype == ms.int32: return self.uniform_adpater(from_alias, kMaxInt32) elif self_dtype == ms.int16: return self.uniform_adpater(from_alias, kMaxInt16) elif self_dtype == ms.int8: return self.uniform_adpater(from_alias, kMaxInt8) return self.uniform_adpater(from_alias, to) def random_(self, from_alias=0, to=None, *, generator=None): output = self.random_adapter(from_alias, to, generator=generator) return _tensor_inplace_assign(self, output, "random_", "random_adapter") def zero_adapter(self): output = tensor_operator_registry.get("fill")(self.dtype, self.shape, 0.0) return cast_to_adapter_tensor(output) def zero_(self): output = self.zero_adapter() return _tensor_inplace_assign(self, output, "zero_", "zero_adapter") def new_zeros(self, size, *, dtype=None, device=None, requires_grad=False, layout=None, pin_memory=False): unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(layout) if layout: raise NotImplementedError("layout is not supported.") unsupported_attr(pin_memory) if pin_memory is True: raise NotImplementedError("pin_memory is not supported to True.") output = tensor_operator_registry.get("fill")(dtype, size, 0.0) return cast_to_adapter_tensor(output) def new_full(self, size, fill_value, *, dtype=None, device=None, requires_grad=False, layout=None, pin_memory=False): unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(layout) if layout: raise NotImplementedError("layout is not supported.") unsupported_attr(pin_memory) if pin_memory is True: raise NotImplementedError("pin_memory is not supported to True.") if not dtype: dtype = self.dtype output = tensor_operator_registry.get("fill")(dtype, size, fill_value) return cast_to_adapter_tensor(output) def add(self, other, *, alpha=1): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = ms.ops.add(input, other*alpha) return cast_to_adapter_tensor(output) def add_(self, other, *, alpha=1): output = self.add(other, alpha=alpha) return _tensor_inplace_assign(self, output, "add_", "add") def erfinv(self): input = cast_to_ms_tensor(self) output = ms.ops.erfinv(input) return cast_to_adapter_tensor(output) def erfinv_(self): output = self.erfinv() return _tensor_inplace_assign(self, output, "erfinv_", "erfinv") def permute(self, *dims): ms_input = cast_to_ms_tensor(self) output = ms_input.transpose(*dims) return cast_to_adapter_tensor(output) def contiguous(self, memory_format=None): #TODO unsupported_attr(memory_format) return self def new_tensor(self, data, *, dtype=None, device=None, requires_grad=False, layout=None, pin_memory=False): unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(layout) unsupported_attr(pin_memory) if isinstance(data, Tensor): raise ValueError("To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() " "or sourceTensor.clone().detach().requires_grad_(True), " "rather than tensor.new_tensor(sourceTensor).") return tensor(data, dtype) def copy_(self, src, non_blocking=False): unsupported_attr(non_blocking) input_ms = cast_to_ms_tensor(src) output = ms.ops.broadcast_to(input_ms, self.shape) output = output.astype(self.dtype) return _tensor_inplace_assign(self, output, "copy_", "new_tensor") def expand(self, *size): input_ms = cast_to_ms_tensor(self) if isinstance(size[0], (list, tuple)): size = size[0] out = ms.ops.broadcast_to(input_ms, size) return cast_to_adapter_tensor(out) def sigmoid(self): input = cast_to_ms_tensor(self) output = _get_cache_prim(P.Sigmoid)()(input) return cast_to_adapter_tensor(output) def sigmoid_(self): output = self.sigmoid() return _tensor_inplace_assign(self, output, "sigmoid_", "sigmoid") def float(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.float()) def flip(self, dims): # TODO ms.numpy.flip -> Tensor.flip input_ms = cast_to_ms_tensor(self) output = ms.numpy.flip(input_ms, dims) return cast_to_adapter_tensor(output) def sign(self): input = cast_to_ms_tensor(self) output = _get_cache_prim(P.Sign)()(input) return cast_to_adapter_tensor(output) def mul(self, value): input = cast_to_ms_tensor(self) ms_value = cast_to_ms_tensor(value) output = ms.ops.mul(input, ms_value) return cast_to_adapter_tensor(output) def mul_(self, value): output = self.mul(value) return _tensor_inplace_assign(self, output, "mul_", "mul") def device(self): #TODO pass def div(self, value, *, rounding_mode=None) : output = _div_calcu(self, value, rounding_mode) return cast_to_adapter_tensor(output) def div_(self, value, *, rounding_mode=None): output = _div_calcu(self, value, rounding_mode) return _tensor_inplace_assign(self, output, "div_", "div") def cpu(self): #TODO return self def min(self, dim=None, keepdim=False): input = cast_to_ms_tensor(self) if dim is None: return cast_to_adapter_tensor(input.min()) #TODO # Until now, P.min do not support when `input` is type of `int32`, `int64``. if self.dtype == mstype.int64 or self.dtype == mstype.int32: if self.dtype == mstype.int64: dtype_name = 'torch.int64' else: dtype_name = 'torch.int32' raise TypeError("For 'Tensor.min', the type of `input` do not support `torch.int64` and " "`torch.int32`, got {}.".format(dtype_name)) indices, result = P.min(input, axis=dim, keep_dims=keepdim) return cast_to_adapter_tensor(result), cast_to_adapter_tensor(indices) def max(self, dim=None, keepdim=False): input = cast_to_ms_tensor(self) if dim is None: return cast_to_adapter_tensor(input.max()) # TODO: Until now, P.max do not support when `input` is type of `int32`, `int64``. if self.dtype == mstype.int64 or self.dtype == mstype.int32: if self.dtype == mstype.int64: dtype_name = 'torch.int64' else: dtype_name = 'torch.int32' raise TypeError("For 'Tensor.max', the type of `input` do not support `torch.int64` and " "`torch.int32`, got {}.".format(dtype_name)) indices, result = P.max(input, axis=dim, keep_dims=keepdim) return cast_to_adapter_tensor(result), cast_to_adapter_tensor(indices) def numel(self): input = cast_to_ms_tensor(self) return P.size(input) def detach(self): input_ms = cast_to_ms_tensor(self) output = ms.ops.stop_gradient(input_ms) return cast_to_adapter_tensor(output) def sum(self, dim=None, keepdim=False, dtype=None): input = cast_to_ms_tensor(self) if not dtype and self.dtype in (mstype.uint8, mstype.uint16, mstype.uint32, mstype.int8, mstype.int16, mstype.int32): dtype = mstype.int64 input = input.astype(dtype) return cast_to_adapter_tensor(input.sum(axis=dim, dtype=dtype, keepdims=keepdim)) def mean(self, dim=None, keepdim=False, dtype=None): if dim is None: axis = () else: axis = dim input = cast_to_adapter_tensor(self) if dtype: input = self.astype(dtype) output = ms.ops.mean(input, axis, keepdim) return cast_to_adapter_tensor(output) def prod(self, dim=None, keepdim=False, dtype=None): if dim is None: axis = () else: axis = dim input = cast_to_adapter_tensor(self) if dtype: input = self.astype(dtype) output = ms.ops.prod(input, axis, keepdim) return cast_to_adapter_tensor(output) def split(self, split_size, dim=0): tensor = cast_to_ms_tensor(self) output = ms.ops.split(tensor, split_size, dim) return cast_to_adapter_tensor(output) def numpy(self): return self.asnumpy() def view(self, *shape): self._init_check() if not shape: raise ValueError("The shape variable should not be empty") if isinstance(shape[0], (tuple, list)): if len(shape) != 1: raise ValueError(f"Only one tuple is needed, but got {shape}") shape = shape[0] if isinstance(shape, list): shape = tuple(shape) input_size = self.shape if input_size[0] == 0: # only support first element is 0 numel = ms.ops.size(self) shape = _infer_size(shape, numel) output = ms.ops.zeros(shape, self.dtype) else: input = cast_to_ms_tensor(self) output = tensor_operator_registry.get('reshape')()(input, shape) return cast_to_adapter_tensor(output) def view_as(self, other): return self.view(other.shape) def ndimension(self): input_ms = cast_to_ms_tensor(self) return input_ms.ndimension() def pow(self, exponent): power = cast_to_ms_tensor(exponent) input_ms = cast_to_ms_tensor(self) output = input_ms.pow(power) return cast_to_adapter_tensor(output) def repeat(self, *sizes): input_x = cast_to_ms_tensor(self) if isinstance(sizes[0], (tuple, list)): output = ms.ops.tile(input_x, *sizes) else: output = ms.ops.tile(input_x, sizes) return cast_to_adapter_tensor(output) def repeat_interleave(self, repeats, dim=None, *, output_size=None): unsupported_attr(output_size) if isinstance(repeats, Tensor): new_repeats = [] for index in repeats: new_repeats.append(int(index)) repeats = new_repeats input_ms = cast_to_ms_tensor(self) output = input_ms.repeat(repeats, dim) return cast_to_adapter_tensor(output) def reshape(self, *shape): input_ms = cast_to_ms_tensor(self) input_size = input_ms.shape if input_size[0] == 0: # only support first element is 0 numel = ms.ops.size(input_ms) shape = _infer_size(shape, numel) output = ms.ops.zeros(shape, input_ms.dtype) else: output = input_ms.reshape(*shape) return cast_to_adapter_tensor(output) def reshape_as(self, other): return self.reshape(other.shape) def arcsinh(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.arcsinh()) def arctanh(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.arctanh()) def det(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.det()) def negative(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.negative()) def negative_(self): output = self.negative() return _tensor_inplace_assign(self, output, "negative_", "negative") def abs(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.abs()) def abs_(self): output = self.abs() return _tensor_inplace_assign(self, output, "abs_", "abs") @property def ndim(self): return len(self.shape) def amax(self, dim=None, keepdim=False): input_ms = cast_to_ms_tensor(self) if dim is not None: return cast_to_adapter_tensor(input_ms.amax(axis=dim, keep_dims=keepdim)) return cast_to_adapter_tensor(input_ms.amax(keep_dims=keepdim)) def amin(self, dim=None, keepdim=False): input_ms = cast_to_ms_tensor(self) if dim is not None: return cast_to_adapter_tensor(input_ms.amin(axis=dim, keep_dims=keepdim)) return cast_to_adapter_tensor(input_ms.amin(keep_dims=keepdim)) def as_strided(self, size, stride, storage_offset=None): warnings.warn("not support output as a view.") input_ms = cast_to_ms_tensor(self) if len(size) != len(stride): raise RuntimeError("mismatch in length of strides and shape.") index = np.arange(0, size[0]*stride[0], stride[0]) for i in range(1, len(size)): tmp = np.arange(0, size[i]*stride[i], stride[i]) index = np.expand_dims(index, -1) index = index + tmp if storage_offset is not None: index = index + storage_offset input_indices = ms.Tensor(index) out = ms.ops.gather(input_ms.reshape(-1), input_indices, 0) return cast_to_adapter_tensor(out) def bmm(self, batch2): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.bmm(batch2)) def clamp(self, min=None, max=None): input_ms = cast_to_ms_tensor(self) type = input_ms.dtype if min is not None and max is not None and min > max: output = ms.ops.ones_like(input_ms).astype(type)*max else: if min is not None: min = ms.Tensor(min, type) if max is not None: max = ms.Tensor(max, type) output = ms.ops.clip_by_value(input_ms, min, max) return cast_to_adapter_tensor(output) def clamp_(self, min=None, max=None): output = self.clamp(min, max) return _tensor_inplace_assign(self, output, "clamp_", "clamp") def dim(self): return len(self.shape) def expand_as(self, other): input_ms = cast_to_ms_tensor(self) output = input_ms.expand_as(other) return cast_to_adapter_tensor(output) def item(self): input_ms = cast_to_ms_tensor(self) if input_ms.size > 1: raise ValueError("only one element tensors can be converted to Python scalars") output = input_ms.reshape(-1).asnumpy().tolist() return output[0] def log(self): input_ms = cast_to_ms_tensor(self) output = input_ms.log() return cast_to_adapter_tensor(output) def log2(self): input = cast_to_ms_tensor(self) output = ms.ops.log2(input) return cast_to_adapter_tensor(output) def matmul(self, tensor2): input_ms = cast_to_ms_tensor(self) tensor2_ms = cast_to_ms_tensor(tensor2) output = ms.ops.matmul(input_ms, tensor2_ms) return cast_to_adapter_tensor(output) def squeeze(self, dim=None): input_ms = cast_to_ms_tensor(self) if dim is not None: if input_ms.shape[dim] != 1: output = input_ms else: output = ms.ops.squeeze(input_ms, dim) else: output = ms.ops.squeeze(input_ms) return cast_to_adapter_tensor(output) def squeeze_(self, dim=None): output = self.squeeze(dim) return _tensor_inplace_assign(self, output, "squeeze_", "squeeze") def stride(self, dim=None): input_ms = cast_to_ms_tensor(self) bytelen = input_ms.nbytes//input_ms.size output = list(input_ms.strides) for i in range(len(output)): output[i] = output[i]//bytelen output = tuple(output) if dim is not None: output = output[dim] return output def sub(self, other, *, alpha=1): input_ms = cast_to_ms_tensor(self) input_other = cast_to_ms_tensor(other) * alpha output = ms.ops.sub(input_ms, input_other) return cast_to_adapter_tensor(output) def sub_(self, other, *, alpha=1): output = self.sub(other, alpha=alpha) return _tensor_inplace_assign(self, output, "sub_", "sub") # TODO: delete it, apply ms.Tensor.is_floating_point def is_floating_point(self): return self._dtype in (mstype.float16, mstype.float32, mstype.float64) def unbind(self, dim=0): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.unbind(dim)) def unsqueeze(self, dim): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.unsqueeze(dim)) def unsqueeze_(self, dim): output = self.unsqueeze(dim) return _tensor_inplace_assign(self, output, "unsqueeze_", "unsqueeze") def is_signed(self): # input_ms = cast_to_ms_tensor(self) # return input_ms.is_signed() #TODO mindspore 11/17 2.0nightly supported pass def transpose(self, dim0, dim1): input_ms = cast_to_ms_tensor(self) dims = list(range(input_ms.ndim)) dims[dim0], dims[dim1] = dim1, dim0 output = input_ms.transpose(dims) return cast_to_adapter_tensor(output) def transpose_(self, dim0, dim1): output = self.transpose(dim0, dim1) return _tensor_inplace_assign(self, output, "transpose_", "transpose") def floor(self): input_ms = cast_to_ms_tensor(self) output = input_ms.floor() return cast_to_adapter_tensor(output) def floor_(self): output = self.floor() return _tensor_inplace_assign(self, output, "floor_", "floor") def isfinite(self): input_ms = cast_to_ms_tensor(self) output = ms.ops.isfinite(input_ms) return cast_to_adapter_tensor(output) def isnan(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.isnan()) def clone(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.copy()) def to(self, *args, **kwargs): # TODO: # Note that this API requires the user to ensure the correctness of the input currently, # and only the function of modifying dtype is available. if len(args) == 0 and len(kwargs) == 0: raise ValueError("Tensor.to is missing inputs, please check.") input_ms = cast_to_ms_tensor(self) if "dtype" in kwargs: set_dtype = kwargs.get("dtype") return cast_to_adapter_tensor(input_ms.astype(set_dtype)) elif "other" in kwargs: set_dtype = kwargs.get("other")._dtype return cast_to_adapter_tensor(input_ms.astype(set_dtype)) elif "device" in kwargs: return self if len(args) == 0: raise ValueError("The inputs of Tensor.to is abnormal, please check.") if args[0] in _dtypeDict.values(): return cast_to_adapter_tensor(input_ms.astype(args[0])) elif isinstance(args[0], Tensor): set_dtype = args[0]._dtype return cast_to_adapter_tensor(input_ms.astype(set_dtype)) elif not isinstance(args[0], (str, Device)): raise ValueError("The inputs of Tensor.to is abnormal, please check.") if len(args) > 1 and args[1] in _dtypeDict.values(): return cast_to_adapter_tensor(input_ms.astype(args[1])) return self def sort(self, dim=-1, descending=False): # TODO: ops.sort() should be replaced. input_ms = cast_to_ms_tensor(self) input_type = input_ms.dtype if 'Int' in str(input_type): input_ms = input_ms.astype(ms.float32) sort_tensor, sort_index = ms.ops.Sort(dim, descending)(input_ms) sort_tensor = sort_tensor.astype(input_type) sort_index = sort_index.astype(ms.int64) return cast_to_adapter_tensor((sort_tensor, sort_index)) else: output = _get_cache_prim(ms.ops.Sort)(dim, descending)(input_ms) return cast_to_adapter_tensor(output) def msort(self): # TODO: ops.sort() should be replaced. input_ms = cast_to_ms_tensor(self) sort_op = _get_cache_prim(ms.ops.Sort)(axis=0) input_type = input_ms.dtype if 'Int' in str(input_type): input_ms = input_ms.astype(ms.float32) output, _ = sort_op(input_ms) output = output.astype(input_type) else: output, _ = sort_op(input_ms) return cast_to_adapter_tensor(output) def argsort(self, dim=-1, descending=False): # TODO: ops.sort() should be replaced. input_ms = cast_to_ms_tensor(self) sort_op = _get_cache_prim(ms.ops.Sort)(dim, descending) input_type = input_ms.dtype if 'Int' in str(input_type): input_ms = input_ms.astype(ms.float32) _, output = sort_op(input_ms) output = output.astype(ms.int64) else: _, output = sort_op(input_ms) return cast_to_adapter_tensor(output) def sqrt(self): input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(ms.ops.sqrt(input_ms)) def sqrt_(self): output = self.sqrt() return _tensor_inplace_assign(self, output, "sqrt_", "sqrt") def rsqrt(self): input = cast_to_ms_tensor(self) if 'Bool' in str(input.dtype) or 'Int' in str(input.dtype): input = input.astype(ms.float32) output = _get_cache_prim(ms.ops.Rsqrt)()(input) return cast_to_adapter_tensor(output) def rsqrt_(self): output = self.rsqrt() return _tensor_inplace_assign(self, output, "rsqrt_", "rsqrt") def resize(self, *size, memory_format=None): unsupported_attr(memory_format) input = cast_to_ms_tensor(self) input_size = input.shape if len(input_size) == 1 and input_size[0] == 0: out = ms.ops.zeros(size, self.dtype) elif len(size) > 0 and isinstance(size[0], tuple): out = input.resize(size[0]) else: out = input.resize(size) return cast_to_adapter_tensor(out) def resize_(self, *size, memory_format=None): output = self.resize(*size, memory_format=memory_format) return _tensor_inplace_assign(self, output, "resize_", "resize") def resize_as(self, tensor, memory_format=None): unsupported_attr(memory_format) if not isinstance(tensor, Tensor): raise TypeError("resize_as(): argument 'tensor' must be Tensor.") input = cast_to_ms_tensor(self) size = tensor.shape input_size = input.shape if len(input_size) == 1 and input_size[0] == 0: out = ms.ops.zeros(size, self.dtype) else: out = input.resize(size) return cast_to_adapter_tensor(out) def resize_as_(self, tensor, memory_format=None): output = self.resize_as(tensor, memory_format) return _tensor_inplace_assign(self, output, "resize_as_", "resize_as") def index_fill(self, dim, index, value): input = cast_to_ms_tensor(self) index = cast_to_ms_tensor(index) index = ms.ops.cast(index, mstype.int32) if is_under_ascend_context(): raise NotImplementedError("for adapter, index_fill not supported on ascend.") out = input.index_fill(dim, index, value) return cast_to_adapter_tensor(out) def index_fill_(self, dim, index, value): output = self.index_fill(dim, index, value) return _tensor_inplace_assign(self, output, "index_fill_", "index_fill") def index_select(self, dim, index): _input_params = cast_to_ms_tensor(self) _input_indices = cast_to_ms_tensor(index) output = ms.ops.gather(_input_params, _input_indices, dim) return cast_to_adapter_tensor(output) @property def data(self): return self.detach() def new(self, *size): if len(size) > 0 and isinstance(size[0], tuple): size = size[0] return Tensor(*size, dtype=self.dtype) def cuda(self, device=None, non_blocking=False, memory_format=None): unsupported_attr(device) unsupported_attr(non_blocking) unsupported_attr(memory_format) if not is_under_gpu_context(): backend = get_backend() warning = f"MsAdater.pytorch.Tensor.cuda() didn't work because it is under {backend} context." warnings.warn(warning) return self def is_cuda(self): return is_under_gpu_context() def le(self, other): input = cast_to_ms_tensor(self) if isinstance(other, Tensor): other = cast_to_ms_tensor(other) out = ms.ops.le(input, other) return cast_to_adapter_tensor(out) def le_(self, other): output = self.le(other) return _tensor_inplace_assign(self, output, "le_", "le") def t(self): input_ms = cast_to_ms_tensor(self) if input_ms.ndim > 2: raise ValueError("t() expects a tensor with <= 2 dimensions, but self is {}D".format(input_ms.ndim)) dims = list(range(input_ms.ndim)).reverse() output = input_ms.transpose(dims) return cast_to_adapter_tensor(output) @property def T(self): input_ms = cast_to_ms_tensor(self) if input_ms.ndim <= 2: warning = ("The use of Tensor.T() on tensors of dimension other than 2 to reverse " "their shape is deprecated and it will throw an error in a future release. ") warnings.warn(warning) dims = list(range(input_ms.ndim)).reverse() output = input_ms.transpose(dims) return cast_to_adapter_tensor(output) @property def requires_grad(self): return True def requires_grad_(self, requires_grad=True): if requires_grad is False: warnings.warn("requires_grad is always True in Tensor.") return self def nonzero(self, *, out=None, as_tuple=False): if out is not None: warnings.warn("Do not support parameter 'out'.") input = cast_to_ms_tensor(self) output = None if as_tuple: if input.ndim == 1: res = ms.ops.nonzero(input) output = (cast_to_adapter_tensor(res.flatten()),) elif input.ndim > 1: output = [] res = ms.ops.nonzero(input) res = res.transpose(1, 0) res = ms.ops.split(res, input.ndim, axis=0) for cur in res: output.append(cast_to_adapter_tensor(cur)) output = tuple(output) elif input.ndim == 0: raise ValueError("Do not support input ndim == 0.") return output return cast_to_adapter_tensor(ms.ops.nonzero(input)) def bool(self, memory_format=None): unsupported_attr(memory_format) input = cast_to_ms_tensor(self) output = input.bool() return cast_to_adapter_tensor(output) def eq(self, other): input_ms = cast_to_ms_tensor(self) other_ms = cast_to_ms_tensor(other) output = input_ms.equal(other_ms) return cast_to_adapter_tensor(output) def eq_(self, other): output = self.eq(other) return _tensor_inplace_assign(self, output, "eq_", "eq") def std(self, dim=None, unbiased=True, keepdim=False): #TODO: not support float64 or complex input input_ms = cast_to_ms_tensor(self) type_float64 = False if input_ms.dtype == ms.float64: input_ms = input_ms.astype(ms.float32) type_float64 = True # TODO: mindspore.ops.std() not supported GPU, use tensor.std() instead, which means ms.ops.var(). if is_under_gpu_context(): _dim = dim if dim is not None else () _ddof = 1 if unbiased else 0 output = input_ms.std(_dim, _ddof, keepdim) else: if dim is not None: output, _ = ms.ops.std(input_ms, dim, unbiased, keepdim) else: output, _ = ms.ops.std(input_ms, unbiased=unbiased, keep_dims=keepdim) if type_float64: output = output.astype(ms.float64) return cast_to_adapter_tensor(output) def exp(self): input_ms = cast_to_ms_tensor(self) output = input_ms.exp() return cast_to_adapter_tensor(output) def masked_fill(self, mask, value): input_ms = cast_to_ms_tensor(self) output = input_ms.masked_fill(mask, value) return cast_to_adapter_tensor(output) def masked_fill_(self, mask, value): output = self.masked_fill(mask, value) return _tensor_inplace_assign(self, output, "masked_fill_", "masked_fill") def tolist(self): return self.numpy().tolist() def bernoulli(self, *, generator=None): unsupported_attr(generator) if generator: raise NotImplementedError("generator is not supported.") input_ms = cast_to_ms_tensor(self) bernoulli_seed = ms.get_seed() if not bernoulli_seed: bernoulli_seed = -1 return cast_to_adapter_tensor(input_ms.bernoulli(input_ms, bernoulli_seed)) def bernoulli_(self, p=0.5, *, generator=None): output = self.bernoulli_adapter(p, generator=generator) return _tensor_inplace_assign(self, output, "bernoulli_", "bernoulli_adapter") def bernoulli_adapter(self, p=0.5, *, generator=None): unsupported_attr(generator) if generator: raise NotImplementedError("generator is not supported.") input_ms = cast_to_ms_tensor(self) bernoulli_seed = ms.get_seed() if not bernoulli_seed: bernoulli_seed = -1 return cast_to_adapter_tensor(input_ms.bernoulli(p, bernoulli_seed)) def round(self, decimals=0): input = cast_to_ms_tensor(self) if decimals == 0: output = ms.ops.round(input) else: p = 10 ** decimals input = input * p output = ms.ops.round(input) / p return cast_to_adapter_tensor(output) def long(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.astype(_dtypeDict["long"])) def half(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.astype(_dtypeDict["half"])) def int(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.int()) def double(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.astype(_dtypeDict["double"])) def char(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.astype(_dtypeDict["char"])) def byte(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.astype(_dtypeDict["byte"])) def short(self, memory_format=None): unsupported_attr(memory_format) if memory_format: raise NotImplementedError("memory_format is not supported.") input_ms = cast_to_ms_tensor(self) return cast_to_adapter_tensor(input_ms.astype(_dtypeDict["short"])) def chunk(self, chunks, dim=0): input = cast_to_ms_tensor(self) output = ms.ops.chunk(input, chunks, dim) return cast_to_adapter_tensor(output) def flatten(self, start_dim=0, end_dim=-1): @constexpr def get_dst_shape(): self_shape = self.shape rank = len(self_shape) start = start_dim end = end_dim if start < 0: start += rank if end < 0: end += rank dst_shape = [] i = 0 while i != start: dst_shape.append(self_shape[i]) i = i + 1 flatten_shape = 1 while i <= end: flatten_shape = flatten_shape * self_shape[i] i = i + 1 dst_shape.append(flatten_shape) while i < rank: dst_shape.append(self_shape[i]) i = i + 1 return tuple(dst_shape) shape = get_dst_shape() input_ms = cast_to_ms_tensor(self) input_ms.reshape(shape) return cast_to_adapter_tensor(input_ms.reshape(shape)) def sin(self): input = cast_to_ms_tensor(self) return cast_to_adapter_tensor(ms.ops.sin(input)) def sin_(self): output = self.sin() return _tensor_inplace_assign(self, output, "sin_", "sin") def ge(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = input.ge(other) return cast_to_adapter_tensor(output) def ge_(self, other): output = self.ge(other) return _tensor_inplace_assign(self, output, "ge_", "ge") def cumsum(self, dim, dtype=None): input = cast_to_ms_tensor(self) output = input.cumsum(axis=dim, dtype=dtype) return cast_to_adapter_tensor(output) def absolute(self): return self.abs() def absolute_(self): output = self.abs() return _tensor_inplace_assign(self, output, "absolute_", "absolute") def acos(self): input = cast_to_ms_tensor(self) output = ms.ops.acos(input) return cast_to_adapter_tensor(output) def acos_(self): output = self.acos() return _tensor_inplace_assign(self, output, "acos_", "acos") def arccos(self): return self.acos() def arccos_(self): output = self.acos() return _tensor_inplace_assign(self, output, "arccos_", "arccos") def asinh(self): input_ms = cast_to_ms_tensor(self) output = ms.ops.asinh(input_ms) return cast_to_adapter_tensor(output) def asinh_(self): output = self.asinh() return _tensor_inplace_assign(self, output, "asinh_", "asinh") def atanh(self): input_ms = cast_to_ms_tensor(self) output = ms.ops.atanh(input_ms) return cast_to_adapter_tensor(output) def atanh_(self): output = self.atanh() return _tensor_inplace_assign(self, output, "atanh_", "atanh") def addcdiv(self, tensor1, tensor2, *, value=1): input = cast_to_ms_tensor(self) tensor1 = cast_to_ms_tensor(tensor1) tensor2 = cast_to_ms_tensor(tensor2) value = ms.Tensor(value) output = ms.ops.addcdiv(input, tensor1, tensor2, value) return cast_to_adapter_tensor(output) def addcdiv_(self, tensor1, tensor2, *, value=1): output = self.addcdiv(tensor1, tensor2, value=value) return _tensor_inplace_assign(self, output, "addcdiv_", "addcdiv") def gather(self, dim, index): input = cast_to_ms_tensor(self) index = cast_to_ms_tensor(index) output = ms.ops.gather_elements(input, dim, index) return cast_to_adapter_tensor(output) def fmod(self, divisor): x = cast_to_ms_tensor(self) other = cast_to_ms_tensor(divisor) #TODO: repalce with ms.ops.fmod if not (isinstance(x, (Tensor, Tensor_)) or isinstance(other, (Tensor, Tensor_))): raise TypeError("At least one of the types of inputs must be tensor, " + \ f"but the type of 'x' got is {type(x)}, " + \ f"and the type of 'other' is {type(other)}.") return x - ms.ops.div(x, other, rounding_mode="trunc") * other def fmod_(self, divisor): output = self.fmod(divisor) return _tensor_inplace_assign(self, output, "fmod_", "fmod") def lt(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = ms.ops.less(input, other) return cast_to_adapter_tensor(output) def lt_(self, other): output = self.lt(other) return _tensor_inplace_assign(self, output, "lt_", "lt") def less(self, other): return self.lt(other) def less_(self, other): output = self.lt(other) return _tensor_inplace_assign(self, output, "less_", "less") def less_equal(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = ms.ops.less_equal(input, other) return cast_to_adapter_tensor(output) def less_equal_(self, other): output = self.less_equal(other) return _tensor_inplace_assign(self, output, "less_equal_", "less_equal") def ne(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = ms.ops.ne(input, other) return cast_to_adapter_tensor(output) def ne_(self, other): output = self.ne(other) return _tensor_inplace_assign(self, output, "ne_", "ne") def not_equal(self, other): return self.ne(other) def not_equal_(self, other): output = self.ne(other) return _tensor_inplace_assign(self, output, "not_equal_", "not_equal") def equal(self, other): if not isinstance(other, Tensor): raise ValueError("`other` must be Tensor") x = cast_to_ms_tensor(self) y = cast_to_ms_tensor(other) if x.dtype != y.dtype: return False if x.shape == y.shape: size = x.size output = ms.ops.equal(x, y) output = output.sum() if output == size: return True return False def greater(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = ms.ops.greater(input, other) return cast_to_adapter_tensor(output) def greater_(self, other): output = self.greater(other) return _tensor_inplace_assign(self, output, "greater_", "greater") def gt(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = ms.ops.gt(input, other) return cast_to_adapter_tensor(output) def gt_(self, other): output = self.greater(other) return _tensor_inplace_assign(self, output, "gt_", "gt") def greater_equal(self, other): x = cast_to_ms_tensor(self) y = cast_to_ms_tensor(other) output = ms.ops.greater_equal(x, y) return cast_to_adapter_tensor(output) def greater_equal_(self, other): output = self.greater_equal(other) return _tensor_inplace_assign(self, output, "greater_equal_", "greater_equal") def argmin(self, dim=None, keepdim=False): input = cast_to_ms_tensor(self) # TODO: output = ms.ops.argmin(input, axis=dim, keepdims=keepdim) if keepdim: raise NotImplementedError("keepdim is not supported.") # TODO: ascend not support argmin if is_under_ascend_context(): input = input * -1 output = ms.ops.argmax(input, axis=dim) else: output = ms.ops.argmin(input, axis=dim) return cast_to_adapter_tensor(output) def argmax(self, dim=None, keepdim=False): input = cast_to_ms_tensor(self) # TODO: output = ms.ops.argmax(input, axis=dim, keepdims=keepdim) if keepdim: raise NotImplementedError("keepdim is not supported.") output = ms.ops.argmax(input, axis=dim) return cast_to_adapter_tensor(output) def type(self, dtype=None, non_blocking=False, **kwargs): def _get_type_from_dtype(dtype): str_dtype = str(dtype).split('.')[-1].lower() _type = _dtype2typeDict.get(str_dtype) return _type def _get_dtype_from_type(type): _dtype = _type2dtypeDict.get(type, 'None') if _dtype == 'None': _dtype = type return _dtype unsupported_attr(non_blocking) unsupported_attr(kwargs) if dtype is None: return _get_type_from_dtype(self.dtype) _dtype = _get_dtype_from_type(dtype) if _dtype == self.dtype: return self x = cast_to_ms_tensor(self) output = x.astype(_dtype) return cast_to_adapter_tensor(output) def type_as(self, tensor): if self.dtype == tensor.dtype: return self x = cast_to_ms_tensor(self) output = x.astype(tensor.dtype) return cast_to_adapter_tensor(output) def get_device(self): return -1 def baddbmm(self, batch1, batch2, *, beta=1, alpha=1): x = cast_to_ms_tensor(self) batch1 = cast_to_ms_tensor(batch1) batch2 = cast_to_ms_tensor(batch2) output = ms.ops.baddbmm(x, batch1, batch2, beta, alpha) return cast_to_adapter_tensor(output) def baddbmm_(self, batch1, batch2, *, beta=1, alpha=1): output = self.baddbmm(batch1, batch2, beta=beta, alpha=alpha) return _tensor_inplace_assign(self, output, "baddbmm_", "baddbmm") def topk(self, k, dim=None, largest=True, sorted=True): unsupported_attr(dim) unsupported_attr(largest) input = cast_to_ms_tensor(self) output = input.top_k(k, sorted=sorted) return cast_to_adapter_tensor(output) def maximum(self, other): x = cast_to_ms_tensor(self) y = cast_to_ms_tensor(other) #TODO: NAN is different output = ms.ops.maximum(x, y) return cast_to_adapter_tensor(output) def minimum(self, other): x = cast_to_ms_tensor(self) y = cast_to_ms_tensor(other) #TODO: NAN is different output = ms.ops.minimum(x, y) return cast_to_adapter_tensor(output) def multiply(self, value): x = cast_to_ms_tensor(self) y = cast_to_ms_tensor(value) output = ms.ops.mul(x, y) return cast_to_adapter_tensor(output) def multiply_(self, value): output = self.multiply(value) return _tensor_inplace_assign(self, output, "multiply_", "multiply") def neg(self): x = cast_to_ms_tensor(self) output = ms.ops.neg(x) return cast_to_adapter_tensor(output) def neg_(self): output = self.neg() return _tensor_inplace_assign(self, output, "neg_", "neg") def ravel(self): x = cast_to_ms_tensor(self) output = x.ravel() return cast_to_adapter_tensor(output) def select(self, dim, index): input = cast_to_ms_tensor(self) _input_indices = ms.Tensor(index) output = ms.ops.gather(input, _input_indices, dim) @constexpr def _get_out_shape(input_shape, dim): shape = [input_shape[i] for i in range(len(input_shape)) if i != dim] return tuple(shape) output_shape = _get_out_shape(input.shape, dim) output = output.reshape(output_shape) return cast_to_adapter_tensor(output) def square(self): x = cast_to_ms_tensor(self) output = ms.ops.square(x) return cast_to_adapter_tensor(output) def broadcast_to(self, shape): input = cast_to_ms_tensor(self) output = ms.ops.broadcast_to(input, shape) return cast_to_adapter_tensor(output) def divide(self, value, *, rounding_mode=None) : output = _div_calcu(self, value, rounding_mode) return cast_to_adapter_tensor(output) def divide_(self, value, *, rounding_mode=None) : output = _div_calcu(self, value, rounding_mode) return _tensor_inplace_assign(self, output, "divide_", "divide") def unique(self, sorted=True, return_inverse=False, return_counts=False, dim=None): unsupported_attr(dim) unsupported_attr(return_counts) input = cast_to_ms_tensor(self) data_type = input.dtype if sorted and return_inverse: raise ValueError("Don't support sorted=True and return_inverse=True.") res, idx = ms.ops.unique(input) if sorted: res = ms.ops.cast(res, ms.float32) res, _ = ms.ops.sort(res) res = ms.ops.cast(res, data_type) if return_inverse: res = cast_to_adapter_tensor(res) idx = cast_to_adapter_tensor(idx) return (res, idx) else: res = cast_to_adapter_tensor(res) return res def mm(self, mat2): input = cast_to_ms_tensor(self) input2 = cast_to_ms_tensor(mat2) input_type = input.dtype if input_type in (mstype.int32,mstype.int64): input = self.astype(mstype.float32) output = ms.ops.matmul(input, input2) output = ms.ops.cast(output, input_type) else: output = ms.ops.matmul(input, input2) return cast_to_adapter_tensor(output) def logsumexp(self, dim, keepdim=False): ms_input = cast_to_ms_tensor(self) if ms_input.dtype != mstype.float32: ms_input = ms_input.astype(mstype.float32) output = ms.ops.logsumexp(ms_input, dim, keepdim) return cast_to_adapter_tensor(output) def addmv(self, mat, vec, *, beta=1, alpha=1): input = cast_to_ms_tensor(self) mat = cast_to_ms_tensor(mat) vec = cast_to_ms_tensor(vec) output = ms.ops.addmv(input, mat, vec, beta=beta, alpha=alpha) return cast_to_adapter_tensor(output) def dot(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) #TODO: ms.ops.tensor_dot only supports float16/float32 input_dtype = input.dtype if input_dtype in (mstype.float32, mstype.float16): output = ms.ops.tensor_dot(input, other, 1) else: input = input.astype(ms.float32) other = other.astype(ms.float32) output = ms.ops.tensor_dot(input, other, 1) output = output.astype(input_dtype) return cast_to_adapter_tensor(output) def inverse(self): input = cast_to_ms_tensor(self) #TODO: Ascend has no ms.ops.inverse # output = ms.ops.inverse(input) # ops.inverse is not ready at 11-13 if self.dtype in msdapter_dtype.all_int_type: input = input.astype(mstype.float32) output = _get_cache_prim(P.MatrixInverse)()(input) return cast_to_adapter_tensor(output) def asin(self): input = cast_to_ms_tensor(self) if self.dtype in msdapter_dtype.all_int_type: input = input.astype(mstype.float32) output = ms.ops.asin(input) return cast_to_adapter_tensor(output) def asin_(self): output = self.asin() return _tensor_inplace_assign(self, output, "asin_", "asin") def atan(self): input = cast_to_ms_tensor(self) if self.dtype in msdapter_dtype.all_int_type: input = input.astype(mstype.float32) output = ms.ops.atan(input) return cast_to_adapter_tensor(output) def atan_(self): output = self.atan() return _tensor_inplace_assign(self, output, "atan_", "atan") def atan2(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) if self.dtype in msdapter_dtype.all_int_type: input = input.astype(mstype.float32) other = other.astype(mstype.float32) output = ms.ops.atan2(input, other) return cast_to_adapter_tensor(output) def atan2_(self, other): output = self.atan2(other) return _tensor_inplace_assign(self, output, "atan2_", "atan2") def count_nonzero(self, dim=None): input = cast_to_ms_tensor(self) if dim is None: dim = () output = ms.ops.count_nonzero(input, axis=dim) return cast_to_adapter_tensor(output) def scatter(self, dim, index, src, reduce=None): if not reduce: reduce = 'none' else: # TODO: to supported 'multiply' if reduce not in ('none', 'add'): raise ValueError("For tensor.scatter or scatter_, `reduce` only support 'none', " f"'add', but got '{reduce}'.") # TODO: add not supported on Ascend yet if reduce == 'add' and is_under_ascend_context(): raise NotImplementedError("For tensor.scatter or scatter_, `reduce` == 'add' not supported on Ascend") input, index, src = cast_to_ms_tensor((self, index, src)) if isinstance(src, numbers.Number): src = ms.ops.scalar_to_tensor(src, dtype=input.dtype) src = ms.ops.broadcast_to(src, index.shape) elif isinstance(src, ms.Tensor): src_shape = src.shape index_shape = index.shape if src_shape != index_shape: # TODO raise NotImplementedError("For scatter, not support src.shape != index.shape yet") else: raise TypeError(f"For scatter, `src` must be number or tensor, but got {type(src)}") if is_under_ascend_context(): input_dtype = input.dtype input = _ascend_tensor_general_cast(input) src = _ascend_tensor_general_cast(src) output = ms.ops.tensor_scatter_elements(input, index, src, dim, reduce) output = output.astype(input_dtype) else: output = ms.ops.tensor_scatter_elements(input, index, src, dim, reduce) return cast_to_adapter_tensor(output) def scatter_(self, dim, index, src, reduce=None): output = self.scatter(dim, index, src, reduce) return _tensor_inplace_assign(self, output, "scatter_", "scatter") def acosh(self, input): input = cast_to_ms_tensor(input) shape = input.shape if len(shape) > 7: input = input.flatten() output = ms.ops.acosh(input) output = output.reshape(shape) else: output = ms.ops.acosh(input) return cast_to_adapter_tensor(output) def new_ones(self, *size, dtype=None, device=None, requires_grad=False, layout=None, pin_memory=False): unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(layout) unsupported_attr(pin_memory) if dtype is None: dtype = self.dtype if isinstance(size[0], (tuple, list)): output = ms.ops.ones(*size, dtype=dtype) else: output = ms.ops.ones(size, dtype=dtype) return cast_to_adapter_tensor(output) def new_empty(self, *size, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False, memory_format=None): unsupported_attr(layout) unsupported_attr(device) unsupported_attr(requires_grad) unsupported_attr(pin_memory) unsupported_attr(memory_format) if dtype is None: dtype = self.dtype _size = size if isinstance(size[0], (tuple, list)): _size = size[0] output = ms.numpy.empty(_size, dtype) return cast_to_adapter_tensor(output) def addcmul(self, tensor1, tensor2, *, value=1): #Todo: use ms.ops.addcmul after it has been fixed input = cast_to_ms_tensor(self) tensor1 = cast_to_ms_tensor(tensor1) tensor2 = cast_to_ms_tensor(tensor2) value = ms.ops.scalar_to_tensor(value) mul = ms.ops.mul(tensor1, tensor2) * value output = ms.ops.add(input, mul) return cast_to_adapter_tensor(output) def arccosh(self): input = cast_to_ms_tensor(self) shape = input.shape if len(shape) > 7: input = input.flatten() output = ms.ops.acosh(input) output = output.reshape(shape) else: output = ms.ops.acosh(input) return cast_to_adapter_tensor(output) def arcsin(self): input = cast_to_ms_tensor(self) output = ms.ops.asin(input) return cast_to_adapter_tensor(output) def arctan(self): input = cast_to_ms_tensor(self) shape = input.shape if len(shape) > 7: input = input.flatten() output = ms.ops.atan(input) output = output.reshape(shape) else: output = ms.ops.atan(input) return cast_to_adapter_tensor(output) def arctan_(self): output = self.arctan() return _tensor_inplace_assign(self, output, "arctan_", "arctan") def arctan2(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) output = ms.ops.atan2(input, other) return cast_to_adapter_tensor(output) def arctan2_(self, other): output = self.arctan2(self, other) return _tensor_inplace_assign(self, output, "arctan2_", "arctan2") def bitwise_not(self): input = cast_to_ms_tensor(self) type = input.dtype if str(type) != 'Bool': output = 0 - input - 1 else: output = 1 - input output = output.astype(ms.bool_) return cast_to_adapter_tensor(output) def bitwise_not_(self): output = self.bitwise_not(self) return _tensor_inplace_assign(self, output, "bitwise_not_", "bitwise_not") def bitwise_and(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) if input.dtype == msdapter_dtype.bool: input = input.astype(ms.int8) output = ms.ops.bitwise_and(input, other) output = output.astype(ms.bool_) else: output = ms.ops.bitwise_and(input, other) return cast_to_adapter_tensor(output) def bitwise_and_(self, other): output = self.bitwise_and(self, other) return _tensor_inplace_assign(self, output, "bitwise_and_", "bitwise_and") def bitwise_or(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) if input.dtype == msdapter_dtype.bool: input = input.astype(ms.int8) output = ms.ops.bitwise_or(input, other) output = output.astype(ms.bool_) else: output = ms.ops.bitwise_or(input, other) return cast_to_adapter_tensor(output) def bitwise_xor(self, other): input = cast_to_ms_tensor(self) other = cast_to_ms_tensor(other) if input.dtype == msdapter_dtype.bool: input = input.astype(ms.int8) output = ms.ops.bitwise_xor(input, other) output = output.astype(ms.bool_) else: output = ms.ops.bitwise_xor(input, other) return cast_to_adapter_tensor(output) def addbmm(self, batch1, batch2, *, beta=1, alpha=1): _input, _batch1, _batch2 = cast_to_ms_tensor((self, batch1, batch2)) output = ms.ops.addbmm(_input, _batch1, _batch2, beta=beta, alpha=alpha) return cast_to_adapter_tensor(output) def addbmm_(self, batch1, batch2, *, beta=1, alpha=1): output = self.addbmm(batch1, batch2, beta=beta, alpha=alpha) return _tensor_inplace_assign(self, output, "addbmm_", "addbmm") def addmm(self, mat1, mat2, *, beta=1, alpha=1): _input, _mat1, _mat2 = cast_to_ms_tensor((self, mat1, mat2)) output = ms.ops.addmm(_input, _mat1, _mat2, beta=beta, alpha=alpha) return cast_to_adapter_tensor(output) def addr(self, vec1, vec2, *, beta=1, alpha=1): _input, _vec1, _vec2 = cast_to_ms_tensor((self, vec1, vec2)) output = ms.ops.addr(_input, _vec1, _vec2, beta=beta, alpha=alpha) return cast_to_adapter_tensor(output) def addr_(self, vec1, vec2, *, beta=1, alpha=1): output = self.addr(vec1, vec2, beta=beta, alpha=alpha) return _tensor_inplace_assign(self, output, "addr_", "addr") def all(self, dim=(), keepdim=False): input = cast_to_ms_tensor(self) if input.dtype != ms.bool_: input = input.astype(ms.bool_) output = input.all(axis=dim, keep_dims=keepdim) return cast_to_adapter_tensor(output) def isclose(self, other, rtol=1e-05, atol=1e-08, equal_nan=False): _input, _other = cast_to_ms_tensor((self, other)) output = ms.ops.isclose(_input, _other, rtol=rtol, atol=atol, equal_nan=equal_nan) return cast_to_adapter_tensor(output) def allclose(self, other, rtol=1e-05, atol=1e-08, equal_nan=False): output = self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan) output = output.all() return output.item() def cholesky(self, upper=False): input = cast_to_ms_tensor(self) output = input.cholesky(upper) return cast_to_adapter_tensor(output) def cholesky_inverse(self, upper=False): input = cast_to_ms_tensor(self) output = input.cholesky_inverse(upper) return cast_to_adapter_tensor(output) def nelement(self): input = cast_to_ms_tensor(self) output = input.nelement() return output def aminmax(self, *, dim=None, keepdim=False): _input = cast_to_ms_tensor(self) _min = _input.min(axis=dim, keepdims=keepdim) _max = _input.max(axis=dim, keepdims=keepdim) return cast_to_adapter_tensor((_min, _max)) def any(self, dim=(), keepdim=False): input = cast_to_ms_tensor(self) if input.dtype != ms.bool_: input = input.astype(ms.bool_) output = input.any(axis=dim, keep_dims=keepdim) return cast_to_adapter_tensor(output) def bincount(self, weights=None, minlength=0): input = cast_to_ms_tensor(self) type = 'int64' if input.dtype == ms.uint8: input = input.astype(ms.int16) if weights is not None: weights = cast_to_ms_tensor(weights) type = weights.dtype output = ms.numpy.bincount(input, weights, minlength).astype(type) return cast_to_adapter_tensor(output) def bitwise_left_shift(self, other): input = cast_to_ms_tensor(self).asnumpy() if isinstance(other, Tensor): other = cast_to_ms_tensor(other).asnumpy() output = Tensor(np.left_shift(input, other), inner=True) return output def bitwise_right_shift(self, other): input = cast_to_ms_tensor(self).asnumpy() if isinstance(other, Tensor): other = cast_to_ms_tensor(other).asnumpy() output = Tensor(np.right_shift(input, other), inner=True) return output def clip(self, min=None, max=None): input = cast_to_ms_tensor(self) output = input.clip(min, max) return cast_to_adapter_tensor(output) def copysign(self, other): input = cast_to_ms_tensor(self) input_type = input.dtype input = input.asnumpy() is_num = True if isinstance(other, Tensor): is_num = False other = cast_to_ms_tensor(other) other_type = other.dtype other = other.asnumpy() output = ms.Tensor(np.copysign(input, other)) if 'Int' in str(input_type): if is_num or 'Int' in str(other_type): output = output.astype(ms.float32) else: output = output.astype(other_type) elif is_num or 'Int' in str(other_type): output = output.astype(input_type) return cast_to_adapter_tensor(output) def cos(self): input = cast_to_ms_tensor(self) output = ms.ops.cos(input) return cast_to_adapter_tensor(output) def cosh(self): input = cast_to_ms_tensor(self) output = ms.ops.cosh(input) return cast_to_adapter_tensor(output) def cummax(self, dim): input = cast_to_ms_tensor(self) output = ms.ops.cummax(input, axis=dim) return cast_to_adapter_tensor(output) def cummin(self, dim): input = cast_to_ms_tensor(self) output = ms.ops.cummin(input, dim) # the output dtype in ms.ops.cummin is different with ms.ops.cummax output[1] = output[1].astype(ms.common.dtype.int64) return cast_to_adapter_tensor(output) def cumprod(self, dim, *, dtype=None): input = cast_to_ms_tensor(self) output = ms.ops.cumprod(input, dim, dtype=dtype) return cast_to_adapter_tensor(output) def deg2rad(self): input = cast_to_ms_tensor(self) if input.dtype not in (ms.float16, ms.float32, ms.float64): input = input.astype(ms.float32) output = ms.ops.deg2rad(input) return cast_to_adapter_tensor(output) def diag(self, diagonal=0): # TODO: ms.ops.diag do not support diagonal input = cast_to_ms_tensor(self) output = ms.numpy.diag(input, diagonal) return cast_to_adapter_tensor(output) def diagflat(self, offset=0): input = cast_to_ms_tensor(self) output = ms.numpy.diagflat(input, offset) return cast_to_adapter_tensor(output) def diagonal(self, offset=0, dim1=0, dim2=1): input = cast_to_ms_tensor(self) #TODO float64 not support if offset != 0 if offset != 0: input = input.astype(mstype.float32) output = ms.ops.diagonal(input, offset, dim1, dim2) return cast_to_adapter_tensor(output) def is_complex(self): input = cast_to_ms_tensor(self) return input.is_complex() def isinf(self): input = cast_to_ms_tensor(self) output = ms.ops.isinf(input) return cast_to_adapter_tensor(output) def isneginf(self): input = cast_to_ms_tensor(self) output = ms.ops.isneginf(input) return cast_to_adapter_tensor(output) def isposinf(self): input = cast_to_ms_tensor(self) output = ms.ops.isposinf(input) return cast_to_adapter_tensor(output) def isreal(self): input = cast_to_ms_tensor(self) output = ms.ops.isreal(input) return cast_to_adapter_tensor(output) class _TypeTensor(Tensor): def __init__(self, *input_data, dtype_name): super(_TypeTensor, self).__init__(*input_data, dtype=dtype_name, inner=False) class ByteTensor(_TypeTensor): def __init__(self, *input_data): super(ByteTensor, self).__init__(*input_data, dtype_name='uint8') class CharTensor(_TypeTensor): def __init__(self, *input_data): super(CharTensor, self).__init__(*input_data, dtype_name='int8') class ShortTensor(_TypeTensor): def __init__(self, *input_data): super(ShortTensor, self).__init__(*input_data, dtype_name='int16') class IntTensor(_TypeTensor): def __init__(self, *input_data): super(IntTensor, self).__init__(*input_data, dtype_name='int32') class HalfTensor(_TypeTensor): def __init__(self, *input_data): super(HalfTensor, self).__init__(*input_data, dtype_name='float16') class FloatTensor(_TypeTensor): def __init__(self, *input_data): super(FloatTensor, self).__init__(*input_data, dtype_name='float32') class DoubleTensor(_TypeTensor): def __init__(self, *input_data): super(DoubleTensor, self).__init__(*input_data, dtype_name='float64') class LongTensor(_TypeTensor): def __init__(self, *input_data): super(LongTensor, self).__init__(*input_data, dtype_name='int64') def tensor(data, dtype=None, device=None, requires_grad=True): unsupported_attr(device) if requires_grad is False: msg = ("In Adapter, Tensor's `requires_grad` is always 'True', can not be set to 'False'. ") warnings.warn(msg) return Tensor(data, dtype=dtype, inner=True) def cast_to_ms_tensor(inputs): """ Cast MSAdapter.Tensor to MindSpore.Tensor before call mindspore API. """ def _cast(inputs): if isinstance(inputs, Tensor): inputs = inner.convert_to_ms_tensor(inputs) elif isinstance(inputs, (tuple, list)): inputs = list(inputs) for id, value in enumerate(inputs): inputs[id] = _cast(value) inputs = tuple(inputs) return inputs inputs = _cast(inputs) return inputs def cast_to_adapter_tensor(outputs): """ Cast MindSpore.Tensor to MSAdapter.Tensor after call mindspore API. """ def _cast(outputs): if isinstance(outputs, (ms.Tensor, Tensor_)): outputs = inner.convert_to_adapter_tensor(outputs) elif isinstance(outputs, (tuple, list)): outputs = list(outputs) for id, value in enumerate(outputs): outputs[id] = _cast(value) outputs = tuple(outputs) return outputs outputs = _cast(outputs) return outputs def _tensor_inplace_assign(input, output, op_name, replace_op): # if pynative_mode_condition(): # TODO: ms_function # input.assign_value(output) # return input # raise RuntimeError('`Tensor.{a}` is an in-place operation and "x.{a}()" is not encouraged to use ' # 'in MindSpore static graph mode. Please use "x = x.{b}()" or other API ' # 'instead.'.format(a=op_name, b=replace_op)) # TODO: tensor api will be used in init data, but it can not be used in graph. unsupported_attr(op_name) unsupported_attr(replace_op) input.assign_value(output) return input def _div_calcu(input, other, rounding_mode): input = cast_to_ms_tensor(input) other = cast_to_ms_tensor(other) if rounding_mode is None: if input.dtype == mstype.int64 or input.dtype == mstype.int32: input = ms.ops.cast(input, mstype.float32) output = ms.ops.div(input, other) if rounding_mode == "trunc": output = ms.ops.div(input, other) if input.dtype == ms.int64: dtype_ = output.dtype output = ms.numpy.trunc(output, dtype=dtype_) else: output = ms.ops.trunc(output) if rounding_mode == "floor": input_dtype = input.dtype output = ms.ops.floor_div(input, other) output = ms.ops.cast(output, input_dtype) return output msadapter/ms_adapter/pytorch/utils/000077500000000000000000000000001436713554100200165ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/__init__.py000066400000000000000000000000521436713554100221240ustar00rootroot00000000000000from ms_adapter.pytorch.utils import data msadapter/ms_adapter/pytorch/utils/data/000077500000000000000000000000001436713554100207275ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/__init__.py000066400000000000000000000033641436713554100230460ustar00rootroot00000000000000from ms_adapter.pytorch.utils.data.sampler import ( BatchSampler, RandomSampler, Sampler, SequentialSampler, SubsetRandomSampler, WeightedRandomSampler, ) from ms_adapter.pytorch.utils.data.dataset import ( ChainDataset, ConcatDataset, Dataset, IterableDataset, Subset, TensorDataset, random_split, ) # from ms_adapter.pytorch.utils.data.datapipes.datapipe import ( # DFIterDataPipe, # DataChunk, # IterDataPipe, # MapDataPipe, # ) from ms_adapter.pytorch.utils.data.dataloader import ( DataLoader, _DatasetKind, get_worker_info, default_collate, default_convert, ) from ms_adapter.pytorch.utils.data.distributed import DistributedSampler # from ms_adapter.pytorch.utils.data import communication __all__ = ['BatchSampler', 'ChainDataset', 'ConcatDataset', # 'DFIterDataPipe', # 'DataChunk', 'DataLoader', # 'DataLoader2', 'Dataset', 'DistributedSampler', # 'IterDataPipe', 'IterableDataset', # 'MapDataPipe', 'RandomSampler', 'Sampler', 'SequentialSampler', 'Subset', 'SubsetRandomSampler', 'TensorDataset', 'WeightedRandomSampler', '_DatasetKind', # 'argument_validation', # 'communication', 'default_collate', 'default_convert', # 'functional_datapipe', 'get_worker_info', # 'guaranteed_datapipes_determinism', # 'non_deterministic', 'random_split', # 'runtime_validation', # 'runtime_validation_disabled' ] assert __all__ == sorted(__all__)msadapter/ms_adapter/pytorch/utils/data/_utils/000077500000000000000000000000001436713554100222265ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/_utils/__init__.py000066400000000000000000000034411436713554100243410ustar00rootroot00000000000000r"""Utility classes & functions for data loading. Code in this folder is mostly used by ../dataloder.py. A lot of multiprocessing is used in data loading, which only supports running functions defined in global environment (py2 can't serialize static methods). Therefore, for code tidiness we put these functions into different files in this folder. """ import sys import atexit # old private location of the ExceptionWrapper that some users rely on: from ms_adapter.pytorch._utils import ExceptionWrapper IS_WINDOWS = sys.platform == "win32" MP_STATUS_CHECK_INTERVAL = 5.0 r"""Interval (in seconds) to check status of processes to avoid hanging in multiprocessing data loading. This is mainly used in getting data from another process, in which case we need to periodically check whether the sender is alive to prevent hanging.""" python_exit_status = False DATAPIPE_SHARED_SEED = "_dl_shared_seed" r"""The key to share the same seed for shuffle DataPipe across distributed processes""" DATAPIPE_SHARED_SEED_COUNTER = "_dl_shared_seed_recv_cnt" r"""The key to count the number of distributed processes that have received the shared seed""" DATAPIPE_SHARED_SEED_DEFAULT_TIMEOUT = 30 * 60 r"""Timeout (in seconds) sending the shared seed from Rank 0 and sending the signal of the shared seed received from other Ranks. It uses the same default timeout for the distributed process group""" DATAPIPE_SHARED_SEED_CHECK_INTERVAL = 0.01 r"""Interval to check if each rank has received the shared seed""" try: import numpy HAS_NUMPY = True except ModuleNotFoundError: HAS_NUMPY = False def _set_python_exit_flag(): global python_exit_status python_exit_status = True atexit.register(_set_python_exit_flag) from . import worker, signal_handling, pin_memory, collate, fetch msadapter/ms_adapter/pytorch/utils/data/_utils/collate.py000066400000000000000000000202021436713554100242170ustar00rootroot00000000000000r""""Contains definitions of the methods used by the _BaseDataLoaderIter workers to collate samples fetched from dataset into Tensor(s). These **needs** to be in global scope since Py2 doesn't support serializing static methods. `default_collate` and `default_convert` are exposed to users via 'dataloader.py'. """ import mindspore as ms from ms_adapter.pytorch.tensor import Tensor from ms_adapter.pytorch.functional import stack import re import collections from ms_adapter.pytorch._six import string_classes from mindspore.common.api import _pynative_executor np_str_obj_array_pattern = re.compile(r'[SaUO]') def default_convert(data): r""" Function that converts each NumPy array element into a :class:`torch.Tensor`. If the input is a `Sequence`, `Collection`, or `Mapping`, it tries to convert each element inside to a :class:`torch.Tensor`. If the input is not an NumPy array, it is left unchanged. This is used as the default function for collation when both `batch_sampler` and `batch_size` are NOT defined in :class:`~torch.utils.data.DataLoader`. The general input type to output type mapping is similar to that of :func:`~torch.utils.data.default_collate`. See the description there for more details. Args: data: a single data point to be converted Examples: >>> # Example with `int` >>> default_convert(0) 0 >>> # Example with NumPy array >>> default_convert(np.array([0, 1])) tensor([0, 1]) >>> # Example with NamedTuple >>> Point = namedtuple('Point', ['x', 'y']) >>> default_convert(Point(0, 0)) Point(x=0, y=0) >>> default_convert(Point(np.array(0), np.array(0))) Point(x=tensor(0), y=tensor(0)) >>> # Example with List >>> default_convert([np.array([0, 1]), np.array([2, 3])]) [tensor([0, 1]), tensor([2, 3])] """ elem_type = type(data) if isinstance(data, ms.Tensor): return Tensor(data) elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \ and elem_type.__name__ != 'string_': # array of string classes and object if elem_type.__name__ == 'ndarray' \ and np_str_obj_array_pattern.search(data.dtype.str) is not None: return data return Tensor(data) elif isinstance(data, collections.abc.Mapping): try: return elem_type({key: default_convert(data[key]) for key in data}) except TypeError: # The mapping type may not support `__init__(iterable)`. return {key: default_convert(data[key]) for key in data} elif isinstance(data, tuple) and hasattr(data, '_fields'): # namedtuple return elem_type(*(default_convert(d) for d in data)) elif isinstance(data, tuple): return [default_convert(d) for d in data] # Backwards compatibility. elif isinstance(data, collections.abc.Sequence) and not isinstance(data, string_classes): try: return elem_type([default_convert(d) for d in data]) except TypeError: # The sequence type may not support `__init__(iterable)` (e.g., `range`). return [default_convert(d) for d in data] else: return data default_collate_err_msg_format = ( "default_collate: batch must contain tensors, numpy arrays, numbers, " "dicts or lists; found {}") def default_collate(batch): r""" Function that takes in a batch of data and puts the elements within the batch into a tensor with an additional outer dimension - batch size. The exact output type can be a :class:`torch.Tensor`, a `Sequence` of :class:`torch.Tensor`, a Collection of :class:`torch.Tensor`, or left unchanged, depending on the input type. This is used as the default function for collation when `batch_size` or `batch_sampler` is defined in :class:`~torch.utils.data.DataLoader`. Here is the general input type (based on the type of the element within the batch) to output type mapping: * :class:`torch.Tensor` -> :class:`torch.Tensor` (with an added outer dimension batch size) * NumPy Arrays -> :class:`torch.Tensor` * `float` -> :class:`torch.Tensor` * `int` -> :class:`torch.Tensor` * `str` -> `str` (unchanged) * `bytes` -> `bytes` (unchanged) * `Mapping[K, V_i]` -> `Mapping[K, default_collate([V_1, V_2, ...])]` * `NamedTuple[V1_i, V2_i, ...]` -> `NamedTuple[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]` * `Sequence[V1_i, V2_i, ...]` -> `Sequence[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]` Args: batch: a single batch to be collated Examples: >>> # Example with a batch of `int`s: >>> default_collate([0, 1, 2, 3]) tensor([0, 1, 2, 3]) >>> # Example with a batch of `str`s: >>> default_collate(['a', 'b', 'c']) ['a', 'b', 'c'] >>> # Example with `Map` inside the batch: >>> default_collate([{'A': 0, 'B': 1}, {'A': 100, 'B': 100}]) {'A': tensor([ 0, 100]), 'B': tensor([ 1, 100])} >>> # Example with `NamedTuple` inside the batch: >>> Point = namedtuple('Point', ['x', 'y']) >>> default_collate([Point(0, 0), Point(1, 1)]) Point(x=tensor([0, 1]), y=tensor([0, 1])) >>> # Example with `Tuple` inside the batch: >>> default_collate([(0, 1), (2, 3)]) [tensor([0, 2]), tensor([1, 3])] >>> # Example with `List` inside the batch: >>> default_collate([[0, 1], [2, 3]]) [tensor([0, 2]), tensor([1, 3])] """ elem = batch[0] elem_type = type(elem) if isinstance(elem, ms.Tensor): return stack(batch, 0) # return batch elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \ and elem_type.__name__ != 'string_': if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap': # array of string classes and object if np_str_obj_array_pattern.search(elem.dtype.str) is not None: raise TypeError(default_collate_err_msg_format.format(elem.dtype)) return default_collate([Tensor(b) for b in batch]) elif elem.shape == (): # scalars return Tensor(batch) elif isinstance(elem, float): return Tensor(batch, dtype=ms.float64) elif isinstance(elem, int): return Tensor(batch, inner=True) elif isinstance(elem, string_classes): return batch elif isinstance(elem, collections.abc.Mapping): try: return elem_type({key: default_collate([d[key] for d in batch]) for key in elem}) except TypeError: # The mapping type may not support `__init__(iterable)`. return {key: default_collate([d[key] for d in batch]) for key in elem} elif isinstance(elem, tuple) and hasattr(elem, '_fields'): # namedtuple return elem_type(*(default_collate(samples) for samples in zip(*batch))) elif isinstance(elem, collections.abc.Sequence): # check to make sure that the elements in batch have consistent size it = iter(batch) elem_size = len(next(it)) if not all(len(elem) == elem_size for elem in it): raise RuntimeError('each element in list of batch should be of equal size') transposed = list(zip(*batch)) # It may be accessed twice, so we use a list. if isinstance(elem, tuple): return [default_collate(samples) for samples in transposed] # Backwards compatibility. else: try: return elem_type([default_collate(samples) for samples in transposed]) except TypeError: # The sequence type may not support `__init__(iterable)` (e.g., `range`). return [default_collate(samples) for samples in transposed] raise TypeError(default_collate_err_msg_format.format(elem_type)) msadapter/ms_adapter/pytorch/utils/data/_utils/fetch.py000066400000000000000000000036061436713554100236760ustar00rootroot00000000000000r""""Contains definitions of the methods used by the _BaseDataLoaderIter to fetch data from an iterable-style or map-style dataset. This logic is shared in both single- and multi-processing data loading. """ class _BaseDatasetFetcher(object): def __init__(self, dataset, auto_collation, collate_fn, drop_last): self.dataset = dataset self.auto_collation = auto_collation self.collate_fn = collate_fn self.drop_last = drop_last def fetch(self, possibly_batched_index): raise NotImplementedError() class _IterableDatasetFetcher(_BaseDatasetFetcher): def __init__(self, dataset, auto_collation, collate_fn, drop_last): super(_IterableDatasetFetcher, self).__init__(dataset, auto_collation, collate_fn, drop_last) self.dataset_iter = iter(dataset) self.ended = False def fetch(self, possibly_batched_index): if self.ended: raise StopIteration if self.auto_collation: data = [] for _ in possibly_batched_index: try: data.append(next(self.dataset_iter)) except StopIteration: self.ended = True break if len(data) == 0 or (self.drop_last and len(data) < len(possibly_batched_index)): raise StopIteration else: data = next(self.dataset_iter) return self.collate_fn(data) class _MapDatasetFetcher(_BaseDatasetFetcher): def __init__(self, dataset, auto_collation, collate_fn, drop_last): super(_MapDatasetFetcher, self).__init__(dataset, auto_collation, collate_fn, drop_last) def fetch(self, possibly_batched_index): if self.auto_collation: data = [self.dataset[idx] for idx in possibly_batched_index] else: data = self.dataset[possibly_batched_index] return self.collate_fn(data) msadapter/ms_adapter/pytorch/utils/data/_utils/pin_memory.py000066400000000000000000000053261436713554100247640ustar00rootroot00000000000000r""""Contains definitions of the methods used by the _BaseDataLoaderIter to put fetched tensors into pinned memory. """ from __future__ import absolute_import import collections import queue import ms_adapter.pytorch as torch import mindspore as ms from ms_adapter.pytorch._six import string_classes from . import MP_STATUS_CHECK_INTERVAL from ms_adapter.pytorch._utils import ExceptionWrapper def _pin_memory_loop(in_queue, out_queue, device_id, done_event, device): # This setting is thread local, and prevents the copy in pin_memory from # consuming all CPU cores. # torch.set_num_threads(1) TODO: ms.set_context(device_id=device_id) # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the # logic of this function. while not done_event.is_set(): try: r = in_queue.get(timeout=MP_STATUS_CHECK_INTERVAL) except queue.Empty: continue idx, data = r if not done_event.is_set() and not isinstance(data, ExceptionWrapper): try: data = pin_memory(data, device) except Exception: data = ExceptionWrapper( where="in pin memory thread for device {}".format(device_id)) r = (idx, data) while not done_event.is_set(): try: out_queue.put(r, timeout=MP_STATUS_CHECK_INTERVAL) break except queue.Full: continue del r # save memory def pin_memory(data, device=None): if isinstance(data, torch.Tensor): return data.pin_memory(device) elif isinstance(data, string_classes): return data elif isinstance(data, collections.abc.Mapping): try: return type(data)({k: pin_memory(sample, device) for k, sample in data.items()}) # type: ignore[call-arg] except TypeError: # The mapping type may not support `__init__(iterable)`. return {k: pin_memory(sample, device) for k, sample in data.items()} elif isinstance(data, tuple) and hasattr(data, '_fields'): # namedtuple return type(data)(*(pin_memory(sample, device) for sample in data)) elif isinstance(data, tuple): return [pin_memory(sample, device) for sample in data] # Backwards compatibility. elif isinstance(data, collections.abc.Sequence): try: return type(data)([pin_memory(sample, device) for sample in data]) # type: ignore[call-arg] except TypeError: # The sequence type may not support `__init__(iterable)` (e.g., `range`). return [pin_memory(sample, device) for sample in data] elif hasattr(data, "pin_memory"): return data.pin_memory() else: return data msadapter/ms_adapter/pytorch/utils/data/_utils/serialization.py000066400000000000000000000006631436713554100254620ustar00rootroot00000000000000try: import dill # XXX: By default, dill writes the Pickler dispatch table to inject its # own logic there. This globally affects the behavior of the standard library # pickler for any user who transitively depends on this module! # Undo this extension to avoid altering the behavior of the pickler globally. dill.extend(use_dill=False) DILL_AVAILABLE = True except ImportError: DILL_AVAILABLE = False msadapter/ms_adapter/pytorch/utils/data/_utils/signal_handling.py000066400000000000000000000070211436713554100257210ustar00rootroot00000000000000r""""Signal handling for multiprocessing data loading. NOTE [ Signal handling in multiprocessing data loading ] In cases like DataLoader, if a worker process dies due to bus error/segfault or just hang, the main process will hang waiting for data. This is difficult to avoid on PyTorch side as it can be caused by limited shm, or other libraries users call in the workers. In this file and `DataLoader.cpp`, we make our best effort to provide some error message to users when such unfortunate events happen. When a _BaseDataLoaderIter starts worker processes, their pids are registered in a defined in `DataLoader.cpp`: id(_BaseDataLoaderIter) => Collection[ Worker pids ] via `_set_worker_pids`. When an error happens in a worker process, the main process received a SIGCHLD, and Python will eventually call the handler registered below (in `_set_SIGCHLD_handler`). In the handler, the `_error_if_any_worker_fails` call checks all registered worker pids and raise proper error message to prevent main process from hanging waiting for data from worker. Additionally, at the beginning of each worker's `_utils.worker._worker_loop`, `_set_worker_signal_handlers` is called to register critical signal handlers (e.g., for SIGSEGV, SIGBUS, SIGFPE, SIGTERM) in C, which just prints an error message to stderr before triggering the default handler. So a message will also be printed from the worker process when it is killed by such signals. See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for the reasoning of this signal handling design and other mechanism we implement to make our multiprocessing data loading robust to errors. """ import signal import threading from . import IS_WINDOWS # Some of the following imported functions are not used in this file, but are to # be used `_utils.signal_handling.XXXXX`. # TODO: # from torch._C import _set_worker_pids, _remove_worker_pids # noqa: F401 # from torch._C import _error_if_any_worker_fails, _set_worker_signal_handlers # noqa: F401 _SIGCHLD_handler_set = False r"""Whether SIGCHLD handler is set for DataLoader worker failures. Only one handler needs to be set for all DataLoaders in a process.""" def _set_worker_pids(): raise ValueError("'_set_worker_pids' is not currently supported.") def _remove_worker_pids(): raise ValueError("'_remove_worker_pids' is not currently supported.") def _error_if_any_worker_fails(): raise ValueError("'_error_if_any_worker_fails' is not currently supported.") def _set_worker_signal_handlers(): raise ValueError("'_set_worker_signal_handlers' is not currently supported.") def _set_SIGCHLD_handler(): # Windows doesn't support SIGCHLD handler if IS_WINDOWS: return # can't set signal in child threads if not isinstance(threading.current_thread(), threading._MainThread): # type: ignore[attr-defined] return global _SIGCHLD_handler_set if _SIGCHLD_handler_set: return previous_handler = signal.getsignal(signal.SIGCHLD) if not callable(previous_handler): # This doesn't catch default handler, but SIGCHLD default handler is a # no-op. previous_handler = None def handler(signum, frame): # This following call uses `waitid` with WNOHANG from C side. Therefore, # Python can still get and update the process status successfully. _error_if_any_worker_fails() if previous_handler is not None: assert callable(previous_handler) previous_handler(signum, frame) signal.signal(signal.SIGCHLD, handler) _SIGCHLD_handler_set = True msadapter/ms_adapter/pytorch/utils/data/_utils/worker.py000066400000000000000000000325031436713554100241140ustar00rootroot00000000000000r""""Contains definitions of the methods used by the _BaseDataLoaderIter workers. These **needs** to be in global scope since Py2 doesn't support serializing static methods. """ from __future__ import absolute_import import mindspore.context as ms_context import ms_adapter.pytorch as torch import random import os import queue from dataclasses import dataclass from ms_adapter.pytorch._utils import ExceptionWrapper from typing import Optional, Union from . import signal_handling, MP_STATUS_CHECK_INTERVAL, IS_WINDOWS, HAS_NUMPY if IS_WINDOWS: import ctypes from ctypes.wintypes import DWORD, BOOL, HANDLE # On Windows, the parent ID of the worker process remains unchanged when the manager process # is gone, and the only way to check it through OS is to let the worker have a process handle # of the manager and ask if the process status has changed. class ManagerWatchdog(object): def __init__(self): self.manager_pid = os.getppid() # mypy cannot detect this code is windows only self.kernel32 = ctypes.WinDLL('kernel32', use_last_error=True) # type: ignore[attr-defined] self.kernel32.OpenProcess.argtypes = (DWORD, BOOL, DWORD) self.kernel32.OpenProcess.restype = HANDLE self.kernel32.WaitForSingleObject.argtypes = (HANDLE, DWORD) self.kernel32.WaitForSingleObject.restype = DWORD # Value obtained from https://msdn.microsoft.com/en-us/library/ms684880.aspx SYNCHRONIZE = 0x00100000 self.manager_handle = self.kernel32.OpenProcess(SYNCHRONIZE, 0, self.manager_pid) if not self.manager_handle: raise ctypes.WinError(ctypes.get_last_error()) # type: ignore[attr-defined] self.manager_dead = False def is_alive(self): if not self.manager_dead: # Value obtained from https://msdn.microsoft.com/en-us/library/windows/desktop/ms687032.aspx self.manager_dead = self.kernel32.WaitForSingleObject(self.manager_handle, 0) == 0 return not self.manager_dead else: class ManagerWatchdog(object): # type: ignore[no-redef] def __init__(self): self.manager_pid = os.getppid() self.manager_dead = False def is_alive(self): if not self.manager_dead: self.manager_dead = os.getppid() != self.manager_pid return not self.manager_dead _worker_info = None class WorkerInfo(object): __initialized = False def __init__(self, **kwargs): for k, v in kwargs.items(): setattr(self, k, v) self.__keys = tuple(kwargs.keys()) self.__initialized = True def __setattr__(self, key, val): if self.__initialized: raise RuntimeError("Cannot assign attributes to {} objects".format(self.__class__.__name__)) return super(WorkerInfo, self).__setattr__(key, val) def __repr__(self): items = [] for k in self.__keys: items.append('{}={}'.format(k, getattr(self, k))) return '{}({})'.format(self.__class__.__name__, ', '.join(items)) def get_worker_info(): r"""Returns the information about the current :class:`~torch.utils.data.DataLoader` iterator worker process. When called in a worker, this returns an object guaranteed to have the following attributes: * :attr:`id`: the current worker id. * :attr:`num_workers`: the total number of workers. * :attr:`seed`: the random seed set for the current worker. This value is determined by main process RNG and the worker id. See :class:`~torch.utils.data.DataLoader`'s documentation for more details. * :attr:`dataset`: the copy of the dataset object in **this** process. Note that this will be a different object in a different process than the one in the main process. When called in the main process, this returns ``None``. .. note:: When used in a :attr:`worker_init_fn` passed over to :class:`~torch.utils.data.DataLoader`, this method can be useful to set up each worker process differently, for instance, using ``worker_id`` to configure the ``dataset`` object to only read a specific fraction of a sharded dataset, or use ``seed`` to seed other libraries used in dataset code. """ return _worker_info r"""Dummy class used to signal the end of an IterableDataset""" @dataclass(frozen=True) class _IterableDatasetStopIteration(object): worker_id: int r"""Dummy class used to resume the fetching when worker reuse is enabled""" @dataclass(frozen=True) class _ResumeIteration(object): seed: Optional[int] = None # The function `_generate_state` is adapted from `numpy.random.SeedSequence` # from https://github.com/numpy/numpy/blob/main/numpy/random/bit_generator.pyx # It's MIT licensed, here is the copyright: # Copyright (c) 2015 Melissa E. O'Neill # Copyright (c) 2019 NumPy Developers # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # This function generates an array of int32 as the seed for # `numpy.random`, in order to prevent state collision due to same # seed and algorithm for `numpy.random` and `random` modules. # TODO: Implement `SeedSequence` like object for `torch.random` def _generate_state(base_seed, worker_id): INIT_A = 0x43b0d7e5 MULT_A = 0x931e8875 INIT_B = 0x8b51f9dd MULT_B = 0x58f38ded MIX_MULT_L = 0xca01f9dd MIX_MULT_R = 0x4973f715 XSHIFT = 4 * 8 // 2 MASK32 = 0xFFFFFFFF entropy = [worker_id, base_seed & MASK32, base_seed >> 32, 0] pool = [0] * 4 hash_const_A = INIT_A def hash(value): nonlocal hash_const_A value = (value ^ hash_const_A) & MASK32 hash_const_A = (hash_const_A * MULT_A) & MASK32 value = (value * hash_const_A) & MASK32 value = (value ^ (value >> XSHIFT)) & MASK32 return value def mix(x, y): result_x = (MIX_MULT_L * x) & MASK32 result_y = (MIX_MULT_R * y) & MASK32 result = (result_x - result_y) & MASK32 result = (result ^ (result >> XSHIFT)) & MASK32 return result # Add in the entropy to the pool. for i in range(len(pool)): pool[i] = hash(entropy[i]) # Mix all bits together so late bits can affect earlier bits. for i_src in range(len(pool)): for i_dst in range(len(pool)): if i_src != i_dst: pool[i_dst] = mix(pool[i_dst], hash(pool[i_src])) hash_const_B = INIT_B state = [] for i_dst in range(4): data_val = pool[i_dst] data_val = (data_val ^ hash_const_B) & MASK32 hash_const_B = (hash_const_B * MULT_B) & MASK32 data_val = (data_val * hash_const_B) & MASK32 data_val = (data_val ^ (data_val >> XSHIFT)) & MASK32 state.append(data_val) return state def _worker_loop(dataset_kind, dataset, index_queue, data_queue, done_event, auto_collation, collate_fn, drop_last, base_seed, init_fn, worker_id, num_workers, persistent_workers, shared_seed): # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the # logic of this function. ms_context.set_context(device_target="CPU", mode=ms_context.PYNATIVE_MODE) try: # Initialize C side signal handlers for SIGBUS and SIGSEGV. Python signal # module's handlers are executed after Python returns from C low-level # handlers, likely when the same fatal signal had already happened # again. # https://docs.python.org/3/library/signal.html#execution-of-python-signal-handlers # signal_handling._set_worker_signal_handlers() # torch.set_num_threads(1) TODO: seed = base_seed + worker_id random.seed(seed) # torch.manual_seed(seed) TODO: if HAS_NUMPY: np_seed = _generate_state(base_seed, worker_id) import numpy as np np.random.seed(np_seed) # from ms_adapter.pytorch.utils.data import IterDataPipe # from ms_adapter.pytorch.utils.data.graph_settings import apply_shuffle_seed shared_rng = None # if isinstance(dataset, IterDataPipe): #TODO data pipe # shared_rng = torch.Generator() # assert shared_seed is not None # shared_rng.manual_seed(shared_seed) # dataset = apply_shuffle_seed(dataset, shared_rng) global _worker_info _worker_info = WorkerInfo(id=worker_id, num_workers=num_workers, seed=seed, dataset=dataset) from ms_adapter.pytorch.utils.data import _DatasetKind init_exception = None try: if init_fn is not None: init_fn(worker_id) fetcher = _DatasetKind.create_fetcher(dataset_kind, dataset, auto_collation, collate_fn, drop_last) except Exception: init_exception = ExceptionWrapper( where="in DataLoader worker process {}".format(worker_id)) # When using Iterable mode, some worker can exit earlier than others due # to the IterableDataset behaving differently for different workers. # When such things happen, an `_IterableDatasetStopIteration` object is # sent over to the main process with the ID of this worker, so that the # main process won't send more tasks to this worker, and will send # `None` to this worker to properly exit it. # # Note that we cannot set `done_event` from a worker as it is shared # among all processes. Instead, we set the `iteration_end` flag to # signify that the iterator is exhausted. When either `done_event` or # `iteration_end` is set, we skip all processing step and just wait for # `None`. iteration_end = False watchdog = ManagerWatchdog() while watchdog.is_alive(): try: r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL) except queue.Empty: continue if isinstance(r, _ResumeIteration): # Acknowledge the main process data_queue.put((r, None)) iteration_end = False # if isinstance(dataset, IterDataPipe): # assert r.seed is not None # shared_rng.manual_seed(r.seed) # dataset = apply_shuffle_seed(dataset, shared_rng) # Recreate the fetcher for worker-reuse policy fetcher = _DatasetKind.create_fetcher( dataset_kind, dataset, auto_collation, collate_fn, drop_last) continue elif r is None: # Received the final signal assert done_event.is_set() or iteration_end break elif done_event.is_set() or iteration_end: # `done_event` is set. But I haven't received the final signal # (None) yet. I will keep continuing until get it, and skip the # processing steps. continue idx, index = r data: Union[_IterableDatasetStopIteration, ExceptionWrapper] if init_exception is not None: data = init_exception init_exception = None else: try: data = fetcher.fetch(index) except Exception as e: if isinstance(e, StopIteration) and dataset_kind == _DatasetKind.Iterable: data = _IterableDatasetStopIteration(worker_id) # Set `iteration_end` # (1) to save future `next(...)` calls, and # (2) to avoid sending multiple `_IterableDatasetStopIteration`s. iteration_end = True else: # It is important that we don't store exc_info in a variable. # `ExceptionWrapper` does the correct thing. # See NOTE [ Python Traceback Reference Cycle Problem ] data = ExceptionWrapper( where="in DataLoader worker process {}".format(worker_id)) data_queue.put((idx, data)) del data, idx, index, r # save memory except KeyboardInterrupt: # Main process will raise KeyboardInterrupt anyways. pass if done_event.is_set(): data_queue.cancel_join_thread() data_queue.close() msadapter/ms_adapter/pytorch/utils/data/backward_compatibility.py000066400000000000000000000003241436713554100260070ustar00rootroot00000000000000import warnings def worker_init_fn(worker_id): warnings.warn("Usage of backward_compatibility.worker_init_fn is deprecated" " as DataLoader automatically applies sharding in every worker") msadapter/ms_adapter/pytorch/utils/data/communication/000077500000000000000000000000001436713554100235745ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/communication/__init__.py000066400000000000000000000001771436713554100257120ustar00rootroot00000000000000from . import eventloop from . import iter from . import map from . import messages from . import protocol from . import queue msadapter/ms_adapter/pytorch/utils/data/communication/eventloop.py000066400000000000000000000053521436713554100261660ustar00rootroot00000000000000from __future__ import absolute_import import ms_adapter.pytorch as torch import threading import pickle from ms_adapter.pytorch.utils.data import IterDataPipe, communication, MapDataPipe try: import dill # XXX: By default, dill writes the Pickler dispatch table to inject its # own logic there. This globally affects the behavior of the standard library # pickler for any user who transitively depends on this module! # Undo this extension to avoid altering the behavior of the pickler globally. dill.extend(use_dill=False) HAS_DILL = True except ImportError: HAS_DILL = False __all__ = [ "DataPipeToQueuesLoop", "SpawnProcessForDataPipeline", "SpawnThreadForDataPipeline", ] def DataPipeToQueuesLoop(source_datapipe, req_queue, res_queue): if isinstance(source_datapipe, IterDataPipe): pipe_type = communication.iter protocol_type = communication.protocol.IterDataPipeQueueProtocolServer elif isinstance(source_datapipe, MapDataPipe): pipe_type = communication.map # type: ignore[misc] protocol_type = communication.protocol.MapDataPipeQueueProtocolServer # type: ignore[assignment] else: raise Exception('Only supports IterDataPipe or MapDataPipe, got', source_datapipe) torch.set_num_threads(1) for _ in pipe_type.DataPipeBehindQueues(source_datapipe, protocol_type(req_queue, res_queue), blocking_request_get=True): pass def SpawnProcessForDataPipeline(multiprocessing_ctx, datapipe): req_queue = multiprocessing_ctx.Queue() res_queue = multiprocessing_ctx.Queue() process = multiprocessing_ctx.Process( target=DataPipeToQueuesLoop, args=(datapipe, req_queue, res_queue)) return process, req_queue, res_queue def SpawnThreadForDataPipeline(datapipe): r""" Given a DataPipe, creates a copy of the DataPipe, starts a new Thread with DataPipeToQueuesLoop as target, and return the process, req_queue, res_queue, thread_local_datapipe. """ req_queue = communication.queue.ThreadingQueue() res_queue = communication.queue.ThreadingQueue() try: new_datapipe = pickle.loads(pickle.dumps(datapipe)) except Exception as pe: if HAS_DILL: try: new_datapipe = dill.loads(dill.dumps(datapipe)) except Exception as de: Exception('Unable to dill DataPipe to make thread local copy', de) else: raise Exception('Unable to pickle DataPipe to make thread local copy (consider installing `dill`)', pe) process = threading.Thread(target=DataPipeToQueuesLoop, args=( new_datapipe, req_queue, res_queue), daemon=True) return process, req_queue, res_queue, new_datapipe msadapter/ms_adapter/pytorch/utils/data/communication/iter.py000066400000000000000000000146621436713554100251220ustar00rootroot00000000000000from __future__ import absolute_import import time import types from ms_adapter.pytorch.utils.data import IterDataPipe, communication DEFAULT_NON_BLOCKING_SLEEP = 0.001 __all__ = [ "DataPipeBehindQueues", "EnsureNonBlockingDataPipe", "InvalidStateResetRequired", "NonBlocking", "NotAvailable", "QueueWrapper", "default_not_available_hook", ] def default_not_available_hook(): time.sleep(DEFAULT_NON_BLOCKING_SLEEP) class NotAvailable(Exception): pass class InvalidStateResetRequired(Exception): """ Returned by DataPipe when it is expecting to get reset request, for example RouterDataPipe expecting all workers to request reset' """ pass class NonBlocking(IterDataPipe): not_available_hook = default_not_available_hook def __iter__(self): self.reset_iterator() return self def __next__(self): while True: try: return self.nonblocking_next() except StopIteration: raise StopIteration except NotAvailable: if NonBlocking.not_available_hook is not None: NonBlocking.not_available_hook() def nonblocking_next(self): raise NotImplementedError( "nonblocking_next is not implemented for %s" % self.__class__) def reset_iterator(self): raise NotImplementedError( "reset_iterator is not implemented for %s" % self.__class__) @staticmethod def register_not_available_hook(hook_function): NonBlocking.not_available_hook = hook_function def EnsureNonBlockingDataPipe(validated_datapipe): if not isinstance(validated_datapipe, IterDataPipe): raise Exception('Not Iterable DataPipe ' + str(validated_datapipe.__class__)) if isinstance(validated_datapipe, NonBlocking): return validated_datapipe if not hasattr(validated_datapipe, '_as_iterator'): validated_datapipe._as_iterator = None # type: ignore[attr-defined] if not hasattr(validated_datapipe, 'nonblocking_next'): def nonblocking_next(self): if self._as_iterator is None: self._as_iterator = iter(self) return next(self._as_iterator) validated_datapipe.nonblocking_next = types.MethodType( # type: ignore[attr-defined] nonblocking_next, validated_datapipe) if not hasattr(validated_datapipe, 'reset_iterator'): def reset_iterator(self): self._as_iterator = None validated_datapipe.reset_iterator = types.MethodType( # type: ignore[attr-defined] reset_iterator, validated_datapipe) return validated_datapipe def DataPipeBehindQueues(source_datapipe, protocol, full_stop=False, blocking_request_get=False): """ Indefinitely iterates over req_queue and passing values from source_datapipe to res_queue If raise_stop is true, raises exception when StopIteration received from the source_datapipe """ if not isinstance(protocol, communication.protocol.IterDataPipeQueueProtocolServer): raise Exception('Expecting IterDataPipeQueueProtocolServer, got', protocol) source_datapipe = EnsureNonBlockingDataPipe(source_datapipe) forever = True while forever: try: # Non-blocking call is Extremely slow here for python.mp, need to figure out a good workaround request = protocol.get_new_request(block=blocking_request_get) except communication.protocol.EmptyQueue: yield True continue if isinstance(request, communication.messages.ResetIteratorRequest): source_datapipe.reset_iterator() protocol.response_reset_iterator() elif isinstance(request, communication.messages.TerminateRequest): forever = False protocol.response_terminate() elif isinstance(request, communication.messages.GetNextRequest): while forever: try: value = source_datapipe.nonblocking_next() except NotAvailable: yield True continue except StopIteration: protocol.response_stop_iteration() if full_stop: forever = False else: yield True break except InvalidStateResetRequired: protocol.response_invalid_state() if full_stop: forever = False else: yield True break protocol.response_next(value) yield True # Returns control break else: raise Exception('Unrecognized type of request received', request) class QueueWrapper(NonBlocking): """ Creates iter.DataPipe which reads data from the DataLoader.Queue """ def __init__(self, protocol, response_wait_time=0.00001): if not isinstance(protocol, communication.protocol.IterDataPipeQueueProtocolClient): raise Exception('Got', protocol) self.protocol = protocol self.counter = 0 self._stop_iteration = False self._response_wait_time = response_wait_time def reset_iterator(self): self._stop_iteration = False self.counter = 0 self.protocol.request_reset_iterator() while True: try: self.protocol.get_response_reset_iterator() break except communication.protocol.EmptyQueue: if NonBlocking.not_available_hook is not None: NonBlocking.not_available_hook() def nonblocking_next(self): if self._stop_iteration: raise Exception( '`next` or `nonblocking_next` called after receiving StopIteration') if self.protocol.can_take_request(): self.protocol.request_next() try: response = self.protocol.get_response_next(block=True, timeout=self._response_wait_time) except communication.protocol.EmptyQueue: raise NotAvailable if isinstance(response, communication.messages.StopIterationResponse): self._stop_iteration = True raise StopIteration if isinstance(response, communication.messages.InvalidStateResponse): raise NotAvailable return response.value msadapter/ms_adapter/pytorch/utils/data/communication/map.py000066400000000000000000000137431436713554100247330ustar00rootroot00000000000000from __future__ import absolute_import import time import types from ms_adapter.pytorch.utils.data import communication, MapDataPipe DEFAULT_NON_BLOCKING_SLEEP = 0.001 __all__ = [ "DataPipeBehindQueues", "EnsureNonBlockingMapDataPipe", "NonBlockingMap", "NotAvailable", "QueueWrapperForMap", "default_not_available_hook", ] def default_not_available_hook(): time.sleep(DEFAULT_NON_BLOCKING_SLEEP) class NotAvailable(Exception): pass class NonBlockingMap(MapDataPipe): not_available_hook = default_not_available_hook def __getitem__(self, index): while True: try: return self.nonblocking_getitem(index) except NotAvailable: if NonBlockingMap.not_available_hook is not None: NonBlockingMap.not_available_hook() def __len__(self): try: return self.nonblocking_len() except NotAvailable: if NonBlockingMap.not_available_hook is not None: NonBlockingMap.not_available_hook() def nonblocking_len(self): raise NotImplementedError( "nonblocking_len is not implemented for %s" % self.__class__) def nonblocking_getitem(self, index): raise NotImplementedError( "nonblocking_getitem is not implemented for %s" % self.__class__) @staticmethod def register_not_available_hook(hook_function): NonBlockingMap.not_available_hook = hook_function def EnsureNonBlockingMapDataPipe(validated_datapipe): if not isinstance(validated_datapipe, MapDataPipe): raise Exception(f'Not Map DataPipe - got {validated_datapipe.__class__}') if isinstance(validated_datapipe, NonBlockingMap): return validated_datapipe if not hasattr(validated_datapipe, 'nonblocking_len'): def nonblocking_len(self): return self.__len__() validated_datapipe.nonblocking_len = types.MethodType( # type: ignore[attr-defined] nonblocking_len, validated_datapipe) if not hasattr(validated_datapipe, 'nonblocking_getitem'): def nonblocking_getitem(self, index): return self.__getitem__(index) validated_datapipe.nonblocking_getitem = types.MethodType( # type: ignore[attr-defined] nonblocking_getitem, validated_datapipe) return validated_datapipe def DataPipeBehindQueues(source_datapipe, protocol, full_stop=False, blocking_request_get=False): """ Indefinitely iterates over req_queue and passing values from source_datapipe to res_queue If raise_stop is true, raises exception when StopIteration received from the source_datapipe """ if not isinstance(protocol, communication.protocol.MapDataPipeQueueProtocolServer): raise Exception('Expecting MapDataPipeQueueProtocolServer, got', protocol) source_datapipe = EnsureNonBlockingMapDataPipe(source_datapipe) forever = True while forever: try: # Non-blocking call is Extremely slow here for python.mp, need to figure out a good workaround request = protocol.get_new_request(block=blocking_request_get) except communication.protocol.EmptyQueue: yield True continue if isinstance(request, communication.messages.TerminateRequest): forever = False protocol.response_terminate() elif isinstance(request, communication.messages.LenRequest): size = source_datapipe.nonblocking_len() protocol.response_len(size) elif isinstance(request, communication.messages.GetItemRequest): while forever: try: value = source_datapipe.nonblocking_getitem(request.key) except NotAvailable: yield True continue except IndexError as e: # Alternatively, we can just allow the underlying DataPipe to throw an exception? protocol.response_index_out_of_bound() if full_stop: forever = False else: yield True break protocol.response_item(request.key, value) yield True # Returns control break else: raise Exception('Unrecognized type of request received', request) class QueueWrapperForMap(NonBlockingMap): """ Creates map.DataPipe which reads data from the DataLoader.Queue """ def __init__(self, protocol, response_wait_time=0.00001): if not isinstance(protocol, communication.protocol.MapDataPipeQueueProtocolClient): raise Exception('Got', protocol) self.protocol = protocol self.counter = 0 self._stop_iteration = False self._response_wait_time = response_wait_time def nonblocking_getitem(self, index): if self._stop_iteration: raise Exception( '`getitem` or `nonblocking_getitem` called after receiving StopIteration') if self.protocol.can_take_request(): self.protocol.request_item(index) try: response = self.protocol.get_response_item(block=True, timeout=self._response_wait_time) except communication.protocol.EmptyQueue: raise NotAvailable if isinstance(response, communication.messages.StopIterationResponse): self._stop_iteration = True raise IndexError(f"Index {index} is out of bound.") return response.key, response.value def nonblocking_len(self): if self._stop_iteration: raise Exception( '`len` or `nonblocking_len` called after receiving StopIteration') if self.protocol.can_take_request(): self.protocol.request_len() try: response = self.protocol.get_response_len(block=True, timeout=self._response_wait_time) except communication.protocol.EmptyQueue: raise NotAvailable return response.len msadapter/ms_adapter/pytorch/utils/data/communication/messages.py000066400000000000000000000022231436713554100257540ustar00rootroot00000000000000class DataLoaderQueueMessage(object): pass class Request(DataLoaderQueueMessage): pass class Response(DataLoaderQueueMessage): pass class ResetIteratorRequest(Request): pass class ResetIteratorResponse(Response): pass class TerminateRequest(Request): pass class TerminateResponse(Response): pass class LenRequest(Request): pass class LenResponse(Response): __slots__ = ('len') def __init__(self, len): self.len = len class GetItemRequest(Request): __slots__ = ('key') def __init__(self, key): self.key = key class GetItemResponse(Response): __slots__ = ('key', 'value') def __init__(self, key, value): self.key = key self.value = value class GetNextRequest(Request): pass class GetNextResponse(Response): __slots__ = ('value') def __init__(self, value): self.value = value class StopIterationResponse(Response): pass class InvalidStateResponse(Response): """ Returned by DataPipe when it is expecting to get reset request, for example RouterDataPipe expecting all workers to request reset' """ pass msadapter/ms_adapter/pytorch/utils/data/communication/protocol.py000066400000000000000000000201671436713554100260150ustar00rootroot00000000000000from __future__ import absolute_import from ms_adapter.pytorch.utils.data import communication class Protocol(object): __slots__ = ('request_queue', 'response_queue') def __init__(self, request_queue, response_queue): self.request_queue = request_queue self.response_queue = response_queue class ProtocolClient(Protocol): """ ProtocolClient takes charge of putting requests into req_queue and returning results from res_queue. """ _req_sent = None def __init__(self, request_queue, response_queue): self.request_queue = request_queue self.response_queue = response_queue self._req_sent = None def can_take_request(self): return self._req_sent is None def waiting_for_response(self): return self._req_sent is not None def request_sent(self, request=True): if not self.can_take_request(): raise Exception('Protocol only supports one request in the Queue') self._req_sent = request def request_served(self, result=None): if not self.waiting_for_response(): raise Exception( 'Expected no peding requests, but something got served', result) self._req_sent = None class ProtocolServer(Protocol): """ ProtocolServer takes charge of getting requests from req_queue and fetching data from source datapipe. """ _req_received = None def __init__(self, request_queue, response_queue): self.request_queue = request_queue self.response_queue = response_queue self._req_received = None def have_pending_request(self): return self._req_received is not None def get_new_request(self, block=False): if self.have_pending_request(): raise Exception( 'Trying to get next request, while having one unserved') try: response = self.request_queue.get(block=block) except Exception as e: # TODO: Catch only timeout exceptions raise EmptyQueue('queue is empty') self._req_received = response return response # TODO: Validate supported requests def response_terminate(self): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") if not isinstance(self._req_received, communication.messages.TerminateRequest): raise Exception( "Replaying with terminate status to other type of message") self.response_queue.put(communication.messages.TerminateResponse()) self._req_received = None class MapDataPipeQueueProtocolServer(ProtocolServer): def response_item(self, key, value): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") self.response_queue.put(communication.messages.GetItemResponse(key, value)) self._req_received = None def response_len(self, size): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") self.response_queue.put(communication.messages.LenResponse(size)) self._req_received = None def response_index_out_of_bound(self): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") self.response_queue.put(communication.messages.StopIterationResponse()) self._req_received = None class MapDataPipeQueueProtocolClient(ProtocolClient): def request_len(self): if not self.can_take_request(): raise Exception('Can not request len while we are still waiting response for previous request') request = communication.messages.LenRequest() self.request_queue.put(request) self.request_sent(request) def request_item(self, index): if not self.can_take_request(): raise Exception('Can not request item while we are still waiting response for previous request') request = communication.messages.GetItemRequest(index) self.request_queue.put(request) self.request_sent(request) def get_response_len(self, block=False, timeout=None): if not self.waiting_for_response(): raise Exception('Can not expect any response without submitted request') try: response = self.response_queue.get(block=block, timeout=timeout) except TimeoutError: raise EmptyQueue('queue is empty') self.request_served(response) if not isinstance(response, communication.messages.LenResponse): raise Exception('Invalid response received') return response def get_response_item(self, block=False, timeout=None): if not self.waiting_for_response(): raise Exception('Can not expect any response without submitted request') try: response = self.response_queue.get(block=block, timeout=timeout) except TimeoutError: raise EmptyQueue('queue is empty') self.request_served(response) # if not isinstance(response, communication.messages.GetItemResponse): # raise Exception('Invalid response received') return response class EmptyQueue(Exception): pass class IterDataPipeQueueProtocolServer(ProtocolServer): def response_reset_iterator(self): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") if not isinstance(self._req_received, communication.messages.ResetIteratorRequest): raise Exception( "Replaying with reset status to other type of message") self.response_queue.put(communication.messages.ResetIteratorResponse()) self._req_received = None def response_next(self, value): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") self.response_queue.put(communication.messages.GetNextResponse(value)) self._req_received = None def response_stop_iteration(self): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") self.response_queue.put(communication.messages.StopIterationResponse()) self._req_received = None def response_invalid_state(self): if not self.have_pending_request(): raise Exception("Attempting to reply with pending request") self.response_queue.put(communication.messages.InvalidStateResponse()) self._req_received = None class IterDataPipeQueueProtocolClient(ProtocolClient): def request_reset_iterator(self): if not self.can_take_request(): raise Exception('Can not reset while we are still waiting response for previous request') request = communication.messages.ResetIteratorRequest() self.request_queue.put(request) self.request_sent(request) def request_next(self): if not self.can_take_request(): raise Exception('Can not request next item while we are still waiting response for previous request') request = communication.messages.GetNextRequest() self.request_queue.put(request) self.request_sent(request) def get_response_reset_iterator(self, block=False): try: response = self.response_queue.get(block=block) except Exception as e: # TODO: Catch only timeout exceptions raise EmptyQueue('queue is empty') self.request_served(response) if not isinstance(response, communication.messages.ResetIteratorResponse): raise Exception('Invalid response received') def get_response_next(self, block=False, timeout=None): if not self.waiting_for_response(): raise Exception( 'Can not expect any response without submitted request') try: response = self.response_queue.get(block=block, timeout=timeout) except Exception as e: # TODO: Catch only timeout exceptions raise EmptyQueue('queue is empty') self.request_served(response) # TODO(VitalyFedyunin): Add possible response types validation here return response msadapter/ms_adapter/pytorch/utils/data/communication/queue.py000066400000000000000000000025711436713554100252770ustar00rootroot00000000000000import threading import time class LocalQueue(): ops = 0 stored = 0 uid = 0 empty = 0 def __init__(self, name='unnamed'): self.items = [] self.name = name self.uid = LocalQueue.uid LocalQueue.uid += 1 def put(self, item, block=True): LocalQueue.ops += 1 LocalQueue.stored += 1 self.items.append(item) def get(self, block=True, timeout=0): # TODO(VitalyFedyunin): Add support of block and timeout arguments LocalQueue.ops += 1 if not len(self.items): LocalQueue.empty += 1 raise Exception('LocalQueue is empty') LocalQueue.stored -= 1 return self.items.pop() class ThreadingQueue(): def __init__(self, name='unnamed'): self.lock = threading.Lock() self.items = [] self.name = name def put(self, item, block=True): with self.lock: self.items.append(item) def get(self, block=True, timeout=0): # TODO(VitalyFedyunin): Add support of block and timeout arguments while True: with self.lock: if len(self.items) > 0: return self.items.pop() if not block: raise Exception("Not available") # TODO(VitalyFedyunin): Figure out what to do if nothing in the queue time.sleep(0.000001) msadapter/ms_adapter/pytorch/utils/data/dataloader.py000066400000000000000000001461671436713554100234200ustar00rootroot00000000000000r"""Definition of the DataLoader and associated iterators that subclass _BaseDataLoaderIter To support these two classes, in `./_utils` we define many utility methods and functions to be run in multiprocessing. E.g., the data loading worker loop is in `./_utils/worker.py`. """ from __future__ import absolute_import import functools import itertools import logging import os import queue import threading import time import warnings import multiprocessing from datetime import timedelta from typing import Any, Callable, Iterable, TypeVar, Generic, Sequence, List, Optional, Union import numpy as np import ms_adapter.pytorch as torch # import ms_adapter.pytorch.utils.data.graph_settings from ms_adapter.pytorch._utils import ExceptionWrapper from ms_adapter.pytorch._six import string_classes # from mindspore import communication as dist from . import ( # IterDataPipe, # MapDataPipe, IterableDataset, Sampler, SequentialSampler, RandomSampler, BatchSampler, Dataset, ) # from ms_adapter.pytorch.utils.data.datapipes.datapipe import _IterDataPipeSerializationWrapper, _MapDataPipeSerializationWrapper from . import _utils __all__ = [ "DataLoader", "get_worker_info", "default_collate", "default_convert", ] multiprocessing.set_start_method('spawn', force=True) T_co = TypeVar('T_co', covariant=True) T = TypeVar('T') _worker_init_fn_t = Callable[[int], None] # Ideally we would parameterize `DataLoader` by the return type of `collate_fn`, but there is currently no way to have that # type parameter set to a default value if the user doesn't pass in a custom 'collate_fn'. # See https://github.com/python/mypy/issues/3737. _collate_fn_t = Callable[[List[T]], Any] default_collate: _collate_fn_t = _utils.collate.default_collate default_convert = _utils.collate.default_convert get_worker_info = _utils.worker.get_worker_info logger = logging.getLogger(__name__) class _DatasetKind(object): Map = 0 Iterable = 1 @staticmethod def create_fetcher(kind, dataset, auto_collation, collate_fn, drop_last): if kind == _DatasetKind.Map: return _utils.fetch._MapDatasetFetcher(dataset, auto_collation, collate_fn, drop_last) else: return _utils.fetch._IterableDatasetFetcher(dataset, auto_collation, collate_fn, drop_last) class _InfiniteConstantSampler(Sampler): r"""Analogous to ``itertools.repeat(None, None)``. Args: data_source (Dataset): dataset to sample from """ def __init__(self): super(_InfiniteConstantSampler, self).__init__(None) def __iter__(self): while True: yield None # def _get_distributed_settings(): # if dist.is_available and dist.is_initialized(): # return dist.get_world_size(), dist.get_rank() #TODO # else: # return 1, 0 # def _sharding_worker_init_fn(worker_init_fn, world_size, rank_id, worker_id): # global_worker_id = worker_id # info = torch.utils.data.get_worker_info() # total_workers = info.num_workers # datapipe = info.dataset # # To distribute elements across distributed process evenly, we should shard data on distributed # # processes first then shard on worker processes # total_workers *= world_size # global_worker_id = global_worker_id * world_size + rank_id # ms_adapter.pytorch.utils.data.graph_settings.apply_sharding(datapipe, total_workers, global_worker_id) # if worker_init_fn is not None: # worker_init_fn(worker_id) class DataLoader(Generic[T_co]): r""" Data loader. Combines a dataset and a sampler, and provides an iterable over the given dataset. The :class:`~torch.utils.data.DataLoader` supports both map-style and iterable-style datasets with single- or multi-process loading, customizing loading order and optional automatic batching (collation) and memory pinning. See :py:mod:`torch.utils.data` documentation page for more details. Args: dataset (Dataset): dataset from which to load the data. batch_size (int, optional): how many samples per batch to load (default: ``1``). shuffle (bool, optional): set to ``True`` to have the data reshuffled at every epoch (default: ``False``). sampler (Sampler or Iterable, optional): defines the strategy to draw samples from the dataset. Can be any ``Iterable`` with ``__len__`` implemented. If specified, :attr:`shuffle` must not be specified. batch_sampler (Sampler or Iterable, optional): like :attr:`sampler`, but returns a batch of indices at a time. Mutually exclusive with :attr:`batch_size`, :attr:`shuffle`, :attr:`sampler`, and :attr:`drop_last`. num_workers (int, optional): how many subprocesses to use for data loading. ``0`` means that the data will be loaded in the main process. (default: ``0``) collate_fn (callable, optional): merges a list of samples to form a mini-batch of Tensor(s). Used when using batched loading from a map-style dataset. pin_memory (bool, optional): If ``True``, the data loader will copy Tensors into device/CUDA pinned memory before returning them. If your data elements are a custom type, or your :attr:`collate_fn` returns a batch that is a custom type, see the example below. drop_last (bool, optional): set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then the last batch will be smaller. (default: ``False``) timeout (numeric, optional): if positive, the timeout value for collecting a batch from workers. Should always be non-negative. (default: ``0``) worker_init_fn (callable, optional): If not ``None``, this will be called on each worker subprocess with the worker id (an int in ``[0, num_workers - 1]``) as input, after seeding and before data loading. (default: ``None``) generator (torch.Generator, optional): If not ``None``, this RNG will be used by RandomSampler to generate random indexes and multiprocessing to generate `base_seed` for workers. (default: ``None``) prefetch_factor (int, optional, keyword-only arg): Number of batches loaded in advance by each worker. ``2`` means there will be a total of 2 * num_workers batches prefetched across all workers. (default: ``2``) persistent_workers (bool, optional): If ``True``, the data loader will not shutdown the worker processes after a dataset has been consumed once. This allows to maintain the workers `Dataset` instances alive. (default: ``False``) pin_memory_device (str, optional): the data loader will copy Tensors into device pinned memory before returning them if pin_memory is set to true. .. warning:: If the ``spawn`` start method is used, :attr:`worker_init_fn` cannot be an unpicklable object, e.g., a lambda function. See :ref:`multiprocessing-best-practices` on more details related to multiprocessing in PyTorch. .. warning:: ``len(dataloader)`` heuristic is based on the length of the sampler used. When :attr:`dataset` is an :class:`~torch.utils.data.IterableDataset`, it instead returns an estimate based on ``len(dataset) / batch_size``, with proper rounding depending on :attr:`drop_last`, regardless of multi-process loading configurations. This represents the best guess PyTorch can make because PyTorch trusts user :attr:`dataset` code in correctly handling multi-process loading to avoid duplicate data. However, if sharding results in multiple workers having incomplete last batches, this estimate can still be inaccurate, because (1) an otherwise complete batch can be broken into multiple ones and (2) more than one batch worth of samples can be dropped when :attr:`drop_last` is set. Unfortunately, PyTorch can not detect such cases in general. See `Dataset Types`_ for more details on these two types of datasets and how :class:`~torch.utils.data.IterableDataset` interacts with `Multi-process data loading`_. .. warning:: See :ref:`reproducibility`, and :ref:`dataloader-workers-random-seed`, and :ref:`data-loading-randomness` notes for random seed related questions. """ dataset: Dataset[T_co] batch_size: Optional[int] num_workers: int pin_memory: bool drop_last: bool timeout: float sampler: Union[Sampler, Iterable] pin_memory_device: str prefetch_factor: int _iterator : Optional['_BaseDataLoaderIter'] __initialized = False def __init__(self, dataset: Dataset[T_co], batch_size: Optional[int] = 1, shuffle: Optional[bool] = None, sampler: Union[Sampler, Iterable, None] = None, batch_sampler: Union[Sampler[Sequence], Iterable[Sequence], None] = None, num_workers: int = 0, collate_fn: Optional[_collate_fn_t] = None, pin_memory: bool = False, drop_last: bool = False, timeout: float = 0, worker_init_fn: Optional[_worker_init_fn_t] = None, multiprocessing_context=None, generator=None, *, prefetch_factor: int = 2, persistent_workers: bool = False, pin_memory_device: str = ""): # torch._C._log_api_usage_once("python.data_loader") if num_workers < 0: raise ValueError('num_workers option should be non-negative; ' 'use num_workers=0 to disable multiprocessing.') if timeout < 0: raise ValueError('timeout option should be non-negative') if num_workers == 0 and prefetch_factor != 2: raise ValueError('prefetch_factor option could only be specified in multiprocessing.' 'let num_workers > 0 to enable multiprocessing.') assert prefetch_factor > 0 if persistent_workers and num_workers == 0: raise ValueError('persistent_workers option needs num_workers > 0') self.dataset = dataset self.num_workers = num_workers self.prefetch_factor = prefetch_factor self.pin_memory = pin_memory self.pin_memory_device = pin_memory_device self.timeout = timeout self.worker_init_fn = worker_init_fn self.multiprocessing_context = multiprocessing_context if self.pin_memory: warnings.warn("Don't support pin_memory now.") if len(pin_memory_device) != 0: warnings.warn("Don't support set pin_memory_device now.") # Adds several forward compatibilities so classic DataLoader can work with DataPipes # 1. _DataPipeSerializationWrapper container makes it easier to serialize without redefining pickler # 2. Additional worker init function will take care of sharding in MP and Distributed # if isinstance(self.dataset, IterDataPipe): # self.dataset = _IterDataPipeSerializationWrapper(self.dataset) # ws, rank = _get_distributed_settings() # if num_workers > 0: # self.worker_init_fn = functools.partial( # _sharding_worker_init_fn, self.worker_init_fn, ws, rank) # else: # ms_adapter.pytorch.utils.data.graph_settings.apply_sharding(self.dataset, ws, rank) # elif isinstance(self.dataset, MapDataPipe): # self.dataset = _MapDataPipeSerializationWrapper(self.dataset) # ws, rank = _get_distributed_settings() # if num_workers > 0: # self.worker_init_fn = functools.partial( # _sharding_worker_init_fn, self.worker_init_fn, ws, rank) # else: # ms_adapter.pytorch.utils.data.graph_settings.apply_sharding(self.dataset, ws, rank) # Arg-check dataset related before checking samplers because we want to # tell users that iterable-style datasets are incompatible with custom # samplers first, so that they don't learn that this combo doesn't work # after spending time fixing the custom sampler errors. if isinstance(dataset, IterableDataset): self._dataset_kind = _DatasetKind.Iterable # NOTE [ Custom Samplers and IterableDataset ] # # `IterableDataset` does not support custom `batch_sampler` or # `sampler` since the key is irrelevant (unless we support # generator-style dataset one day...). # # For `sampler`, we always create a dummy sampler. This is an # infinite sampler even when the dataset may have an implemented # finite `__len__` because in multi-process data loading, naive # settings will return duplicated data (which may be desired), and # thus using a sampler with length matching that of dataset will # cause data lost (you may have duplicates of the first couple # batches, but never see anything afterwards). Therefore, # `Iterabledataset` always uses an infinite sampler, an instance of # `_InfiniteConstantSampler` defined above. # # A custom `batch_sampler` essentially only controls the batch size. # However, it is unclear how useful it would be since an iterable-style # dataset can handle that within itself. Moreover, it is pointless # in multi-process data loading as the assignment order of batches # to workers is an implementation detail so users can not control # how to batchify each worker's iterable. Thus, we disable this # option. If this turns out to be useful in future, we can re-enable # this, and support custom samplers that specify the assignments to # specific workers. # if isinstance(dataset, IterDataPipe): # if shuffle is not None: # dataset = ms_adapter.pytorch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle) # We cannot check `shuffle is not None` here, since previously `shuffle=False` was the default. if shuffle not in {False, None}: raise ValueError( "DataLoader with IterableDataset: expected unspecified " "shuffle option, but got shuffle={}".format(shuffle)) if sampler is not None: # See NOTE [ Custom Samplers and IterableDataset ] raise ValueError( "DataLoader with IterableDataset: expected unspecified " "sampler option, but got sampler={}".format(sampler)) elif batch_sampler is not None: # See NOTE [ Custom Samplers and IterableDataset ] raise ValueError( "DataLoader with IterableDataset: expected unspecified " "batch_sampler option, but got batch_sampler={}".format(batch_sampler)) else: shuffle = bool(shuffle) self._dataset_kind = _DatasetKind.Map if sampler is not None and shuffle: raise ValueError('sampler option is mutually exclusive with ' 'shuffle') if batch_sampler is not None: # auto_collation with custom batch_sampler if batch_size != 1 or shuffle or sampler is not None or drop_last: raise ValueError('batch_sampler option is mutually exclusive ' 'with batch_size, shuffle, sampler, and ' 'drop_last') batch_size = None drop_last = False elif batch_size is None: # no auto_collation if drop_last: raise ValueError('batch_size=None option disables auto-batching ' 'and is mutually exclusive with drop_last') if sampler is None: # give default samplers if self._dataset_kind == _DatasetKind.Iterable: # See NOTE [ Custom Samplers and IterableDataset ] sampler = _InfiniteConstantSampler() else: # map-style if shuffle: sampler = RandomSampler(dataset, generator=generator) # type: ignore[arg-type] else: sampler = SequentialSampler(dataset) # type: ignore[arg-type] if batch_size is not None and batch_sampler is None: # auto_collation without custom batch_sampler batch_sampler = BatchSampler(sampler, batch_size, drop_last) self.batch_size = batch_size self.drop_last = drop_last self.sampler = sampler self.batch_sampler = batch_sampler self.generator = generator if collate_fn is None: if self._auto_collation: collate_fn = _utils.collate.default_collate else: collate_fn = _utils.collate.default_convert self.collate_fn = collate_fn self.persistent_workers = persistent_workers self.__initialized = True self._IterableDataset_len_called = None # See NOTE [ IterableDataset and __len__ ] self._iterator = None self.check_worker_number_rationality() # torch.set_vital('Dataloader', 'enabled', 'True') # type: ignore[attr-defined] def _get_iterator(self) -> '_BaseDataLoaderIter': if self.num_workers == 0: return _SingleProcessDataLoaderIter(self) else: self.check_worker_number_rationality() return _MultiProcessingDataLoaderIter(self) @property def multiprocessing_context(self): return self.__multiprocessing_context @multiprocessing_context.setter def multiprocessing_context(self, multiprocessing_context): if multiprocessing_context is not None: if self.num_workers > 0: if isinstance(multiprocessing_context, string_classes): valid_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_context not in valid_start_methods: raise ValueError( ('multiprocessing_context option ' 'should specify a valid start method in {!r}, but got ' 'multiprocessing_context={!r}').format(valid_start_methods, multiprocessing_context)) # error: Argument 1 to "get_context" has incompatible type "Union[str, bytes]"; expected "str" [arg-type] multiprocessing_context = multiprocessing.get_context(multiprocessing_context) # type: ignore[arg-type] if not isinstance(multiprocessing_context, multiprocessing.context.BaseContext): raise TypeError(('multiprocessing_context option should be a valid context ' 'object or a string specifying the start method, but got ' 'multiprocessing_context={}').format(multiprocessing_context)) else: raise ValueError(('multiprocessing_context can only be used with ' 'multi-process loading (num_workers > 0), but got ' 'num_workers={}').format(self.num_workers)) self.__multiprocessing_context = multiprocessing_context def __setattr__(self, attr, val): if self.__initialized and attr in ( 'batch_size', 'batch_sampler', 'sampler', 'drop_last', 'dataset', 'persistent_workers'): raise ValueError('{} attribute should not be set after {} is ' 'initialized'.format(attr, self.__class__.__name__)) super(DataLoader, self).__setattr__(attr, val) # We quote '_BaseDataLoaderIter' since it isn't defined yet and the definition can't be moved up # since '_BaseDataLoaderIter' references 'DataLoader'. def __iter__(self) -> '_BaseDataLoaderIter': # When using a single worker the returned iterator should be # created everytime to avoid reseting its state # However, in the case of a multiple workers iterator # the iterator is only created once in the lifetime of the # DataLoader object so that workers can be reused if self.persistent_workers and self.num_workers > 0: if self._iterator is None: self._iterator = self._get_iterator() else: self._iterator._reset(self) return self._iterator else: return self._get_iterator() @property def _auto_collation(self): return self.batch_sampler is not None @property def _index_sampler(self): # The actual sampler used for generating indices for `_DatasetFetcher` # (see _utils/fetch.py) to read data at each time. This would be # `.batch_sampler` if in auto-collation mode, and `.sampler` otherwise. # We can't change `.sampler` and `.batch_sampler` attributes for BC # reasons. if self._auto_collation: return self.batch_sampler else: return self.sampler def __len__(self) -> int: if self._dataset_kind == _DatasetKind.Iterable: length = self._IterableDataset_len_called = len(self.dataset) # type: ignore[assignment, arg-type] if self.batch_size is not None: # IterableDataset doesn't allow custom sampler or batch_sampler from math import ceil if self.drop_last: length = length // self.batch_size else: length = ceil(length / self.batch_size) return length else: return len(self._index_sampler) def check_worker_number_rationality(self): def _create_warning_msg(num_worker_suggest, num_worker_created, cpuset_checked): suggested_max_worker_msg = (( "Our suggested max number of worker in current system is {}{}, which is smaller " "than what this DataLoader is going to create.").format( num_worker_suggest, ("" if cpuset_checked else " (`cpuset` is not taken into account)")) ) if num_worker_suggest is not None else ( "DataLoader is not able to compute a suggested max number of worker in current system.") warn_msg = ( "This DataLoader will create {} worker processes in total. {} " "Please be aware that excessive worker creation might get DataLoader running slow or even freeze, " "lower the worker number to avoid potential slowness/freeze if necessary.").format( num_worker_created, suggested_max_worker_msg) return warn_msg if not self.num_workers or self.num_workers == 0: return # try to compute a suggested max number of worker based on system's resource max_num_worker_suggest = None cpuset_checked = False if hasattr(os, 'sched_getaffinity'): try: max_num_worker_suggest = len(os.sched_getaffinity(0)) cpuset_checked = True except Exception: pass if max_num_worker_suggest is None: # os.cpu_count() could return Optional[int] # get cpu count first and check None in order to satify mypy check cpu_count = os.cpu_count() if cpu_count is not None: max_num_worker_suggest = cpu_count if max_num_worker_suggest is None: warnings.warn(_create_warning_msg( max_num_worker_suggest, self.num_workers, cpuset_checked)) return if self.num_workers > max_num_worker_suggest: warnings.warn(_create_warning_msg( max_num_worker_suggest, self.num_workers, cpuset_checked)) def _get_shared_seed(self): # if isinstance(self.dataset, IterDataPipe): # # _shared_seed = torch.empty((), dtype=torch.int64).random_(generator=self.generator).item() # _shared_seed = np.random.randint(2**31) # # TODO torch.dist # if dist.is_available() and dist.is_initialized(): # rank = dist.get_rank() # ws = dist.get_world_size() # store = dist.distributed_c10d._get_default_store() # if rank == 0: # _shared_seed_str = str(_shared_seed) # store.set(_utils.DATAPIPE_SHARED_SEED, _shared_seed_str) # logger.info(f"Shared seed ({_shared_seed_str}) sent to store on rank 0") # # Use 'add' instead of 'get' since for some store implementations 'add' # # doesn't work well with 'get'. # _shared_seed_recv_cnt = store.add(_utils.DATAPIPE_SHARED_SEED_COUNTER, 1) # start = time.time() # while _shared_seed_recv_cnt < ws: # time.sleep(_utils.DATAPIPE_SHARED_SEED_CHECK_INTERVAL) # _shared_seed_recv_cnt = store.add(_utils.DATAPIPE_SHARED_SEED_COUNTER, 0) # if timedelta(seconds=(time.time() - start)) > \ # timedelta(seconds=_utils.DATAPIPE_SHARED_SEED_DEFAULT_TIMEOUT): # raise RuntimeError("Timed out receiving the signal from the distribtued store on " # "Rank 0 that all other Ranks have received the shared seed. " # f"(world_size={ws}, received={_shared_seed_recv_cnt}, " # f"timeout={_utils.DATAPIPE_SHARED_SEED_DEFAULT_TIMEOUT})") # # Reset after all distributed processes have received the shared seed # store.set(_utils.DATAPIPE_SHARED_SEED, "") # _shared_seed_recv_cnt = store.add(_utils.DATAPIPE_SHARED_SEED_COUNTER, -ws) # assert _shared_seed_recv_cnt == 0 # else: # _shared_seed_str = "" # start = time.time() # while len(_shared_seed_str) == 0: # time.sleep(_utils.DATAPIPE_SHARED_SEED_CHECK_INTERVAL) # _shared_seed_str = store.get(_utils.DATAPIPE_SHARED_SEED) # if timedelta(seconds=(time.time() - start)) > \ # timedelta(seconds=_utils.DATAPIPE_SHARED_SEED_DEFAULT_TIMEOUT): # raise RuntimeError("Timed out receiving the shared seed from the distribtued store " # f"on Rank {rank}. (world_size={ws}, " # f"timeout={_utils.DATAPIPE_SHARED_SEED_DEFAULT_TIMEOUT})") # logger.info(f"Shared seed ({_shared_seed_str}) received from store on rank {rank}") # _shared_seed_recv_cnt = store.add(_utils.DATAPIPE_SHARED_SEED_COUNTER, 1) # # Exit only when all ranks received seed, otherwise we are at risk that current rank # # will reach same section of the code again while rank zero still in the previous iteration # while _shared_seed_recv_cnt > 0: # time.sleep(_utils.DATAPIPE_SHARED_SEED_CHECK_INTERVAL) # _shared_seed_recv_cnt = store.add(_utils.DATAPIPE_SHARED_SEED_COUNTER, 0) # _shared_seed = int(_shared_seed_str) # return _shared_seed # else: return None class _BaseDataLoaderIter(object): def __init__(self, loader: DataLoader) -> None: self._dataset = loader.dataset self._shared_seed = loader._get_shared_seed() # if isinstance(self._dataset, IterDataPipe): #TODO datapipe # shared_rng = torch.Generator() # shared_rng.manual_seed(self._shared_seed) # self._dataset = ms_adapter.pytorch.utils.data.graph_settings.apply_shuffle_seed(self._dataset, shared_rng) self._dataset_kind = loader._dataset_kind self._IterableDataset_len_called = loader._IterableDataset_len_called self._auto_collation = loader._auto_collation self._drop_last = loader.drop_last self._index_sampler = loader._index_sampler self._num_workers = loader.num_workers self._prefetch_factor = loader.prefetch_factor # for other backends, pin_memory_device need to set. if not set # default behaviour is CUDA device. if pin_memory_device is selected # and pin_memory is not set, the default behaviour false. # if (len(loader.pin_memory_device) == 0):#TODO pin_memory # # self._pin_memory = loader.pin_memory and torch.cuda.is_available() #TODO torch.cuda.is_available() 无对应mindspore 算子 # self._pin_memory = loader.pin_memory # self._pin_memory_device = None # else: # if not loader.pin_memory: # warn_msg = ("pin memory device is set and pin_memory flag is not used then device pinned memory won't be used" # "please set pin_memory to true, if you need to use the device pin memory") # warnings.warn(warn_msg) # # self._pin_memory = loader.pin_memory # self._pin_memory_device = loader.pin_memory_device self._timeout = loader.timeout self._collate_fn = loader.collate_fn self._sampler_iter = iter(self._index_sampler) # self._base_seed = torch.empty((), dtype=torch.int64).random_(generator=loader.generator).item() #TODO self._base_seed = np.random.randint(2**31) self._persistent_workers = loader.persistent_workers self._num_yielded = 0 self._profile_name = "enumerate(DataLoader)#{}.__next__".format(self.__class__.__name__) def __iter__(self) -> '_BaseDataLoaderIter': return self def _reset(self, loader, first_iter=False): self._sampler_iter = iter(self._index_sampler) self._num_yielded = 0 self._IterableDataset_len_called = loader._IterableDataset_len_called self._shared_seed = loader._get_shared_seed() # if isinstance(self._dataset, IterDataPipe): #TODO pipe # shared_rng = torch.Generator() # shared_rng.manual_seed(self._shared_seed) # self._dataset = ms_adapter.pytorch.utils.data.graph_settings.apply_shuffle_seed(self._dataset, shared_rng) def _next_index(self): return next(self._sampler_iter) # may raise StopIteration def _next_data(self): raise NotImplementedError def __next__(self) -> Any: if self._sampler_iter is None: self._reset() # type: ignore[call-arg] data = self._next_data() self._num_yielded += 1 if self._dataset_kind == _DatasetKind.Iterable and \ self._IterableDataset_len_called is not None and \ self._num_yielded > self._IterableDataset_len_called: warn_msg = ("Length of IterableDataset {} was reported to be {} (when accessing len(dataloader)), but {} " "samples have been fetched. ").format(self._dataset, self._IterableDataset_len_called, self._num_yielded) if self._num_workers > 0: warn_msg += ("For multiprocessing data-loading, this could be caused by not properly configuring the " "IterableDataset replica at each worker. ") warnings.warn(warn_msg) return data next = __next__ # Python 2 compatibility def __len__(self) -> int: return len(self._index_sampler) def __getstate__(self): raise NotImplementedError("{} cannot be pickled", self.__class__.__name__) class _SingleProcessDataLoaderIter(_BaseDataLoaderIter): def __init__(self, loader): super(_SingleProcessDataLoaderIter, self).__init__(loader) assert self._timeout == 0 assert self._num_workers == 0 self._dataset_fetcher = _DatasetKind.create_fetcher( self._dataset_kind, self._dataset, self._auto_collation, self._collate_fn, self._drop_last) def _next_data(self): index = self._next_index() # may raise StopIteration data = self._dataset_fetcher.fetch(index) # may raise StopIteration # if self._pin_memory:#TODO pin_memory # data = _utils.pin_memory.pin_memory(data, self._pin_memory_device) return data class _MultiProcessingDataLoaderIter(_BaseDataLoaderIter): def __init__(self, loader): super(_MultiProcessingDataLoaderIter, self).__init__(loader) assert self._num_workers > 0 assert self._prefetch_factor > 0 if loader.multiprocessing_context is None: multiprocessing_context = multiprocessing else: multiprocessing_context = loader.multiprocessing_context self._worker_init_fn = loader.worker_init_fn # No certainty which module multiprocessing_context is self._worker_result_queue = multiprocessing_context.Queue() # type: ignore[var-annotated] self._worker_pids_set = False self._shutdown = False self._workers_done_event = multiprocessing_context.Event() self._index_queues = [] self._workers = [] for i in range(self._num_workers): # No certainty which module multiprocessing_context is index_queue = multiprocessing_context.Queue() # type: ignore[var-annotated] # Need to `cancel_join_thread` here! # See sections (2) and (3b) above. index_queue.cancel_join_thread() w = multiprocessing_context.Process( target=_utils.worker._worker_loop, args=(self._dataset_kind, self._dataset, index_queue, self._worker_result_queue, self._workers_done_event, self._auto_collation, self._collate_fn, self._drop_last, self._base_seed, self._worker_init_fn, i, self._num_workers, self._persistent_workers, self._shared_seed)) w.daemon = True # NB: Process.start() actually take some time as it needs to # start a process and pass the arguments over via a pipe. # Therefore, we only add a worker to self._workers list after # it started, so that we do not call .join() if program dies # before it starts, and __del__ tries to join but will get: # AssertionError: can only join a started process. w.start() self._index_queues.append(index_queue) self._workers.append(w) # if self._pin_memory:#TODO pin_memory # self._pin_memory_thread_done_event = threading.Event() # # # Queue is not type-annotated # self._data_queue = queue.Queue() # type: ignore[var-annotated] # pin_memory_thread = threading.Thread( # target=_utils.pin_memory._pin_memory_loop, # args=(self._worker_result_queue, self._data_queue, # torch.cuda.current_device(), # self._pin_memory_thread_done_event, self._pin_memory_device)) #TODO # pin_memory_thread.daemon = True # pin_memory_thread.start() # # Similar to workers (see comment above), we only register # # pin_memory_thread once it is started. # self._pin_memory_thread = pin_memory_thread # else: self._data_queue = self._worker_result_queue # In some rare cases, persistent workers (daemonic processes) # would be terminated before `__del__` of iterator is invoked # when main process exits # It would cause failure when pin_memory_thread tries to read # corrupted data from worker_result_queue # atexit is used to shutdown thread and child processes in the # right sequence before main process exits # if self._persistent_workers and self._pin_memory: #TODO pin_memory # import atexit # for w in self._workers: # atexit.register(_MultiProcessingDataLoaderIter._clean_up_worker, w) # .pid can be None only before process is spawned (not the case, so ignore) # _utils.signal_handling._set_worker_pids(id(self), tuple(w.pid for w in self._workers)) # type: ignore[misc] # _utils.signal_handling._set_SIGCHLD_handler() self._worker_pids_set = True self._reset(loader, first_iter=True) def _reset(self, loader, first_iter=False): super()._reset(loader, first_iter) self._send_idx = 0 # idx of the next task to be sent to workers self._rcvd_idx = 0 # idx of the next task to be returned in __next__ # information about data not yet yielded, i.e., tasks w/ indices in range [rcvd_idx, send_idx). # map: task idx => - (worker_id,) if data isn't fetched (outstanding) # \ (worker_id, data) if data is already fetched (out-of-order) self._task_info = {} self._tasks_outstanding = 0 # always equal to count(v for v in task_info.values() if len(v) == 1) # A list of booleans representing whether each worker still has work to # do, i.e., not having exhausted its iterable dataset object. It always # contains all `True`s if not using an iterable-style dataset # (i.e., if kind != Iterable). # Not that this indicates that a worker still has work to do *for this epoch*. # It does not mean that a worker is dead. In case of `_persistent_workers`, # the worker will be reset to available in the next epoch. self._workers_status = [True for i in range(self._num_workers)] # Reset the worker queue cycle so it resumes next epoch at worker 0 self._worker_queue_idx_cycle = itertools.cycle(range(self._num_workers)) # We resume the prefetching in case it was enabled if not first_iter: for idx in range(self._num_workers): self._index_queues[idx].put(_utils.worker._ResumeIteration(self._shared_seed)) resume_iteration_cnt = self._num_workers while resume_iteration_cnt > 0: return_idx, return_data = self._get_data() if isinstance(return_idx, _utils.worker._ResumeIteration): assert return_data is None resume_iteration_cnt -= 1 # prime the prefetch loop for _ in range(self._prefetch_factor * self._num_workers): self._try_put_index() def _try_get_data(self, timeout=_utils.MP_STATUS_CHECK_INTERVAL): # Tries to fetch data from `self._data_queue` once for a given timeout. # This can also be used as inner loop of fetching without timeout, with # the sender status as the loop condition. # # This raises a `RuntimeError` if any worker died expectedly. This error # can come from either the SIGCHLD handler in `_utils/signal_handling.py` # (only for non-Windows platforms), or the manual check below on errors # and timeouts. # # Returns a 2-tuple: # (bool: whether successfully get data, any: data if successful else None) try: data = self._data_queue.get(timeout=timeout) return (True, data) except Exception as e: # At timeout and error, we manually check whether any worker has # failed. Note that this is the only mechanism for Windows to detect # worker failures. failed_workers = [] for worker_id, w in enumerate(self._workers): if self._workers_status[worker_id] and not w.is_alive(): failed_workers.append(w) self._mark_worker_as_unavailable(worker_id) if len(failed_workers) > 0: pids_str = ', '.join(str(w.pid) for w in failed_workers) raise RuntimeError('DataLoader worker (pid(s) {}) exited unexpectedly'.format(pids_str)) from e if isinstance(e, queue.Empty): return (False, None) import tempfile import errno try: # Raise an exception if we are this close to the FDs limit. # Apparently, trying to open only one file is not a sufficient # test. # See NOTE [ DataLoader on Linux and open files limit ] fds_limit_margin = 10 fs = [tempfile.NamedTemporaryFile() for i in range(fds_limit_margin)] except OSError as e: if e.errno == errno.EMFILE: raise RuntimeError( "Too many open files. Communication with the" " workers is no longer possible.") from None raise def _get_data(self): # Fetches data from `self._data_queue`. # # We check workers' status every `MP_STATUS_CHECK_INTERVAL` seconds, # which we achieve by running `self._try_get_data(timeout=MP_STATUS_CHECK_INTERVAL)` # in a loop. This is the only mechanism to detect worker failures for # Windows. For other platforms, a SIGCHLD handler is also used for # worker failure detection. # # If `pin_memory=True`, we also need check if `pin_memory_thread` had # died at timeouts. if self._timeout > 0: success, data = self._try_get_data(self._timeout) if success: return data else: raise RuntimeError('DataLoader timed out after {} seconds'.format(self._timeout)) # elif self._pin_memory: # while self._pin_memory_thread.is_alive(): # success, data = self._try_get_data() # if success: # return data # else: # # while condition is false, i.e., pin_memory_thread died. # raise RuntimeError('Pin memory thread exited unexpectedly') # # In this case, `self._data_queue` is a `queue.Queue`,. But we don't # # need to call `.task_done()` because we don't use `.join()`. else: while True: success, data = self._try_get_data() if success: return data def _next_data(self): while True: # If the worker responsible for `self._rcvd_idx` has already ended # and was unable to fulfill this task (due to exhausting an `IterableDataset`), # we try to advance `self._rcvd_idx` to find the next valid index. # # This part needs to run in the loop because both the `self._get_data()` # call and `_IterableDatasetStopIteration` check below can mark # extra worker(s) as dead. while self._rcvd_idx < self._send_idx: info = self._task_info[self._rcvd_idx] worker_id = info[0] if len(info) == 2 or self._workers_status[worker_id]: # has data or is still active break del self._task_info[self._rcvd_idx] self._rcvd_idx += 1 else: # no valid `self._rcvd_idx` is found (i.e., didn't break) if not self._persistent_workers: self._shutdown_workers() raise StopIteration # Now `self._rcvd_idx` is the batch index we want to fetch # Check if the next sample has already been generated if len(self._task_info[self._rcvd_idx]) == 2: data = self._task_info.pop(self._rcvd_idx)[1] return self._process_data(data) assert not self._shutdown and self._tasks_outstanding > 0 idx, data = self._get_data() self._tasks_outstanding -= 1 if self._dataset_kind == _DatasetKind.Iterable: # Check for _IterableDatasetStopIteration if isinstance(data, _utils.worker._IterableDatasetStopIteration): if self._persistent_workers: self._workers_status[data.worker_id] = False else: self._mark_worker_as_unavailable(data.worker_id) self._try_put_index() continue if idx != self._rcvd_idx: # store out-of-order samples self._task_info[idx] += (data,) else: del self._task_info[idx] return self._process_data(data) def _try_put_index(self): assert self._tasks_outstanding < self._prefetch_factor * self._num_workers try: index = self._next_index() except StopIteration: return for _ in range(self._num_workers): # find the next active worker, if any worker_queue_idx = next(self._worker_queue_idx_cycle) if self._workers_status[worker_queue_idx]: break else: # not found (i.e., didn't break) return self._index_queues[worker_queue_idx].put((self._send_idx, index)) self._task_info[self._send_idx] = (worker_queue_idx,) self._tasks_outstanding += 1 self._send_idx += 1 def _process_data(self, data): self._rcvd_idx += 1 self._try_put_index() if isinstance(data, ExceptionWrapper): data.reraise() return data def _mark_worker_as_unavailable(self, worker_id, shutdown=False): # Mark a worker as having finished its work e.g., due to # exhausting an `IterableDataset`. This should be used only when this # `_MultiProcessingDataLoaderIter` is going to continue running. assert self._workers_status[worker_id] or (self._persistent_workers and shutdown) # Signal termination to that specific worker. q = self._index_queues[worker_id] # Indicate that no more data will be put on this queue by the current # process. q.put(None) # Note that we don't actually join the worker here, nor do we remove the # worker's pid from C side struct because (1) joining may be slow, and # (2) since we don't join, the worker may still raise error, and we # prefer capturing those, rather than ignoring them, even though they # are raised after the worker has finished its job. # Joinning is deferred to `_shutdown_workers`, which it is called when # all workers finish their jobs (e.g., `IterableDataset` replicas) or # when this iterator is garbage collected. self._workers_status[worker_id] = False assert self._workers_done_event.is_set() == shutdown def _shutdown_workers(self): # Called when shutting down this `_MultiProcessingDataLoaderIter`. # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on # the logic of this function. python_exit_status = _utils.python_exit_status if python_exit_status is True or python_exit_status is None: # See (2) of the note. If Python is shutting down, do no-op. return # Normal exit when last reference is gone / iterator is depleted. # See (1) and the second half of the note. if not self._shutdown: self._shutdown = True try: # Normal exit when last reference is gone / iterator is depleted. # See (1) and the second half of the note. # Exit `pin_memory_thread` first because exiting workers may leave # corrupted data in `worker_result_queue` which `pin_memory_thread` # reads from. if hasattr(self, '_pin_memory_thread'): # Use hasattr in case error happens before we set the attribute. self._pin_memory_thread_done_event.set() # Send something to pin_memory_thread in case it is waiting # so that it can wake up and check `pin_memory_thread_done_event` self._worker_result_queue.put((None, None)) self._pin_memory_thread.join() self._worker_result_queue.cancel_join_thread() self._worker_result_queue.close() # Exit workers now. self._workers_done_event.set() for worker_id in range(len(self._workers)): # Get number of workers from `len(self._workers)` instead of # `self._num_workers` in case we error before starting all # workers. # If we are using workers_status with persistent_workers # we have to shut it down because the worker is paused if self._persistent_workers or self._workers_status[worker_id]: self._mark_worker_as_unavailable(worker_id, shutdown=True) for w in self._workers: # We should be able to join here, but in case anything went # wrong, we set a timeout and if the workers fail to join, # they are killed in the `finally` block. w.join(timeout=_utils.MP_STATUS_CHECK_INTERVAL) for q in self._index_queues: q.cancel_join_thread() q.close() finally: # Even though all this function does is putting into queues that # we have called `cancel_join_thread` on, weird things can # happen when a worker is killed by a signal, e.g., hanging in # `Event.set()`. So we need to guard this with SIGCHLD handler, # and remove pids from the C side data structure only at the # end. # # FIXME: Unfortunately, for Windows, we are missing a worker # error detection mechanism here in this function, as it # doesn't provide a SIGCHLD handler. if self._worker_pids_set: # _utils.signal_handling._remove_worker_pids(id(self)) self._worker_pids_set = False for w in self._workers: if w.is_alive(): # Existing mechanisms try to make the workers exit # peacefully, but in case that we unfortunately reach # here, which we shouldn't, (e.g., pytorch/pytorch#39570), # we kill the worker. w.terminate() # staticmethod is used to remove reference to `_MultiProcessingDataLoaderIter` @staticmethod def _clean_up_worker(w): try: w.join(timeout=_utils.MP_STATUS_CHECK_INTERVAL) finally: if w.is_alive(): w.terminate() def __del__(self): self._shutdown_workers() msadapter/ms_adapter/pytorch/utils/data/dataloader_experimental.py000066400000000000000000000160361436713554100261640ustar00rootroot00000000000000# from __future__ import absolute_import # # import time # # from typing import Any, List # # import ms_adapter.pytorch.utils.data.backward_compatibility # # import ms_adapter.pytorch.utils.data.graph_settings # from ms_adapter.pytorch.utils.data import DataLoader, IterDataPipe, communication # from ms_adapter.pytorch.utils.data.datapipes.iter import IterableWrapper # # __all__ = [ # "DataLoader2", # ] # # # class _ThreadingDataLoader2: # # def __init__(self, datapipe, num_workers=0, collate_fn=None): # self.threads = [] # self.datapipes = [] # self.collate_fn = collate_fn # for worker_id in range(num_workers): # (thread, req_queue, res_queue, thread_localdatapipe) = communication.eventloop.SpawnThreadForDataPipeline(datapipe) # ms_adapter.pytorch.utils.data.graph_settings.apply_sharding(thread_localdatapipe, num_workers, worker_id) # thread.start() # self.threads.append((thread, req_queue, res_queue)) # These queues are independent # local_datapipe = communication.iter.QueueWrapper( # communication.protocol.IterDataPipeQueueProtocolClient(req_queue, res_queue)) # self.datapipes.append(local_datapipe) # # def __iter__(self): # not_available = False # forever = True # exclude_datapipes: List[Any] = [] # while len(exclude_datapipes) < len(self.datapipes): # for dp in self.datapipes: # if dp not in exclude_datapipes: # try: # value = dp.nonblocking_next() # yield value # except StopIteration: # exclude_datapipes.append(dp) # except communication.iter.NotAvailable: # not_available = True # if not_available: # time.sleep(0.001) # # def __del__(self): # self._cleanup_all_threads() # # def _cleanup_all_threads(self): # def clean_me(thread, req_queue, res_queue): # req_queue.put(communication.messages.TerminateRequest()) # _ = res_queue.get() # thread.join() # # for thread, req_queue, res_queue in self.threads: # clean_me(thread, req_queue, res_queue) # # class DataLoader2: # def __new__(cls, # dataset, # batch_size=1, # shuffle=None, # sampler=None, # batch_sampler=None, # num_workers=0, # collate_fn=None, # pin_memory=False, # drop_last=False, # timeout=0, # worker_init_fn=None, # *, # prefetch_factor=2, # persistent_workers=False, # batch_outside_worker=False, # parallelism_mode='mp'): # if isinstance(dataset, IterDataPipe): # data_loader: Any = None # if batch_sampler is not None: # raise Exception( # 'batch_sampler is not yet supported by DataPipes') # if sampler is not None: # raise Exception( # 'sampler is not yet supported by DataPipes') # datapipe = dataset # datapipe = ms_adapter.pytorch.utils.data.graph_settings.apply_shuffle_settings(datapipe, shuffle=shuffle) # if batch_outside_worker and pin_memory: # raise Exception( # 'pin_memory is not yet compatible with batch_outside_worker') # if not batch_outside_worker: # if batch_size is not None: # datapipe = datapipe.batch(batch_size, drop_last=drop_last) # if collate_fn is None: # collate_fn = ms_adapter.pytorch.utils.data._utils.collate.default_collate # # # Note: It is safe to pass shuffle=True to the old DataLoader, as shuffle does nothing # # for Iterable, but required to set Pipes correctly. # data_loader = DataLoader(datapipe, # batch_size=None, # Replaced by .batch DataPipe # shuffle=shuffle, # sampler=None, # batch_sampler=None, # num_workers=num_workers, # collate_fn=collate_fn, # pin_memory=pin_memory, # drop_last=False, # Replaced by .batch DataPipe # timeout=timeout, # worker_init_fn=worker_init_fn, # prefetch_factor=prefetch_factor, # persistent_workers=persistent_workers) # elif parallelism_mode == 'thread': # if collate_fn is not None and not batch_outside_worker: # datapipe = datapipe.map(collate_fn) # if pin_memory: # raise Exception( # 'pin_memory is not yet supported by DataPipes with Threading') # if worker_init_fn is not None: # raise Exception( # 'worker_init_fn is not yet supported by DataPipes with Threading') # data_loader = _ThreadingDataLoader2(datapipe, # num_workers=num_workers, # collate_fn=collate_fn) # else: # raise Exception('Unsupported parallelism mode', parallelism_mode) # if not batch_outside_worker: # return data_loader # else: # if collate_fn is None: # collate_fn = ms_adapter.pytorch.utils.data._utils.collate.default_collate # datapipe = IterableWrapper(data_loader).batch( # batch_size, drop_last=drop_last).map(collate_fn) # return datapipe # else: # if parallelism_mode == 'thread': # raise Exception( # 'thread parallelism mode is not supported for old DataSets') # return DataLoader(dataset, # batch_size=batch_size, # shuffle=shuffle, # sampler=sampler, # batch_sampler=batch_sampler, # num_workers=num_workers, # collate_fn=collate_fn, # pin_memory=pin_memory, # drop_last=drop_last, # timeout=timeout, # worker_init_fn=worker_init_fn, # prefetch_factor=prefetch_factor, # persistent_workers=persistent_workers) msadapter/ms_adapter/pytorch/utils/data/datapipes/000077500000000000000000000000001436713554100227015ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/datapipes/__init__.py000066400000000000000000000000751436713554100250140ustar00rootroot00000000000000from . import iter from . import map from . import dataframe msadapter/ms_adapter/pytorch/utils/data/datapipes/_decorator.py000066400000000000000000000171011436713554100253740ustar00rootroot00000000000000from __future__ import absolute_import import inspect from functools import wraps from typing import Any, Callable, Optional, Type, Union, get_type_hints from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe, MapDataPipe from ms_adapter.pytorch.utils.data.datapipes._typing import _DataPipeMeta ###################################################### # Functional API ###################################################### class functional_datapipe(object): name: str def __init__(self, name: str, enable_df_api_tracing=False) -> None: """ Args: enable_df_api_tracing - if set, any returned DataPipe would accept DataFrames API in tracing mode. """ self.name = name self.enable_df_api_tracing = enable_df_api_tracing def __call__(self, cls): if issubclass(cls, IterDataPipe): if isinstance(cls, Type): # type: ignore[arg-type] if not isinstance(cls, _DataPipeMeta): raise TypeError('`functional_datapipe` can only decorate IterDataPipe') # with non_deterministic decorator else: if not isinstance(cls, non_deterministic) and \ not (hasattr(cls, '__self__') and isinstance(cls.__self__, non_deterministic)): raise TypeError('`functional_datapipe` can only decorate IterDataPipe') IterDataPipe.register_datapipe_as_function(self.name, cls, enable_df_api_tracing=self.enable_df_api_tracing) elif issubclass(cls, MapDataPipe): MapDataPipe.register_datapipe_as_function(self.name, cls) return cls ###################################################### # Determinism ###################################################### _determinism: bool = False class guaranteed_datapipes_determinism(object): prev: bool def __init__(self) -> None: global _determinism self.prev = _determinism _determinism = True def __enter__(self) -> None: pass def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: global _determinism _determinism = self.prev class non_deterministic(object): cls: Optional[Type[IterDataPipe]] = None # TODO: Lambda for picking deterministic_fn: Callable[[], bool] def __init__(self, arg: Union[Type[IterDataPipe], Callable[[], bool]]) -> None: # 1. Decorator doesn't have any argument if isinstance(arg, Type): # type: ignore[arg-type] if not issubclass(arg, IterDataPipe): # type: ignore[arg-type] raise TypeError("Only `IterDataPipe` can be decorated with `non_deterministic`" ", but {} is found".format(arg.__name__)) self.cls = arg # type: ignore[assignment] # 2. Decorator has an argument of a function # This class should behave differently given different inputs. Use this # function to verify the determinism for each instance. # When the function returns True, the instance is non-deterministic. Otherwise, # the instance is a deterministic DataPipe. elif isinstance(arg, Callable): # type:ignore[arg-type] self.deterministic_fn = arg # type: ignore[assignment, misc] else: raise TypeError("{} can not be decorated by non_deterministic".format(arg)) def __call__(self, *args, **kwargs): global _determinism # Decorate IterDataPipe if self.cls is not None: if _determinism: raise TypeError("{} is non-deterministic, but you set 'guaranteed_datapipes_determinism'. " "You can turn off determinism for this DataPipe if that is acceptable " "for your application".format(self.cls.__name__)) return self.cls(*args, **kwargs) # type: ignore[call-arg] # Decorate with a functional argument if not (isinstance(args[0], Type) and # type: ignore[arg-type] issubclass(args[0], IterDataPipe)): raise TypeError("Only `IterDataPipe` can be decorated, but {} is found" .format(args[0].__name__)) self.cls = args[0] return self.deterministic_wrapper_fn def deterministic_wrapper_fn(self, *args, **kwargs) -> IterDataPipe: res = self.deterministic_fn(*args, **kwargs) # type: ignore[call-arg, misc] if not isinstance(res, bool): raise TypeError("deterministic_fn of `non_deterministic` decorator is required " "to return a boolean value, but {} is found".format(type(res))) global _determinism if _determinism and res: raise TypeError("{} is non-deterministic with the inputs, but you set " "'guaranteed_datapipes_determinism'. You can turn off determinism " "for this DataPipe if that is acceptable for your application" .format(self.cls.__name__)) # type: ignore[union-attr] return self.cls(*args, **kwargs) # type: ignore[call-arg, misc] ###################################################### # Type validation ###################################################### # Validate each argument of DataPipe with hint as a subtype of the hint. def argument_validation(f): signature = inspect.signature(f) hints = get_type_hints(f) @wraps(f) def wrapper(*args, **kwargs): bound = signature.bind(*args, **kwargs) for argument_name, value in bound.arguments.items(): if argument_name in hints and isinstance(hints[argument_name], _DataPipeMeta): hint = hints[argument_name] if not isinstance(value, IterDataPipe): raise TypeError("Expected argument '{}' as a IterDataPipe, but found {}" .format(argument_name, type(value))) if not value.type.issubtype(hint.type): raise TypeError("Expected type of argument '{}' as a subtype of " "hint {}, but found {}" .format(argument_name, hint.type, value.type)) return f(*args, **kwargs) return wrapper # Default value is True _runtime_validation_enabled: bool = True class runtime_validation_disabled(object): prev: bool def __init__(self) -> None: global _runtime_validation_enabled self.prev = _runtime_validation_enabled _runtime_validation_enabled = False def __enter__(self) -> None: pass def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: global _runtime_validation_enabled _runtime_validation_enabled = self.prev # Runtime checking # Validate output data is subtype of return hint def runtime_validation(f): # TODO: # Can be extended to validate '__getitem__' and nonblocking if f.__name__ != '__iter__': raise TypeError("Can not decorate function {} with 'runtime_validation'" .format(f.__name__)) @wraps(f) def wrapper(self): global _runtime_validation_enabled if not _runtime_validation_enabled: yield from f(self) else: it = f(self) for d in it: if not self.type.issubtype_of_instance(d): raise RuntimeError("Expected an instance as subtype of {}, but found {}({})" .format(self.type, d, type(d))) yield d return wrapper msadapter/ms_adapter/pytorch/utils/data/datapipes/_typing.py000066400000000000000000000602341436713554100247310ustar00rootroot00000000000000# Taking reference from official Python typing # https://github.com/python/cpython/blob/master/Lib/typing.py import collections import functools import inspect import numbers import sys from typing import (Any, Dict, Iterator, Generic, List, Set, Tuple, TypeVar, Union, get_type_hints) from typing import _eval_type, _tp_cache, _type_check, _type_repr # type: ignore[attr-defined] from typing import ForwardRef # TODO: Use TypeAlias when Python 3.6 is deprecated # Please check [Note: TypeMeta and TypeAlias] # In case of metaclass conflict due to ABCMeta or _ProtocolMeta # For Python 3.9, only Protocol in typing uses metaclass from abc import ABCMeta from typing import _GenericAlias # type: ignore[attr-defined, no-redef] class GenericMeta(ABCMeta): # type: ignore[no-redef] pass import torch class Integer(numbers.Integral): pass class Boolean(numbers.Integral): pass # Python 'type' object is not subscriptable # Tuple[int, List, dict] -> valid # tuple[int, list, dict] -> invalid # Map Python 'type' to abstract base class TYPE2ABC = { bool: Boolean, int: Integer, float: numbers.Real, complex: numbers.Complex, dict: Dict, list: List, set: Set, tuple: Tuple, None: type(None), } def issubtype(left, right, recursive=True): r""" Check if the left-side type is a subtype of the right-side type. If any of type is a composite type like `Union` and `TypeVar` with bounds, it would be expanded into a list of types and check all of left-side types are subtypes of either one from right-side types. """ left = TYPE2ABC.get(left, left) right = TYPE2ABC.get(right, right) if right is Any or left == right: return True if isinstance(right, _GenericAlias): if getattr(right, '__origin__', None) is Generic: return True if right == type(None): return False # Right-side type constraints = _decompose_type(right) if len(constraints) == 0 or Any in constraints: return True if left is Any: return False # Left-side type variants = _decompose_type(left) # all() will return True for empty variants if len(variants) == 0: return False return all(_issubtype_with_constraints(variant, constraints, recursive) for variant in variants) def _decompose_type(t, to_list=True): if isinstance(t, TypeVar): if t.__bound__ is not None: ts = [t.__bound__] else: # For T_co, __constraints__ is () ts = list(t.__constraints__) elif hasattr(t, '__origin__') and t.__origin__ == Union: ts = t.__args__ else: if not to_list: return None ts = [t] # Ignored: Generator has incompatible item type "object"; expected "Type[Any]" ts = list(TYPE2ABC.get(_t, _t) for _t in ts) # type: ignore[misc] return ts def _issubtype_with_constraints(variant, constraints, recursive=True): r""" Check if the variant is a subtype of either one from constraints. For composite types like `Union` and `TypeVar` with bounds, they would be expanded for testing. """ if variant in constraints: return True # [Note: Subtype for Union and TypeVar] # Python typing is able to flatten Union[Union[...]] or Union[TypeVar]. # But it couldn't flatten the following scenarios: # - Union[int, TypeVar[Union[...]]] # - TypeVar[TypeVar[...]] # So, variant and each constraint may be a TypeVar or a Union. # In these cases, all of inner types from the variant are required to be # extraced and verified as a subtype of any constraint. And, all of # inner types from any constraint being a TypeVar or a Union are # also required to be extracted and verified if the variant belongs to # any of them. # Variant vs = _decompose_type(variant, to_list=False) # Variant is TypeVar or Union if vs is not None: return all(_issubtype_with_constraints(v, constraints, recursive) for v in vs) # Variant is not TypeVar or Union if hasattr(variant, '__origin__') and variant.__origin__ is not None: v_origin = variant.__origin__ # In Python-3.9 typing library untyped generics do not have args v_args = getattr(variant, "__args__", None) else: v_origin = variant v_args = None # Constraints for constraint in constraints: cs = _decompose_type(constraint, to_list=False) # Constraint is TypeVar or Union if cs is not None: if _issubtype_with_constraints(variant, cs, recursive): return True # Constraint is not TypeVar or Union else: # __origin__ can be None for plain list, tuple, ... in Python 3.6 if hasattr(constraint, '__origin__') and constraint.__origin__ is not None: c_origin = constraint.__origin__ if v_origin == c_origin: if not recursive: return True # In Python-3.9 typing library untyped generics do not have args c_args = getattr(constraint, "__args__", None) if c_args is None or len(c_args) == 0: return True if v_args is not None and len(v_args) == len(c_args) and \ all(issubtype(v_arg, c_arg) for v_arg, c_arg in zip(v_args, c_args)): return True # Tuple[int] -> Tuple else: if v_origin == constraint: return True return False def issubinstance(data, data_type): if not issubtype(type(data), data_type, recursive=False): return False # In Python-3.9 typing library __args__ attribute is not defined for untyped generics dt_args = getattr(data_type, "__args__", None) if isinstance(data, tuple): if dt_args is None or len(dt_args) == 0: return True if len(dt_args) != len(data): return False return all(issubinstance(d, t) for d, t in zip(data, dt_args)) elif isinstance(data, (list, set)): if dt_args is None or len(dt_args) == 0: return True t = dt_args[0] return all(issubinstance(d, t) for d in data) elif isinstance(data, dict): if dt_args is None or len(dt_args) == 0: return True kt, vt = dt_args return all(issubinstance(k, kt) and issubinstance(v, vt) for k, v in data.items()) return True # [Note: TypeMeta and TypeAlias] # In order to keep compatibility for Python 3.6, use Meta for the typing. # TODO: When PyTorch drops the support for Python 3.6, it can be converted # into the Alias system and using `__class_getitem__` for DataPipe. The # typing system will gain benefit of performance and resolving metaclass # conflicts as elaborated in https://www.python.org/dev/peps/pep-0560/ class _DataPipeType: r""" Save type annotation in `param` """ def __init__(self, param): self.param = param def __repr__(self): return _type_repr(self.param) def __eq__(self, other): if isinstance(other, _DataPipeType): return self.param == other.param return NotImplemented def __hash__(self): return hash(self.param) def issubtype(self, other): if isinstance(other.param, _GenericAlias): if getattr(other.param, '__origin__', None) is Generic: return True if isinstance(other, _DataPipeType): return issubtype(self.param, other.param) if isinstance(other, type): return issubtype(self.param, other) raise TypeError("Expected '_DataPipeType' or 'type', but found {}".format(type(other))) def issubtype_of_instance(self, other): return issubinstance(other, self.param) # Default type for DataPipe without annotation T_co = TypeVar('T_co', covariant=True) _DEFAULT_TYPE = _DataPipeType(Generic[T_co]) class _DataPipeMeta(GenericMeta): r""" Metaclass for `DataPipe`. Add `type` attribute and `__init_subclass__` based on the type, and validate the return hint of `__iter__`. Note that there is subclass `_IterDataPipeMeta` specifically for `IterDataPipe`. """ type: _DataPipeType def __new__(cls, name, bases, namespace, **kwargs): return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] # TODO: the statements below are not reachable by design as there is a bug and typing is low priority for now. cls.__origin__ = None if 'type' in namespace: return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] namespace['__type_class__'] = False # For plain derived class without annotation for base in bases: if isinstance(base, _DataPipeMeta): return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] namespace.update({'type': _DEFAULT_TYPE, '__init_subclass__': _dp_init_subclass}) return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] def __init__(self, name, bases, namespace, **kwargs): super().__init__(name, bases, namespace, **kwargs) # type: ignore[call-overload] # TODO: Fix isinstance bug @_tp_cache def _getitem_(self, params): if params is None: raise TypeError('{}[t]: t can not be None'.format(self.__name__)) if isinstance(params, str): params = ForwardRef(params) if not isinstance(params, tuple): params = (params, ) msg = "{}[t]: t must be a type".format(self.__name__) params = tuple(_type_check(p, msg) for p in params) if isinstance(self.type.param, _GenericAlias): orig = getattr(self.type.param, '__origin__', None) if isinstance(orig, type) and orig is not Generic: p = self.type.param[params] # type: ignore[index] t = _DataPipeType(p) l = len(str(self.type)) + 2 name = self.__name__[:-l] name = name + '[' + str(t) + ']' bases = (self,) + self.__bases__ return self.__class__(name, bases, {'__init_subclass__': _dp_init_subclass, 'type': t, '__type_class__': True}) if len(params) > 1: raise TypeError('Too many parameters for {} actual {}, expected 1'.format(self, len(params))) t = _DataPipeType(params[0]) if not t.issubtype(self.type): raise TypeError('Can not subclass a DataPipe[{}] from DataPipe[{}]' .format(t, self.type)) # Types are equal, fast path for inheritance if self.type == t: return self name = self.__name__ + '[' + str(t) + ']' bases = (self,) + self.__bases__ return self.__class__(name, bases, {'__init_subclass__': _dp_init_subclass, '__type_class__': True, 'type': t}) # TODO: Fix isinstance bug def _eq_(self, other): if not isinstance(other, _DataPipeMeta): return NotImplemented if self.__origin__ is None or other.__origin__ is None: # type: ignore[has-type] return self is other return (self.__origin__ == other.__origin__ # type: ignore[has-type] and self.type == other.type) # TODO: Fix isinstance bug def _hash_(self): return hash((self.__name__, self.type)) class _IterDataPipeMeta(_DataPipeMeta): r""" Metaclass for `IterDataPipe` and inherits from `_DataPipeMeta`. Aad various functions for behaviors specific to `IterDataPipe`. """ def __new__(cls, name, bases, namespace, **kwargs): if 'reset' in namespace: reset_func = namespace['reset'] @functools.wraps(reset_func) def conditional_reset(*args, **kwargs): r""" Only execute DataPipe's `reset()` method if `_restored` is False. This allows recently restored DataPipe to preserve its restored state during the initial `__iter__` call. """ datapipe = args[0] if datapipe._restored is True: datapipe._restored = False else: reset_func(*args, **kwargs) namespace['reset'] = conditional_reset if '__setstate__' in namespace: setstate_func = namespace['__setstate__'] @functools.wraps(setstate_func) def wrap_setstate(*args, **kwargs): r""" Set `_restored` to True during `__setstate__`, such that the next `reset()` call during iterator creation will not actually reset the state of the DataPipe. """ datapipe = args[0] datapipe._restored = True return setstate_func(*args, **kwargs) namespace['__setstate__'] = wrap_setstate if '__iter__' in namespace: hook_iterator(namespace, 'enumerate(DataPipe)#{}'.format(name)) return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] def _simplify_obj_name(obj) -> str: """ Simplify the display strings of objects for the purpose of rendering within DataPipe error messages. """ if inspect.isfunction(obj): return obj.__name__ else: return repr(obj) def _generate_input_args_string(obj): """ Generate a string for the input arguments of an object. """ signature = inspect.signature(obj.__class__) input_param_names = set() for param_name, _ in signature.parameters.items(): input_param_names.add(param_name) result = [] for name, obj in inspect.getmembers(obj): if name in input_param_names: result.append((name, _simplify_obj_name(obj))) return ', '.join([f'{name}={value}' for name, value in result]) def _generate_iterdatapipe_msg(datapipe): return f"{datapipe.__class__.__name__}({_generate_input_args_string(datapipe)})" def _gen_invalid_iterdatapipe_msg(datapipe): return ("This iterator has been invalidated because another iterator has been created" f"from the same IterDataPipe: {_generate_iterdatapipe_msg(datapipe)}\n" "This may be caused multiple references to the same IterDataPipe. We recommend " "using `.fork()` if that is necessary.") _feedback_msg = ("\nFor feedback regarding this single iterator per IterDataPipe constraint, feel free " "to comment on this issue: https://github.com/pytorch/data/issues/45.") def _check_iterator_valid(datapipe, iterator_id, next_method_exists=False) -> None: r""" Given an instance of a DataPipe and an iterator ID, check if the IDs match, and if not, raises an exception. In the case of ChildDataPipe, the ID gets compared to the one stored in `main_datapipe` as well. """ if next_method_exists: # This is the case where `IterDataPipe` has both `__iter__` and `__next__`. # The `_valid_iterator_id` should either be never set (`None`), or set by at most one # iterator (`0`). Otherwise, it means there are multiple iterators. if datapipe._valid_iterator_id is not None and datapipe._valid_iterator_id != 0: extra_msg = "\nNote that this exception is raised inside your IterDataPipe's a `__next__` method" raise RuntimeError(_gen_invalid_iterdatapipe_msg(datapipe) + extra_msg + _feedback_msg) elif hasattr(datapipe, "_is_child_datapipe") and datapipe._is_child_datapipe is True: if hasattr(datapipe, "_check_valid_iterator_id"): if not datapipe._check_valid_iterator_id(iterator_id): raise RuntimeError("This iterator has been invalidated, because a new iterator has been created " f"from one of the ChildDataPipes of " f"{_generate_iterdatapipe_msg(datapipe.main_datapipe)}." + _feedback_msg) else: raise RuntimeError("ChildDataPipe must have method `_check_valid_iterator_id`.") elif datapipe._valid_iterator_id != iterator_id: raise RuntimeError(_gen_invalid_iterdatapipe_msg(datapipe) + _feedback_msg) def _set_datapipe_valid_iterator_id(datapipe): r""" Given a DataPipe, updates its valid iterator ID and reset the DataPipe. """ if hasattr(datapipe, "_is_child_datapipe") and datapipe._is_child_datapipe is True: if hasattr(datapipe, "_set_main_datapipe_valid_iterator_id"): datapipe._set_main_datapipe_valid_iterator_id() # reset() is called within this method when appropriate else: raise RuntimeError("ChildDataPipe must have method `_set_main_datapipe_valid_iterator_id`.") else: if datapipe._valid_iterator_id is None: datapipe._valid_iterator_id = 0 else: datapipe._valid_iterator_id += 1 datapipe.reset() return datapipe._valid_iterator_id def hook_iterator(namespace, profile_name): r""" Hook that is applied to all `__iter__` of metaclass `_DataPipeMeta`. This is done for the purpose of profiling and checking if an iterator is still valid. """ def profiler_record_fn_context(): return torch.autograd.profiler.record_function(profile_name) class IteratorDecorator: """Wrap the iterator and modifying its `__next__` method""" def __init__(self, iterator, source_dp, iterator_id): self.iterator = iterator self.source_dp = source_dp self.iterator_id = iterator_id self._profiler_enabled = torch.autograd._profiler_enabled() def __iter__(self): return self def __next__(self): # TODO: Add try-except to in-place reduce traceback from the Exception # See: https://github.com/pytorch/data/issues/284 if self._profiler_enabled: with profiler_record_fn_context(): _check_iterator_valid(self.source_dp, self.iterator_id) return next(self.iterator) else: # Decided against using `contextlib.nullcontext` for performance reasons _check_iterator_valid(self.source_dp, self.iterator_id) return next(self.iterator) def __getattr__(self, name): return getattr(self.iterator, name) func = namespace['__iter__'] # ``__iter__`` of IterDataPipe is a generator function if inspect.isgeneratorfunction(func): @functools.wraps(func) def wrap_generator(*args, **kwargs): gen = func(*args, **kwargs) datapipe = args[0] iterator_id = _set_datapipe_valid_iterator_id(datapipe) # This ID is tied to each created iterator _profiler_enabled = torch.autograd._profiler_enabled() try: if _profiler_enabled: with profiler_record_fn_context(): response = gen.send(None) else: response = gen.send(None) while True: request = yield response # Pass through here every time `__next__` is called if _profiler_enabled: with profiler_record_fn_context(): _check_iterator_valid(datapipe, iterator_id) response = gen.send(request) else: # Decided against using `contextlib.nullcontext` for performance reasons _check_iterator_valid(datapipe, iterator_id) response = gen.send(request) except StopIteration as e: return e.value except Exception as e: # TODO: Simplify the traceback message to skip over `response = gen.send(None)` # Part of https://github.com/pytorch/data/issues/284 datapipe = args[0] msg = "thrown by __iter__ of" full_msg = f"{msg} {datapipe.__class__.__name__}({_generate_input_args_string(datapipe)})" if len(e.args) >= 1 and msg not in e.args[0]: e.args = (e.args[0] + f'\nThis exception is {full_msg}',) + e.args[1:] raise namespace['__iter__'] = wrap_generator else: # ``__iter__`` of IterDataPipe is NOT a generator function # IterDataPipe is an iterator with both ``__iter__`` and ``__next__`` # And ``__iter__`` may or may not return `self` if '__next__' in namespace: # If `__next__` exists, put a wrapper around it next_func = namespace['__next__'] @functools.wraps(next_func) def wrap_next(*args, **kwargs): if torch.autograd._profiler_enabled(): return next_func(*args, **kwargs) else: return next_func(*args, **kwargs) namespace['__next__'] = wrap_next # Note that if the `__next__` and `__iter__` do something completely unrelated? It may cause issue but # the user will be violating the iterator protocol # Regardless if `__next__` exists or not, `__iter__` needs a wrapper to track the number of valid iterators @functools.wraps(func) def wrap_iter(*args, **kwargs): iter_ret = func(*args, **kwargs) datapipe = args[0] iterator_id = _set_datapipe_valid_iterator_id(datapipe) # This ID is tied to each created iterator return IteratorDecorator(iter_ret, datapipe, iterator_id) namespace['__iter__'] = wrap_iter def _dp_init_subclass(sub_cls, *args, **kwargs): # Add function for datapipe instance to reinforce the type sub_cls.reinforce_type = reinforce_type # TODO: # - add global switch for type checking at compile-time # Ignore internal type class if getattr(sub_cls, '__type_class__', False): return # Check if the string type is valid if isinstance(sub_cls.type.param, ForwardRef): base_globals = sys.modules[sub_cls.__module__].__dict__ try: param = _eval_type(sub_cls.type.param, base_globals, locals()) sub_cls.type.param = param except TypeError as e: raise TypeError("{} is not supported by Python typing" .format(sub_cls.type.param.__forward_arg__)) from e if '__iter__' in sub_cls.__dict__: iter_fn = sub_cls.__dict__['__iter__'] hints = get_type_hints(iter_fn) if 'return' in hints: return_hint = hints['return'] # Plain Return Hint for Python 3.6 if return_hint == Iterator: return if not (hasattr(return_hint, '__origin__') and (return_hint.__origin__ == Iterator or return_hint.__origin__ == collections.abc.Iterator)): raise TypeError("Expected 'Iterator' as the return annotation for `__iter__` of {}" ", but found {}".format(sub_cls.__name__, _type_repr(hints['return']))) data_type = return_hint.__args__[0] if not issubtype(data_type, sub_cls.type.param): raise TypeError("Expected return type of '__iter__' as a subtype of {}, but found {}" " for {}".format(sub_cls.type, _type_repr(data_type), sub_cls.__name__)) def reinforce_type(self, expected_type): r""" Reinforce the type for DataPipe instance. And the 'expected_type' is required to be a subtype of the original type hint to restrict the type requirement of DataPipe instance. """ if isinstance(expected_type, tuple): expected_type = Tuple[expected_type] _type_check(expected_type, msg="'expected_type' must be a type") if not issubtype(expected_type, self.type.param): raise TypeError("Expected 'expected_type' as subtype of {}, but found {}" .format(self.type, _type_repr(expected_type))) self.type = _DataPipeType(expected_type) return self msadapter/ms_adapter/pytorch/utils/data/datapipes/dataframe/000077500000000000000000000000001436713554100246255ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/datapipes/dataframe/__init__.py000066400000000000000000000005511436713554100267370ustar00rootroot00000000000000from ms_adapter.pytorch.utils.data.datapipes.dataframe.dataframes import ( CaptureDataFrame, DFIterDataPipe, ) from ms_adapter.pytorch.utils.data.datapipes.dataframe.datapipes import ( DataFramesAsTuplesPipe, ) __all__ = ['CaptureDataFrame', 'DFIterDataPipe', 'DataFramesAsTuplesPipe'] # Please keep this list sorted assert __all__ == sorted(__all__) msadapter/ms_adapter/pytorch/utils/data/datapipes/dataframe/dataframe_wrapper.py000066400000000000000000000055031436713554100306660ustar00rootroot00000000000000_pandas = None _WITH_PANDAS = None def _try_import_pandas() -> bool: try: import pandas # type: ignore[import] global _pandas _pandas = pandas return True except ImportError: return False # pandas used only for prototyping, will be shortly replaced with TorchArrow def _with_pandas() -> bool: global _WITH_PANDAS if _WITH_PANDAS is None: _WITH_PANDAS = _try_import_pandas() return _WITH_PANDAS class PandasWrapper: @classmethod def create_dataframe(cls, data, columns): if not _with_pandas(): raise Exception("DataFrames prototype requires pandas to function") return _pandas.DataFrame(data, columns=columns) # type: ignore[union-attr] @classmethod def is_dataframe(cls, data): if not _with_pandas(): return False return isinstance(data, _pandas.core.frame.DataFrame) # type: ignore[union-attr] @classmethod def is_column(cls, data): if not _with_pandas(): return False return isinstance(data, _pandas.core.series.Series) # type: ignore[union-attr] @classmethod def iterate(cls, data): if not _with_pandas(): raise Exception("DataFrames prototype requires pandas to function") for d in data: yield d @classmethod def concat(cls, buffer): if not _with_pandas(): raise Exception("DataFrames prototype requires pandas to function") return _pandas.concat(buffer) # type: ignore[union-attr] @classmethod def get_item(cls, data, idx): if not _with_pandas(): raise Exception("DataFrames prototype requires pandas to function") return data[idx : idx + 1] @classmethod def get_len(cls, df): if not _with_pandas(): raise Exception("DataFrames prototype requires pandas to function") return len(df.index) # When you build own implementation just override it with dataframe_wrapper.set_df_wrapper(new_wrapper_class) default_wrapper = PandasWrapper def get_df_wrapper(): return default_wrapper def set_df_wrapper(wrapper): global default_wrapper default_wrapper = wrapper def create_dataframe(data, columns=None): wrapper = get_df_wrapper() return wrapper.create_dataframe(data, columns) def is_dataframe(data): wrapper = get_df_wrapper() return wrapper.is_dataframe(data) def is_column(data): wrapper = get_df_wrapper() return wrapper.is_column(data) def concat(buffer): wrapper = get_df_wrapper() return wrapper.concat(buffer) def iterate(data): wrapper = get_df_wrapper() return wrapper.iterate(data) def get_item(data, idx): wrapper = get_df_wrapper() return wrapper.get_item(data, idx) def get_len(df): wrapper = get_df_wrapper() return wrapper.get_len(df) msadapter/ms_adapter/pytorch/utils/data/datapipes/dataframe/dataframes.py000066400000000000000000000211731436713554100273120ustar00rootroot00000000000000from typing import Any, Dict, List from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import DFIterDataPipe, IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.dataframe.structures import DataChunkDF # TODO(VitalyFedyunin): Add error when two different traces get combined __all__ = [ "Capture", "CaptureAdd", "CaptureCall", "CaptureDataFrame", "CaptureDataFrameWithDataPipeOps", "CaptureF", "CaptureGetAttr", "CaptureGetItem", "CaptureInitial", "CaptureMul", "CaptureSetItem", "CaptureSub", "CaptureVariable", "CaptureVariableAssign", "DataFrameTracer", "DataFrameTracedOps", "get_val", ] class DataFrameTracedOps(DFIterDataPipe): def __init__(self, source_datapipe, output_var): self.source_datapipe = source_datapipe self.output_var = output_var def __iter__(self): for item in self.source_datapipe: yield self.output_var.apply_ops(item) # TODO(VitalyFedyunin): Extract this list from the DFIterDataPipe registred functions DATAPIPES_OPS = ['_dataframes_as_tuples', 'groupby', '_dataframes_filter', 'map', 'to_datapipe', 'shuffle', 'concat', 'batch', '_dataframes_per_row', '_dataframes_concat', '_dataframes_shuffle'] class Capture(object): # TODO: All operations are shared across entire InitialCapture, need to figure out what if we join two captures ctx: Dict[str, List[Any]] def __init__(self): self.ctx = {'operations': [], 'variables': []} def __str__(self): return self._ops_str() def _ops_str(self): res = "" for op in self.ctx['operations']: if len(res) > 0: res += "\n" res += str(op) return res def __getattr__(self, attrname): if attrname == 'kwarg': raise Exception('no kwargs!') return CaptureGetAttr(self, attrname, ctx=self.ctx) def __getitem__(self, key): return CaptureGetItem(self, key, ctx=self.ctx) def __setitem__(self, key, value): self.ctx['operations'].append( CaptureSetItem(self, key, value, ctx=self.ctx)) def __add__(self, add_val): res = CaptureAdd(self, add_val, ctx=self.ctx) var = CaptureVariable(res, ctx=self.ctx) self.ctx['operations'].append( CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)) return var def __sub__(self, add_val): res = CaptureSub(self, add_val, ctx=self.ctx) var = CaptureVariable(res, ctx=self.ctx) self.ctx['operations'].append( CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)) return var def __mul__(self, add_val): res = CaptureMul(self, add_val, ctx=self.ctx) var = CaptureVariable(res, ctx=self.ctx) t = CaptureVariableAssign(variable=var, value=res, ctx=self.ctx) self.ctx['operations'].append(t) return var class CaptureF(Capture): def __init__(self, ctx=None, **kwargs): if ctx is None: self.ctx = {'operations': [], 'variables': []} self.ctx = ctx self.kwargs = kwargs class CaptureCall(CaptureF): def __str__(self): return "{variable}({args},{kwargs})".format(**self.kwargs) def execute(self): return (get_val(self.kwargs['variable']))(*self.kwargs['args'], **self.kwargs['kwargs']) class CaptureVariableAssign(CaptureF): def __str__(self): return "{variable} = {value}".format(**self.kwargs) def execute(self): self.kwargs['variable'].calculated_value = self.kwargs['value'].execute() class CaptureVariable(Capture): value = None name = None calculated_value = None names_idx = 0 def __init__(self, value, ctx): self.ctx = ctx self.value = value self.name = 'var_%s' % CaptureVariable.names_idx CaptureVariable.names_idx += 1 self.ctx['variables'].append(self) def __str__(self): return self.name def execute(self): return self.calculated_value def apply_ops(self, dataframe): # TODO(VitalyFedyunin): Make this calculation thread safe (as currently it updates pointer) self.ctx['variables'][0].calculated_value = dataframe for op in self.ctx['operations']: op.execute() return self.calculated_value class CaptureGetItem(Capture): left: Capture key: Any def __init__(self, left, key, ctx): self.ctx = ctx self.left = left self.key = key def __str__(self): return "%s[%s]" % (self.left, get_val(self.key)) def execute(self): return (self.left.execute())[self.key] class CaptureSetItem(Capture): left: Capture key: Any value: Capture def __init__(self, left, key, value, ctx): self.ctx = ctx self.left = left self.key = key self.value = value def __str__(self): return "%s[%s] = %s" % (self.left, get_val(self.key), self.value) def execute(self): (self.left.execute())[ self.key] = self.value.execute() class CaptureAdd(Capture): left = None right = None def __init__(self, left, right, ctx): self.ctx = ctx self.left = left self.right = right def __str__(self): return "%s + %s" % (self.left, self.right) def execute(self): return get_val(self.left) + get_val(self.right) class CaptureMul(Capture): left = None right = None def __init__(self, left, right, ctx): self.ctx = ctx self.left = left self.right = right def __str__(self): return "%s * %s" % (self.left, self.right) def execute(self): return get_val(self.left) * get_val(self.right) class CaptureSub(Capture): left = None right = None def __init__(self, left, right, ctx): self.ctx = ctx self.left = left self.right = right def __str__(self): return "%s - %s" % (self.left, self.right) def execute(self): return get_val(self.left) - get_val(self.right) class CaptureGetAttr(Capture): source = None name: str def __init__(self, src, name, ctx): self.ctx = ctx self.src = src self.name = name def __str__(self): return "%s.%s" % (self.src, self.name) def execute(self): val = get_val(self.src) return getattr(val, self.name) def get_val(capture): if isinstance(capture, Capture): return capture.execute() elif isinstance(capture, str): return '"%s"' % capture else: return capture class CaptureInitial(CaptureVariable): def __init__(self): new_ctx: Dict[str, List[Any]] = {'operations': [], 'variables': []} super().__init__(None, new_ctx) self.name = 'input_%s' % self.name class CaptureDataFrame(CaptureInitial): pass class CaptureDataFrameWithDataPipeOps(CaptureDataFrame): def as_datapipe(self): return DataFrameTracedOps( self.ctx['variables'][0].source_datapipe, self) def raw_iterator(self): return self.as_datapipe().__iter__() def __iter__(self): return iter(self._dataframes_as_tuples()) def batch(self, batch_size=10, drop_last: bool = False, wrapper_class=DataChunkDF): dp = self._dataframes_per_row()._dataframes_concat(batch_size) dp = dp.as_datapipe().batch(1, drop_last=drop_last, wrapper_class=wrapper_class) dp._dp_contains_dataframe = True return dp def groupby(self, group_key_fn, *, buffer_size=10000, group_size=None, guaranteed_group_size=None, drop_remaining=False): dp = self._dataframes_per_row() dp = dp.as_datapipe().groupby(group_key_fn, buffer_size=buffer_size, group_size=group_size, guaranteed_group_size=guaranteed_group_size, drop_remaining=drop_remaining) return dp def shuffle(self, *args, **kwargs): return self._dataframes_shuffle(*args, **kwargs) def filter(self, *args, **kwargs): return self._dataframes_filter(*args, **kwargs) def __getattr__(self, attrname): # ? if attrname in DATAPIPES_OPS: return (self.as_datapipe()).__getattr__(attrname) return super().__getattr__(attrname) @functional_datapipe('trace_as_dataframe') class DataFrameTracer(CaptureDataFrameWithDataPipeOps, IterDataPipe): source_datapipe = None def __init__(self, source_datapipe): super().__init__() self.source_datapipe = source_datapipe msadapter/ms_adapter/pytorch/utils/data/datapipes/dataframe/datapipes.py000066400000000000000000000103771436713554100271610ustar00rootroot00000000000000import random from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import DFIterDataPipe, IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper __all__ = [ "ConcatDataFramesPipe", "DataFramesAsTuplesPipe", "ExampleAggregateAsDataFrames", "FilterDataFramesPipe", "PerRowDataFramesPipe", "ShuffleDataFramesPipe", ] @functional_datapipe('_dataframes_as_tuples') class DataFramesAsTuplesPipe(IterDataPipe): def __init__(self, source_datapipe): self.source_datapipe = source_datapipe def __iter__(self): for df in self.source_datapipe: for record in df.to_records(index=False): yield record @functional_datapipe('_dataframes_per_row', enable_df_api_tracing=True) class PerRowDataFramesPipe(DFIterDataPipe): def __init__(self, source_datapipe): self.source_datapipe = source_datapipe def __iter__(self): for df in self.source_datapipe: for i in range(len(df.index)): yield df[i:i + 1] @functional_datapipe('_dataframes_concat', enable_df_api_tracing=True) class ConcatDataFramesPipe(DFIterDataPipe): def __init__(self, source_datapipe, batch=3): self.source_datapipe = source_datapipe self.n_batch = batch def __iter__(self): buffer = [] for df in self.source_datapipe: buffer.append(df) if len(buffer) == self.n_batch: yield df_wrapper.concat(buffer) buffer = [] if len(buffer): yield df_wrapper.concat(buffer) @functional_datapipe('_dataframes_shuffle', enable_df_api_tracing=True) class ShuffleDataFramesPipe(DFIterDataPipe): def __init__(self, source_datapipe): self.source_datapipe = source_datapipe def __iter__(self): size = None all_buffer = [] for df in self.source_datapipe: if size is None: size = df_wrapper.get_len(df) for i in range(df_wrapper.get_len(df)): all_buffer.append(df_wrapper.get_item(df, i)) random.shuffle(all_buffer) buffer = [] for df in all_buffer: buffer.append(df) if len(buffer) == size: yield df_wrapper.concat(buffer) buffer = [] if len(buffer): yield df_wrapper.concat(buffer) @functional_datapipe('_dataframes_filter', enable_df_api_tracing=True) class FilterDataFramesPipe(DFIterDataPipe): def __init__(self, source_datapipe, filter_fn): self.source_datapipe = source_datapipe self.filter_fn = filter_fn def __iter__(self): size = None all_buffer = [] filter_res = [] for df in self.source_datapipe: if size is None: size = len(df.index) for i in range(len(df.index)): all_buffer.append(df[i:i + 1]) filter_res.append(self.filter_fn(df.iloc[i])) buffer = [] for df, res in zip(all_buffer, filter_res): if res: buffer.append(df) if len(buffer) == size: yield df_wrapper.concat(buffer) buffer = [] if len(buffer): yield df_wrapper.concat(buffer) @functional_datapipe('_to_dataframes_pipe', enable_df_api_tracing=True) class ExampleAggregateAsDataFrames(DFIterDataPipe): def __init__(self, source_datapipe, dataframe_size=10, columns=None): self.source_datapipe = source_datapipe self.columns = columns self.dataframe_size = dataframe_size def _as_list(self, item): try: return list(item) except Exception: # TODO(VitalyFedyunin): Replace with better iterable exception return [item] def __iter__(self): aggregate = [] for item in self.source_datapipe: aggregate.append(self._as_list(item)) if len(aggregate) == self.dataframe_size: yield df_wrapper.create_dataframe(aggregate, columns=self.columns) aggregate = [] if len(aggregate) > 0: yield df_wrapper.create_dataframe(aggregate, columns=self.columns) msadapter/ms_adapter/pytorch/utils/data/datapipes/dataframe/structures.py000066400000000000000000000010511436713554100274170ustar00rootroot00000000000000from ms_adapter.pytorch.utils.data.datapipes.datapipe import DataChunk __all__ = ["DataChunkDF", ] class DataChunkDF(DataChunk): """ DataChunkDF iterating over individual items inside of DataFrame containers, to access DataFrames user `raw_iterator` """ def __iter__(self): for df in self.items: for record in df.to_records(index=False): yield record def __len__(self): total_len = 0 for df in self.items: total_len += len(df) return total_len msadapter/ms_adapter/pytorch/utils/data/datapipes/datapipe.py000066400000000000000000000363651436713554100250570ustar00rootroot00000000000000import functools import pickle from typing import Dict, Callable, Optional, TypeVar, Generic, Iterator from ms_adapter.pytorch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta from ms_adapter.pytorch.utils.data.datapipes.utils.common import ( _deprecation_warning, _iter_deprecated_functional_names, _map_deprecated_functional_names, ) from ms_adapter.pytorch.utils.data.dataset import Dataset, IterableDataset try: import dill # XXX: By default, dill writes the Pickler dispatch table to inject its # own logic there. This globally affects the behavior of the standard library # pickler for any user who transitively depends on this module! # Undo this extension to avoid altering the behavior of the pickler globally. dill.extend(use_dill=False) HAS_DILL = True except ImportError: HAS_DILL = False __all__ = [ "DataChunk", "DFIterDataPipe", "IterDataPipe", "MapDataPipe", ] T = TypeVar('T') T_co = TypeVar('T_co', covariant=True) UNTRACABLE_DATAFRAME_PIPES = ['batch', # As it returns DataChunks 'groupby', # As it returns DataChunks '_dataframes_as_tuples', # As it unpacks DF 'trace_as_dataframe', # As it used to mark DF for tracing ] class IterDataPipe(IterableDataset[T_co], metaclass=_IterDataPipeMeta): r""" Iterable-style DataPipe. All DataPipes that represent an iterable of data samples should subclass this. This style of DataPipes is particularly useful when data come from a stream, or when the number of samples is too large to fit them all in memory. ``IterDataPipe`` is lazily initialized and its elements are computed only when ``next()`` is called on the iterator of an ``IterDataPipe``. All subclasses should overwrite :meth:`__iter__`, which would return an iterator of samples in this DataPipe. Calling ``__iter__`` of an ``IterDataPipe`` automatically invokes its method ``reset()``, which by default performs no operation. When writing a custom ``IterDataPipe``, users should override ``reset()`` if necessary. The common usages include resetting buffers, pointers, and various state variables within the custom ``IterDataPipe``. Note: Only `one` iterator can be valid for each ``IterDataPipe`` at a time, and the creation a second iterator will invalidate the first one. This constraint is necessary because some ``IterDataPipe`` have internal buffers, whose states can become invalid if there are multiple iterators. The code example below presents details on how this constraint looks in practice. If you have any feedback related to this constraint, please see `GitHub IterDataPipe Single Iterator Issue`_. These DataPipes can be invoked in two ways, using the class constructor or applying their functional form onto an existing ``IterDataPipe`` (recommended, available to most but not all DataPipes). You can chain multiple `IterDataPipe` together to form a pipeline that will perform multiple operations in succession. .. _GitHub IterDataPipe Single Iterator Issue: https://github.com/pytorch/data/issues/45 Note: When a subclass is used with :class:`~torch.utils.data.DataLoader`, each item in the DataPipe will be yielded from the :class:`~torch.utils.data.DataLoader` iterator. When :attr:`num_workers > 0`, each worker process will have a different copy of the DataPipe object, so it is often desired to configure each copy independently to avoid having duplicate data returned from the workers. :func:`~torch.utils.data.get_worker_info`, when called in a worker process, returns information about the worker. It can be used in either the dataset's :meth:`__iter__` method or the :class:`~torch.utils.data.DataLoader` 's :attr:`worker_init_fn` option to modify each copy's behavior. Examples: General Usage: >>> from torchdata.datapipes.iter import IterableWrapper, Mapper >>> dp = IterableWrapper(range(10)) >>> map_dp_1 = Mapper(dp, lambda x: x + 1) # Using class constructor >>> map_dp_2 = dp.map(lambda x: x + 1) # Using functional form (recommended) >>> list(map_dp_1) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] >>> list(map_dp_2) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] >>> filter_dp = map_dp_1.filter(lambda x: x % 2 == 0) >>> list(filter_dp) [2, 4, 6, 8, 10] Single Iterator Constraint Example: >>> from torchdata.datapipes.iter import IterableWrapper, Mapper >>> dp = IterableWrapper(range(10)) >>> it1 = iter(source_dp) >>> list(it1) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] >>> it1 = iter(source_dp) >>> it2 = iter(source_dp) # The creation of a new iterator invalidates `it1` >>> next(it2) 0 >>> next(it1) # Further usage of `it1` will raise a `RunTimeError` """ functions: Dict[str, Callable] = {} reduce_ex_hook: Optional[Callable] = None getstate_hook: Optional[Callable] = None str_hook: Optional[Callable] = None repr_hook: Optional[Callable] = None _valid_iterator_id: Optional[int] = None _restored: bool = False def __getattr__(self, attribute_name): if attribute_name in IterDataPipe.functions: if attribute_name in _iter_deprecated_functional_names: kwargs = _iter_deprecated_functional_names[attribute_name] _deprecation_warning(**kwargs) function = functools.partial(IterDataPipe.functions[attribute_name], self) return function else: raise AttributeError("'{0}' object has no attribute '{1}".format(self.__class__.__name__, attribute_name)) @classmethod def register_function(cls, function_name, function): cls.functions[function_name] = function @classmethod def register_datapipe_as_function(cls, function_name, cls_to_register, enable_df_api_tracing=False): if function_name in cls.functions: raise Exception("Unable to add DataPipe function name {} as it is already taken".format(function_name)) def class_function(cls, enable_df_api_tracing, source_dp, *args, **kwargs): result_pipe = cls(source_dp, *args, **kwargs) if isinstance(result_pipe, IterDataPipe): if enable_df_api_tracing or isinstance(source_dp, DFIterDataPipe): if function_name not in UNTRACABLE_DATAFRAME_PIPES: result_pipe = result_pipe.trace_as_dataframe() return result_pipe function = functools.partial(class_function, cls_to_register, enable_df_api_tracing) cls.functions[function_name] = function def __getstate__(self): """ This contains special logic to serialize `lambda` functions when `dill` is available. If this doesn't cover your custom DataPipe's use case, consider writing custom methods for `__getstate__` and `__setstate__`, or use `pickle.dumps` for serialization. """ if IterDataPipe.getstate_hook is not None: return IterDataPipe.getstate_hook(self) return self.__dict__ def __reduce_ex__(self, *args, **kwargs): if IterDataPipe.reduce_ex_hook is not None: try: return IterDataPipe.reduce_ex_hook(self) except NotImplementedError: pass return super().__reduce_ex__(*args, **kwargs) @classmethod def set_getstate_hook(cls, hook_fn): if IterDataPipe.getstate_hook is not None and hook_fn is not None: raise Exception("Attempt to override existing getstate_hook") IterDataPipe.getstate_hook = hook_fn @classmethod def set_reduce_ex_hook(cls, hook_fn): if IterDataPipe.reduce_ex_hook is not None and hook_fn is not None: raise Exception("Attempt to override existing reduce_ex_hook") IterDataPipe.reduce_ex_hook = hook_fn def __repr__(self): if self.repr_hook is not None: return self.repr_hook(self) # Instead of showing , return the class name return str(self.__class__.__qualname__) def __str__(self): if self.str_hook is not None: return self.str_hook(self) # Instead of showing , return the class name return str(self.__class__.__qualname__) def reset(self): r""" Reset the `IterDataPipe` to the initial state. By default, no-op. For subclasses of `IterDataPipe`, depending on their functionalities, they may want to override this method with implementations that may clear the buffers and reset pointers of the DataPipe. The `reset` method is always called when `__iter__` is called as part of `hook_iterator`. """ pass class DFIterDataPipe(IterDataPipe): def _is_dfpipe(self): return True class MapDataPipe(Dataset[T_co], metaclass=_DataPipeMeta): r""" Map-style DataPipe. All datasets that represent a map from keys to data samples should subclass this. Subclasses should overwrite :meth:`__getitem__`, supporting fetching a data sample for a given, unique key. Subclasses can also optionally overwrite :meth:`__len__`, which is expected to return the size of the dataset by many :class:`~torch.utils.data.Sampler` implementations and the default options of :class:`~torch.utils.data.DataLoader`. These DataPipes can be invoked in two ways, using the class constructor or applying their functional form onto an existing `MapDataPipe` (recommend, available to most but not all DataPipes). Note: :class:`~torch.utils.data.DataLoader` by default constructs an index sampler that yields integral indices. To make it work with a map-style DataPipe with non-integral indices/keys, a custom sampler must be provided. Example: >>> from torchdata.datapipes.map import SequenceWrapper, Mapper >>> dp = SequenceWrapper(range(10)) >>> map_dp_1 = dp.map(lambda x: x + 1) # Using functional form (recommended) >>> list(map_dp_1) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] >>> map_dp_2 = Mapper(dp, lambda x: x + 1) # Using class constructor >>> list(map_dp_2) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] >>> batch_dp = map_dp_1.batch(batch_size=2) >>> list(batch_dp) [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]] """ functions: Dict[str, Callable] = {} reduce_ex_hook: Optional[Callable] = None getstate_hook: Optional[Callable] = None str_hook: Optional[Callable] = None repr_hook: Optional[Callable] = None def __getattr__(self, attribute_name): if attribute_name in MapDataPipe.functions: if attribute_name in _map_deprecated_functional_names: kwargs = _map_deprecated_functional_names[attribute_name] _deprecation_warning(**kwargs) function = functools.partial(MapDataPipe.functions[attribute_name], self) return function else: raise AttributeError("'{0}' object has no attribute '{1}".format(self.__class__.__name__, attribute_name)) @classmethod def register_function(cls, function_name, function): cls.functions[function_name] = function @classmethod def register_datapipe_as_function(cls, function_name, cls_to_register): if function_name in cls.functions: raise Exception("Unable to add DataPipe function name {} as it is already taken".format(function_name)) def class_function(cls, source_dp, *args, **kwargs): result_pipe = cls(source_dp, *args, **kwargs) return result_pipe function = functools.partial(class_function, cls_to_register) cls.functions[function_name] = function def __getstate__(self): """ This contains special logic to serialize `lambda` functions when `dill` is available. If this doesn't cover your custom DataPipe's use case, consider writing custom methods for `__getstate__` and `__setstate__`, or use `pickle.dumps` for serialization. """ if MapDataPipe.getstate_hook is not None: return MapDataPipe.getstate_hook(self) return self.__dict__ def __reduce_ex__(self, *args, **kwargs): if MapDataPipe.reduce_ex_hook is not None: try: return MapDataPipe.reduce_ex_hook(self) except NotImplementedError: pass return super().__reduce_ex__(*args, **kwargs) @classmethod def set_getstate_hook(cls, hook_fn): if MapDataPipe.getstate_hook is not None and hook_fn is not None: raise Exception("Attempt to override existing getstate_hook") MapDataPipe.getstate_hook = hook_fn @classmethod def set_reduce_ex_hook(cls, hook_fn): if MapDataPipe.reduce_ex_hook is not None and hook_fn is not None: raise Exception("Attempt to override existing reduce_ex_hook") MapDataPipe.reduce_ex_hook = hook_fn def __repr__(self): if self.repr_hook is not None: return self.repr_hook(self) # Instead of showing , return the class name return str(self.__class__.__qualname__) def __str__(self): if self.str_hook is not None: return self.str_hook(self) # Instead of showing , return the class name return str(self.__class__.__qualname__) class _DataPipeSerializationWrapper: def __init__(self, datapipe): self._datapipe = datapipe def __getstate__(self): use_dill = False try: value = pickle.dumps(self._datapipe) except Exception: if HAS_DILL: value = dill.dumps(self._datapipe) use_dill = True else: raise return (value, use_dill) def __setstate__(self, state): value, use_dill = state if use_dill: self._datapipe = dill.loads(value) else: self._datapipe = pickle.loads(value) def __len__(self): try: return len(self._datapipe) except Exception: raise TypeError( "{} instance doesn't have valid length".format(type(self).__name__) ) class _IterDataPipeSerializationWrapper(_DataPipeSerializationWrapper, IterDataPipe): def __iter__(self): yield from self._datapipe class _MapDataPipeSerializationWrapper(_DataPipeSerializationWrapper, MapDataPipe): def __getitem__(self, idx): return self._datapipe[idx] class DataChunk(list, Generic[T]): def __init__(self, items): super().__init__(items) self.items = items def as_str(self, indent=''): res = indent + "[" + ", ".join(str(i) for i in iter(self)) + "]" return res def __iter__(self) -> Iterator[T]: for i in super().__iter__(): yield i def raw_iterator(self) -> T: # type: ignore[misc] for i in self.items: yield i msadapter/ms_adapter/pytorch/utils/data/datapipes/gen_pyi.py000066400000000000000000000233031436713554100247060ustar00rootroot00000000000000import os import pathlib from typing import Any, Dict, List, Set, Tuple, Union def materialize_lines(lines: List[str], indentation: int) -> str: output = "" new_line_with_indent = "\n" + " " * indentation for i, line in enumerate(lines): if i != 0: output += new_line_with_indent output += line.replace('\n', new_line_with_indent) return output def gen_from_template(dir: str, template_name: str, output_name: str, replacements: List[Tuple[str, Any, int]]): template_path = os.path.join(dir, template_name) output_path = os.path.join(dir, output_name) with open(template_path, "r") as f: content = f.read() for placeholder, lines, indentation in replacements: with open(output_path, "w") as f: content = content.replace(placeholder, materialize_lines(lines, indentation)) f.write(content) def find_file_paths(dir_paths: List[str], files_to_exclude: Set[str]) -> Set[str]: """ When given a path to a directory, returns the paths to the relevant files within it. This function does NOT recursive traverse to subdirectories. """ paths: Set[str] = set() for dir_path in dir_paths: all_files = os.listdir(dir_path) python_files = {fname for fname in all_files if ".py" == fname[-3:]} filter_files = {fname for fname in python_files if fname not in files_to_exclude} paths.update({os.path.join(dir_path, fname) for fname in filter_files}) return paths def extract_method_name(line: str) -> str: """ Extracts method name from decorator in the form of "@functional_datapipe({method_name})" """ if "(\"" in line: start_token, end_token = "(\"", "\")" elif "(\'" in line: start_token, end_token = "(\'", "\')" else: raise RuntimeError(f"Unable to find appropriate method name within line:\n{line}") start, end = line.find(start_token) + len(start_token), line.find(end_token) return line[start:end] def extract_class_name(line: str) -> str: """ Extracts class name from class definition in the form of "class {CLASS_NAME}({Type}):" """ start_token = "class " end_token = "(" start, end = line.find(start_token) + len(start_token), line.find(end_token) return line[start:end] def parse_datapipe_file(file_path: str) -> Tuple[Dict[str, str], Dict[str, str], Set[str]]: """ Given a path to file, parses the file and returns a dictionary of method names to function signatures. """ method_to_signature, method_to_class_name, special_output_type = {}, {}, set() with open(file_path) as f: open_paren_count = 0 method_name, class_name, signature = "", "", "" skip = False for line in f.readlines(): if line.count("\"\"\"") % 2 == 1: skip = not skip if skip or "\"\"\"" in line: # Skipping comment/example blocks continue if "@functional_datapipe" in line: method_name = extract_method_name(line) continue if method_name and "class " in line: class_name = extract_class_name(line) continue if method_name and ("def __init__(" in line or "def __new__(" in line): if "def __new__(" in line: special_output_type.add(method_name) open_paren_count += 1 start = line.find("(") + len("(") line = line[start:] if open_paren_count > 0: open_paren_count += line.count('(') open_paren_count -= line.count(')') if open_paren_count == 0: end = line.rfind(')') signature += line[:end] method_to_signature[method_name] = process_signature(signature) method_to_class_name[method_name] = class_name method_name, class_name, signature = "", "", "" elif open_paren_count < 0: raise RuntimeError("open parenthesis count < 0. This shouldn't be possible.") else: signature += line.strip('\n').strip(' ') return method_to_signature, method_to_class_name, special_output_type def parse_datapipe_files(file_paths: Set[str]) -> Tuple[Dict[str, str], Dict[str, str], Set[str]]: methods_and_signatures, methods_and_class_names, methods_with_special_output_types = {}, {}, set() for path in file_paths: method_to_signature, method_to_class_name, methods_needing_special_output_types = parse_datapipe_file(path) methods_and_signatures.update(method_to_signature) methods_and_class_names.update(method_to_class_name) methods_with_special_output_types.update(methods_needing_special_output_types) return methods_and_signatures, methods_and_class_names, methods_with_special_output_types def split_outside_bracket(line: str, delimiter: str = ",") -> List[str]: """ Given a line of text, split it on comma unless the comma is within a bracket '[]'. """ bracket_count = 0 curr_token = "" res = [] for char in line: if char == "[": bracket_count += 1 elif char == "]": bracket_count -= 1 elif char == delimiter and bracket_count == 0: res.append(curr_token) curr_token = "" continue curr_token += char res.append(curr_token) return res def process_signature(line: str) -> str: """ Given a raw function signature, clean it up by removing the self-referential datapipe argument, default arguments of input functions, newlines, and spaces. """ tokens: List[str] = split_outside_bracket(line) for i, token in enumerate(tokens): tokens[i] = token.strip(' ') if token == "cls": tokens[i] = "self" elif i > 0 and ("self" == tokens[i - 1]) and (tokens[i][0] != "*"): # Remove the datapipe after 'self' or 'cls' unless it has '*' tokens[i] = "" elif "Callable =" in token: # Remove default argument if it is a function head, default_arg = token.rsplit("=", 2) tokens[i] = head.strip(' ') + "= ..." tokens = [t for t in tokens if t != ""] line = ', '.join(tokens) return line def get_method_definitions(file_path: Union[str, List[str]], files_to_exclude: Set[str], deprecated_files: Set[str], default_output_type: str, method_to_special_output_type: Dict[str, str], root: str = "") -> List[str]: """ .pyi generation for functional DataPipes Process # 1. Find files that we want to process (exclude the ones who don't) # 2. Parse method name and signature # 3. Remove first argument after self (unless it is "*datapipes"), default args, and spaces """ if root == "": root = str(pathlib.Path(__file__).parent.resolve()) file_path = [file_path] if isinstance(file_path, str) else file_path file_path = [os.path.join(root, path) for path in file_path] file_paths = find_file_paths(file_path, files_to_exclude=files_to_exclude.union(deprecated_files)) methods_and_signatures, methods_and_class_names, methods_w_special_output_types = \ parse_datapipe_files(file_paths) method_definitions = [] for method_name, arguments in methods_and_signatures.items(): class_name = methods_and_class_names[method_name] if method_name in methods_w_special_output_types: output_type = method_to_special_output_type[method_name] else: output_type = default_output_type method_definitions.append(f"# Functional form of '{class_name}'\n" f"def {method_name}({arguments}) -> {output_type}: ...") method_definitions.sort(key=lambda s: s.split('\n')[1]) # sorting based on method_name return method_definitions # Defined outside of main() so they can be imported by TorchData iterDP_file_path: str = "iter" iterDP_files_to_exclude: Set[str] = {"__init__.py", "utils.py"} iterDP_deprecated_files: Set[str] = set() iterDP_method_to_special_output_type: Dict[str, str] = {"demux": "List[IterDataPipe]", "fork": "List[IterDataPipe]"} mapDP_file_path: str = "map" mapDP_files_to_exclude: Set[str] = {"__init__.py", "utils.py"} mapDP_deprecated_files: Set[str] = set() mapDP_method_to_special_output_type: Dict[str, str] = {} def main() -> None: """ # Inject file into template datapipe.pyi.in TODO: The current implementation of this script only generates interfaces for built-in methods. To generate interface for user-defined DataPipes, consider changing `IterDataPipe.register_datapipe_as_function`. """ iter_method_definitions = get_method_definitions(iterDP_file_path, iterDP_files_to_exclude, iterDP_deprecated_files, "IterDataPipe", iterDP_method_to_special_output_type) map_method_definitions = get_method_definitions(mapDP_file_path, mapDP_files_to_exclude, mapDP_deprecated_files, "MapDataPipe", mapDP_method_to_special_output_type) path = pathlib.Path(__file__).parent.resolve() replacements = [('${IterDataPipeMethods}', iter_method_definitions, 4), ('${MapDataPipeMethods}', map_method_definitions, 4)] gen_from_template(dir=str(path), template_name="datapipe.pyi.in", output_name="datapipe.pyi", replacements=replacements) if __name__ == '__main__': print("Generating Python interface file 'datapipe.pyi'...") main() msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/000077500000000000000000000000001436713554100236445ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/__init__.py000066400000000000000000000040421436713554100257550ustar00rootroot00000000000000from ms_adapter.pytorch.utils.data.datapipes.iter.utils import ( IterableWrapperIterDataPipe as IterableWrapper, ) from ms_adapter.pytorch.utils.data.datapipes.iter.callable import ( CollatorIterDataPipe as Collator, MapperIterDataPipe as Mapper, ) from ms_adapter.pytorch.utils.data.datapipes.iter.combinatorics import ( SamplerIterDataPipe as Sampler, ShufflerIterDataPipe as Shuffler, ) from ms_adapter.pytorch.utils.data.datapipes.iter.combining import ( ConcaterIterDataPipe as Concater, DemultiplexerIterDataPipe as Demultiplexer, ForkerIterDataPipe as Forker, MultiplexerIterDataPipe as Multiplexer, ZipperIterDataPipe as Zipper, ) from ms_adapter.pytorch.utils.data.datapipes.iter.filelister import ( FileListerIterDataPipe as FileLister, ) from ms_adapter.pytorch.utils.data.datapipes.iter.fileopener import ( FileLoaderIterDataPipe as FileLoader, FileOpenerIterDataPipe as FileOpener, ) from ms_adapter.pytorch.utils.data.datapipes.iter.grouping import ( BatcherIterDataPipe as Batcher, GrouperIterDataPipe as Grouper, ShardingFilterIterDataPipe as ShardingFilter, UnBatcherIterDataPipe as UnBatcher, ) from ms_adapter.pytorch.utils.data.datapipes.iter.routeddecoder import ( RoutedDecoderIterDataPipe as RoutedDecoder, ) from ms_adapter.pytorch.utils.data.datapipes.iter.selecting import ( FilterIterDataPipe as Filter, ) from ms_adapter.pytorch.utils.data.datapipes.iter.streamreader import ( StreamReaderIterDataPipe as StreamReader, ) __all__ = ['Batcher', 'Collator', 'Concater', 'Demultiplexer', 'FileLister', 'FileLoader', 'FileOpener', 'Filter', 'Forker', 'Grouper', 'IterableWrapper', 'Mapper', 'Multiplexer', 'RoutedDecoder', 'Sampler', 'ShardingFilter', 'Shuffler', 'StreamReader', 'UnBatcher', 'Zipper'] # Please keep this list sorted assert __all__ == sorted(__all__) msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/callable.py000066400000000000000000000145131436713554100257610ustar00rootroot00000000000000from typing import Callable, Iterator, Sized, TypeVar from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data._utils.collate import default_collate from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.utils.common import _check_lambda_fn __all__ = [ "CollatorIterDataPipe", "MapperIterDataPipe", ] T_co = TypeVar("T_co", covariant=True) @functional_datapipe("map") class MapperIterDataPipe(IterDataPipe[T_co]): r""" Applies a function over each item from the source DataPipe (functional name: ``map``). The function can be any regular Python function or partial object. Lambda function is not recommended as it is not supported by pickle. Args: datapipe: Source Iterable DataPipe fn: Function being applied over each item input_col: Index or indices of data which ``fn`` is applied, such as: - ``None`` as default to apply ``fn`` to the data directly. - Integer(s) is used for list/tuple. - Key(s) is used for dict. output_col: Index of data where result of ``fn`` is placed. ``output_col`` can be specified only when ``input_col`` is not ``None`` - ``None`` as default to replace the index that ``input_col`` specified; For ``input_col`` with multiple indices, the left-most one is used, and other indices will be removed. - Integer is used for list/tuple. ``-1`` represents to append result at the end. - Key is used for dict. New key is acceptable. Example: >>> from torchdata.datapipes.iter import IterableWrapper, Mapper >>> def add_one(x): ... return x + 1 >>> dp = IterableWrapper(range(10)) >>> map_dp_1 = dp.map(add_one) # Invocation via functional form is preferred >>> list(map_dp_1) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] >>> # We discourage the usage of `lambda` functions as they are not serializable with `pickle` >>> # Use `functools.partial` or explicitly define the function instead >>> map_dp_2 = Mapper(dp, lambda x: x + 1) >>> list(map_dp_2) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] """ datapipe: IterDataPipe fn: Callable def __init__( self, datapipe: IterDataPipe, fn: Callable, input_col=None, output_col=None, ) -> None: super().__init__() self.datapipe = datapipe _check_lambda_fn(fn) self.fn = fn # type: ignore[assignment] self.input_col = input_col if input_col is None and output_col is not None: raise ValueError("`output_col` must be None when `input_col` is None.") if isinstance(output_col, (list, tuple)): if len(output_col) > 1: raise ValueError("`output_col` must be a single-element list or tuple") output_col = output_col[0] self.output_col = output_col def _apply_fn(self, data): if self.input_col is None and self.output_col is None: return self.fn(data) if self.input_col is None: res = self.fn(data) elif isinstance(self.input_col, (list, tuple)): args = tuple(data[col] for col in self.input_col) res = self.fn(*args) else: res = self.fn(data[self.input_col]) # Copy tuple to list and run in-place modification because tuple is immutable. if isinstance(data, tuple): t_flag = True data = list(data) else: t_flag = False if self.output_col is None: if isinstance(self.input_col, (list, tuple)): data[self.input_col[0]] = res for idx in sorted(self.input_col[1:], reverse=True): del data[idx] else: data[self.input_col] = res else: if self.output_col == -1: data.append(res) else: data[self.output_col] = res # Convert list back to tuple return tuple(data) if t_flag else data def __iter__(self) -> Iterator[T_co]: for data in self.datapipe: yield self._apply_fn(data) def __len__(self) -> int: if isinstance(self.datapipe, Sized): return len(self.datapipe) raise TypeError( "{} instance doesn't have valid length".format(type(self).__name__) ) @functional_datapipe("collate") class CollatorIterDataPipe(MapperIterDataPipe): r""" Collates samples from DataPipe to Tensor(s) by a custom collate function (functional name: ``collate``). By default, it uses :func:`ms_adapter.pytorch.utils.data.default_collate`. .. note:: While writing a custom collate function, you can import :func:`ms_adapter.pytorch.utils.data.default_collate` for the default behavior and `functools.partial` to specify any additional arguments. Args: datapipe: Iterable DataPipe being collated collate_fn: Customized collate function to collect and combine data or a batch of data. Default function collates to Tensor(s) based on data type. Example: Convert integer data to float Tensor >>> class MyIterDataPipe(ms_adapter.pytorch.utils.data.IterDataPipe): ... def __init__(self, start, end): ... super(MyIterDataPipe).__init__() ... assert end > start, "this example code only works with end >= start" ... self.start = start ... self.end = end ... ... def __iter__(self): ... return iter(range(self.start, self.end)) ... ... def __len__(self): ... return self.end - self.start ... >>> ds = MyIterDataPipe(start=3, end=7) >>> print(list(ds)) [3, 4, 5, 6] >>> def collate_fn(batch): ... return torch.tensor(batch, dtype=torch.float) ... >>> collated_ds = CollateIterDataPipe(ds, collate_fn=collate_fn) >>> print(list(collated_ds)) [tensor(3.), tensor(4.), tensor(5.), tensor(6.)] """ def __init__( self, datapipe: IterDataPipe, collate_fn: Callable = default_collate, ) -> None: super().__init__(datapipe, fn=collate_fn) msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/combinatorics.py000066400000000000000000000142711436713554100270570ustar00rootroot00000000000000import random import torch from ms_adapter.pytorch.utils.data import Sampler, SequentialSampler from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe from typing import Dict, Iterator, List, Optional, Sized, Tuple, Type, TypeVar __all__ = [ "SamplerIterDataPipe", "ShufflerIterDataPipe", ] T_co = TypeVar('T_co', covariant=True) class SamplerIterDataPipe(IterDataPipe[T_co]): r""" Generates sample elements using the provided ``Sampler`` (defaults to :class:`SequentialSampler`). Args: datapipe: IterDataPipe to sample from sampler: Sampler class to generate sample elements from input DataPipe. Default is :class:`SequentialSampler` for IterDataPipe """ datapipe: IterDataPipe sampler: Sampler def __init__(self, datapipe: IterDataPipe, sampler: Type[Sampler] = SequentialSampler, sampler_args: Optional[Tuple] = None, sampler_kwargs: Optional[Dict] = None ) -> None: assert isinstance(datapipe, Sized), \ "Sampler class requires input datapipe implemented `__len__`" super().__init__() self.datapipe = datapipe self.sampler_args = () if sampler_args is None else sampler_args self.sampler_kwargs = {} if sampler_kwargs is None else sampler_kwargs # https://github.com/python/mypy/pull/9629 will solve self.sampler = sampler(data_source=self.datapipe, *self.sampler_args, **self.sampler_kwargs) # type: ignore[misc] def __iter__(self) -> Iterator[T_co]: return iter(self.sampler) def __len__(self) -> int: # Dataset has been tested as `Sized` if isinstance(self.sampler, Sized) and len(self.sampler) >= 0: return len(self.sampler) raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) @functional_datapipe('shuffle') class ShufflerIterDataPipe(IterDataPipe[T_co]): r""" Shuffles the input DataPipe with a buffer (functional name: ``shuffle``). The buffer with ``buffer_size`` is filled with elements from the datapipe first. Then, each item will be yielded from the buffer by reservoir sampling via iterator. ``buffer_size`` is required to be larger than ``0``. For ``buffer_size == 1``, the datapipe is not shuffled. In order to fully shuffle all elements from datapipe, ``buffer_size`` is required to be greater than or equal to the size of datapipe. When it is used with :class:`ms_adapter.pytorch.utils.data.DataLoader`, the methods to set up random seed are different based on :attr:`num_workers`. For single-process mode (:attr:`num_workers == 0`), the random seed is set before the :class:`~ms_adapter.pytorch.utils.data.DataLoader` in the main process. For multi-process mode (:attr:`num_worker > 0`), `worker_init_fn` is used to set up a random seed for each worker process. Args: datapipe: The IterDataPipe being shuffled buffer_size: The buffer size for shuffling (default to ``10000``) unbatch_level: Specifies if it is necessary to unbatch source data before applying the shuffle Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> dp = IterableWrapper(range(10)) >>> shuffle_dp = dp.shuffle() >>> list(shuffle_dp) [0, 4, 1, 6, 3, 2, 9, 5, 7, 8] """ datapipe: IterDataPipe[T_co] buffer_size: int _buffer: List[T_co] _enabled: bool _seed: Optional[int] _rng: random.Random def __init__(self, datapipe: IterDataPipe[T_co], *, buffer_size: int = 10000, unbatch_level: int = 0 ) -> None: super().__init__() # TODO: Performance optimization # buffer can be a fixed size and remove expensive `append()` and `len()` operations self._buffer: List[T_co] = [] assert buffer_size > 0, "buffer_size should be larger than 0" if unbatch_level == 0: self.datapipe = datapipe else: self.datapipe = datapipe.unbatch(unbatch_level=unbatch_level) self.buffer_size = buffer_size self._enabled = True self._seed = None self._rng = random.Random() def set_shuffle(self, shuffle=True): self._enabled = shuffle return self def set_seed(self, seed: int): self._seed = seed def __iter__(self) -> Iterator[T_co]: if not self._enabled: for x in self.datapipe: yield x else: self._rng.seed(self._seed) self._seed = None for x in self.datapipe: if len(self._buffer) == self.buffer_size: idx = self._rng.randint(0, len(self._buffer) - 1) val, self._buffer[idx] = self._buffer[idx], x yield val else: self._buffer.append(x) self._rng.shuffle(self._buffer) while self._buffer: yield self._buffer.pop() def __len__(self) -> int: if isinstance(self.datapipe, Sized): return len(self.datapipe) raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) def reset(self) -> None: self._buffer = [] if self._enabled and self._seed is None: self._seed = int(torch.empty((), dtype=torch.int64).random_().item()) def __getstate__(self): if IterDataPipe.getstate_hook is not None: return IterDataPipe.getstate_hook(self) state = ( self.datapipe, self.buffer_size, self._enabled, self._seed, self._rng.getstate(), ) return state def __setstate__(self, state): ( self.datapipe, self.buffer_size, self._enabled, self._seed, rng_state, ) = state self._rng = random.Random() self._rng.setstate(rng_state) self._buffer = [] def __del__(self): self._buffer.clear() msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/combining.py000066400000000000000000000534351436713554100261750ustar00rootroot00000000000000import warnings from collections import deque from typing import Any, Callable, Iterator, List, Optional, Sized, Tuple, TypeVar, Deque from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.utils.common import _check_lambda_fn __all__ = [ "ConcaterIterDataPipe", "DemultiplexerIterDataPipe", "ForkerIterDataPipe", "MultiplexerIterDataPipe", "ZipperIterDataPipe", ] T_co = TypeVar('T_co', covariant=True) @functional_datapipe('concat') class ConcaterIterDataPipe(IterDataPipe): r""" Concatenates multiple Iterable DataPipes (functional name: ``concat``). The resulting DataPipe will yield all the elements from the first input DataPipe, before yielding from the subsequent ones. Args: datapipes: Iterable DataPipes being concatenated Example: >>> import random >>> from torchdata.datapipes.iter import IterableWrapper >>> dp1 = IterableWrapper(range(3)) >>> dp2 = IterableWrapper(range(5)) >>> list(dp1.concat(dp2)) [0, 1, 2, 0, 1, 2, 3, 4] """ datapipes: Tuple[IterDataPipe] length: Optional[int] def __init__(self, *datapipes: IterDataPipe): if len(datapipes) == 0: raise ValueError("Expected at least one DataPipe, but got nothing") if not all(isinstance(dp, IterDataPipe) for dp in datapipes): raise TypeError("Expected all inputs to be `IterDataPipe`") self.datapipes = datapipes # type: ignore[assignment] self.length = None def __iter__(self) -> Iterator: for dp in self.datapipes: for data in dp: yield data def __len__(self) -> int: if self.length is not None: if self.length == -1: raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) return self.length if all(isinstance(dp, Sized) for dp in self.datapipes): self.length = sum(len(dp) for dp in self.datapipes) else: self.length = -1 return len(self) @functional_datapipe('fork') class ForkerIterDataPipe(IterDataPipe): r""" Creates multiple instances of the same Iterable DataPipe (functional name: ``fork``). Args: datapipe: Iterable DataPipe being copied num_instances: number of instances of the datapipe to create buffer_size: this restricts how far ahead the leading child DataPipe can read relative to the slowest child DataPipe. Defaults to ``1000``. Use ``-1`` for the unlimited buffer. Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> source_dp = IterableWrapper(range(5)) >>> dp1, dp2 = source_dp.fork(num_instances=2) >>> list(dp1) [0, 1, 2, 3, 4] >>> list(dp2) [0, 1, 2, 3, 4] """ def __new__(cls, datapipe: IterDataPipe, num_instances: int, buffer_size: int = 1000): if num_instances < 1: raise ValueError(f"Expected `num_instaces` larger than 0, but {num_instances} is found") if num_instances == 1: return datapipe container = _ForkerIterDataPipe(datapipe, num_instances, buffer_size) return [_ChildDataPipe(container, i) for i in range(num_instances)] class _ForkerIterDataPipe(IterDataPipe): r""" Container to hold instance-specific information on behalf of ForkerIterDataPipe. It tracks the state of its child DataPipes, maintains the buffer, and yields the next value as requested by the child DataPipes. """ def __init__(self, datapipe: IterDataPipe, num_instances: int, buffer_size: int = 1000): self.main_datapipe = datapipe self._datapipe_iterator: Optional[Iterator[Any]] = None self.num_instances = num_instances self.buffer: Deque = deque() self.buffer_size = buffer_size if self.buffer_size < 0: warnings.warn( "Unlimited buffer size is set for `fork`, " "please be aware of OOM at random places", UserWarning ) self.child_pointers: List[int] = [0] * num_instances # Indicate the indices of the next element to get self.slowest_ptr = 0 self.leading_ptr = 0 self.end_ptr: Optional[int] = None def __len__(self): return len(self.main_datapipe) def get_next_element_by_instance(self, instance_id: int): if self._datapipe_iterator is None: self._datapipe_iterator = iter(self.main_datapipe) while self.end_ptr is None or self.child_pointers[instance_id] < self.end_ptr: if not self.buffer or self.child_pointers[instance_id] > self.leading_ptr: self.leading_ptr = self.child_pointers[instance_id] if self.buffer_size >= 0 and self.leading_ptr - self.slowest_ptr + 1 > self.buffer_size: raise BufferError("ForkerIterDataPipe buffer overflow," + f"buffer size {self.buffer_size} is insufficient.") try: self.buffer.append(next(self._datapipe_iterator)) self.child_pointers[instance_id] += 1 yield self.buffer[-1] except StopIteration: self.end_ptr = self.leading_ptr else: # Child pointer is slower than or equal to the leading_ptr buffer_index = self.child_pointers[instance_id] - self.slowest_ptr return_val = self.buffer[buffer_index] self.child_pointers[instance_id] += 1 if self.child_pointers[instance_id] - 1 == self.slowest_ptr: new_min = min(self.child_pointers) # Can optimize by avoiding the call to min() if self.slowest_ptr < new_min: self.slowest_ptr = new_min self.buffer.popleft() yield return_val if self.end_ptr and self.child_pointers[instance_id] == self.end_ptr and\ all(p == self.end_ptr for p in self.child_pointers): self._datapipe_iterator = None def is_every_instance_exhausted(self) -> bool: return all(self.end_ptr == ptr for ptr in self.child_pointers) def reset(self) -> None: self._datapipe_iterator = iter(self.main_datapipe) self.buffer = deque() self.child_pointers = [0] * self.num_instances self.slowest_ptr = 0 self.leading_ptr = 0 self.end_ptr = None def __getstate__(self): if IterDataPipe.getstate_hook is not None: return IterDataPipe.getstate_hook(self) state = ( self.main_datapipe, self.num_instances, self.buffer_size, ) return state def __setstate__(self, state): ( self.main_datapipe, self.num_instances, self.buffer_size, ) = state self._datapipe_iterator = None self.buffer = deque() self.child_pointers = [0] * self.num_instances self.slowest_ptr = 0 self.leading_ptr = 0 self.end_ptr = None def __del__(self): self.buffer.clear() class _ChildDataPipe(IterDataPipe): r""" Iterable Datapipe that is a child of a main DataPipe. The instance of this class will pass its instance_id to get the next value from its main DataPipe. Note: ChildDataPipe, like all other IterDataPipe, follows the single iterator per IterDataPipe constraint. Since ChildDataPipes share a common buffer, when an iterator is created for one of the ChildDataPipes, the previous iterators for all ChildDataPipes must be invalidated, with the exception when a ChildDataPipe hasn't had an iterator created from it since the last invalidation. See the example below. Singler Iterator per IteraDataPipe Invalidation Example: >>> source_dp = IterableWrapper(range(10)) >>> cdp1, cdp2 = source_dp.fork(num_instances=2) >>> it1, it2 = iter(cdp1), iter(cdp2) >>> it3 = iter(cdp1) The line above invalidates `it1` and `it2`, and resets `ForkerIterDataPipe`. >>> it4 = iter(cdp2) The line above doesn't invalidate `it3`, because an iterator for `cdp2` hasn't been created since the last invalidation. Args: main_datapipe: Main DataPipe with a method 'get_next_element_by_instance(instance_id)' instance_id: integer identifier of this instance """ _is_child_datapipe: bool = True def __init__(self, main_datapipe: IterDataPipe, instance_id: int): required_attrs = ["get_next_element_by_instance", "is_every_instance_exhausted", "reset"] required_ops = [getattr(main_datapipe, attr) for attr in required_attrs] if any(not callable(op) for op in required_ops): raise NotImplementedError(f"Main Datapipe must have methods {required_attrs} implemented.") self.main_datapipe: IterDataPipe = main_datapipe self.instance_id = instance_id def __iter__(self): # Note that the logic behind setting iterator ID and `reset` are handled within `hook_iterator` # We want to separate the code for reset and yield, so that 'reset' executes before __next__ is called return self.main_datapipe.get_next_element_by_instance(self.instance_id) def __len__(self): return len(self.main_datapipe) # This method is called by `hook_iterator` in `_typing.py`. def _set_main_datapipe_valid_iterator_id(self) -> int: r""" Update the valid iterator ID for both this DataPipe object and `main_datapipe`. `main_datapipe.reset()` is called when the ID is incremented to a new generation. """ # 1. First time any child iterator is created if self.main_datapipe._valid_iterator_id is None: self.main_datapipe._valid_iterator_id = 0 # type: ignore[attr-defined] # 2. This instance was already in the same generation as `main_datapipe`, # we need to increment the ID further by 1 elif self.main_datapipe._valid_iterator_id == self._valid_iterator_id: # type: ignore[has-type] self.main_datapipe._valid_iterator_id += 1 # type: ignore[attr-defined] # Whenever a new generation of iterator is created, the `main_datapipe` must reset if not self.main_datapipe.is_every_instance_exhausted(): warnings.warn("Some child DataPipes are not exhausted when __iter__ is called. We are resetting " "the buffer and each child DataPipe will read from the start again.", UserWarning) self.main_datapipe.reset() # 3. Otherwise, the iterator is behind the others, so it will just need to catch up by setting # the instance's iterator to match that of `main_datapipe` self._valid_iterator_id = self.main_datapipe._valid_iterator_id return self._valid_iterator_id # This method is called by `hook_iterator` in `_typing.py`. def _check_valid_iterator_id(self, iterator_id) -> bool: r""" Check the valid iterator ID against that of DataPipe object and that of `main_datapipe`. """ return iterator_id == self._valid_iterator_id and iterator_id == self.main_datapipe._valid_iterator_id @functional_datapipe('demux') class DemultiplexerIterDataPipe(IterDataPipe): r""" Splits the input DataPipe into multiple child DataPipes, using the given classification function (functional name: ``demux``). A list of the child DataPipes is returned from this operation. Args: datapipe: Iterable DataPipe being filtered num_instances: number of instances of the DataPipe to create classifier_fn: a function that maps values to an integer within the range ``[0, num_instances - 1]`` or ``None`` drop_none: defaults to ``False``, if ``True``, the function will skip over elements classified as ``None`` buffer_size: this defines the maximum number of inputs that the buffer can hold across all child DataPipes while waiting for their values to be yielded. Defaults to ``1000``. Use ``-1`` for the unlimited buffer. Examples: >>> from torchdata.datapipes.iter import IterableWrapper >>> def odd_or_even(n): ... return n % 2 >>> source_dp = IterableWrapper(range(5)) >>> dp1, dp2 = source_dp.demux(num_instances=2, classifier_fn=odd_or_even) >>> list(dp1) [0, 2, 4] >>> list(dp2) [1, 3] >>> # It can also filter out any element that gets `None` from the `classifier_fn` >>> def odd_or_even_no_zero(n): ... return n % 2 if n != 0 else None >>> dp1, dp2 = source_dp.demux(num_instances=2, classifier_fn=odd_or_even_no_zero, drop_none=True) >>> list(dp1) [2, 4] >>> list(dp2) [1, 3] """ def __new__(cls, datapipe: IterDataPipe, num_instances: int, classifier_fn: Callable[[T_co], Optional[int]], drop_none: bool = False, buffer_size: int = 1000): if num_instances < 1: raise ValueError(f"Expected `num_instaces` larger than 0, but {num_instances} is found") _check_lambda_fn(classifier_fn) # When num_instances == 1, demux can be replaced by filter, # but keep it as Demultiplexer for the sake of consistency # like throwing Error when classification result is out of o range container = _DemultiplexerIterDataPipe(datapipe, num_instances, classifier_fn, drop_none, buffer_size) return [_ChildDataPipe(container, i) for i in range(num_instances)] class _DemultiplexerIterDataPipe(IterDataPipe): r""" Container to hold instance-specific information on behalf of DemultiplexerIterDataPipe. It tracks the state of its child DataPipes, maintains the buffer, classifies and yields the next correct value as requested by the child DataPipes. """ def __init__(self, datapipe: IterDataPipe[T_co], num_instances: int, classifier_fn: Callable[[T_co], Optional[int]], drop_none: bool, buffer_size: int): self.main_datapipe = datapipe self._datapipe_iterator: Optional[Iterator[Any]] = None self.num_instances = num_instances self.buffer_size = buffer_size if self.buffer_size < 0: warnings.warn( "Unlimited buffer size is set for `demux`, " "please be aware of OOM at random places", UserWarning ) self.current_buffer_usage = 0 self.child_buffers: List[Deque[T_co]] = [deque() for _ in range(num_instances)] self.classifier_fn = classifier_fn self.drop_none = drop_none self.main_datapipe_exhausted = False def _find_next(self, instance_id: int) -> T_co: while True: if self.main_datapipe_exhausted: raise StopIteration if self._datapipe_iterator is None: raise ValueError( "_datapipe_iterator has not been set, likely because this private method is called directly " "without invoking get_next_element_by_instance() first.") value = next(self._datapipe_iterator) classification = self.classifier_fn(value) if classification is None and self.drop_none: continue if classification is None or classification >= self.num_instances or classification < 0: raise ValueError(f"Output of the classification fn should be between 0 and {self.num_instances - 1}. " + f"{classification} is returned.") if classification == instance_id: return value self.child_buffers[classification].append(value) self.current_buffer_usage += 1 if self.buffer_size >= 0 and self.current_buffer_usage > self.buffer_size: raise BufferError( f"DemultiplexerIterDataPipe buffer overflow, buffer size {self.buffer_size} is insufficient.") def get_next_element_by_instance(self, instance_id: int): if self._datapipe_iterator is None and not self.main_datapipe_exhausted: self._datapipe_iterator = iter(self.main_datapipe) stop = False while not stop: if self.child_buffers[instance_id]: self.current_buffer_usage -= 1 yield self.child_buffers[instance_id].popleft() else: try: yield self._find_next(instance_id) except StopIteration: stop = True self.main_datapipe_exhausted = True self._datapipe_iterator = None def is_every_instance_exhausted(self) -> bool: return self.main_datapipe_exhausted and all(not child_buffer for child_buffer in self.child_buffers) def reset(self) -> None: self._datapipe_iterator = None self.current_buffer_usage = 0 self.child_buffers = [deque() for _ in range(self.num_instances)] self.main_datapipe_exhausted = False def __getstate__(self): if IterDataPipe.getstate_hook is not None: return IterDataPipe.getstate_hook(self) state = ( self.main_datapipe, self.num_instances, self.buffer_size, self.classifier_fn, self.drop_none, ) return state def __setstate__(self, state): ( self.main_datapipe, self.num_instances, self.buffer_size, self.classifier_fn, self.drop_none, ) = state self._datapipe_iterator = None self.current_buffer_usage = 0 self.child_buffers = [deque() for _ in range(self.num_instances)] self.main_datapipe_exhausted = False def __del__(self): for dq in self.child_buffers: dq.clear() @functional_datapipe('mux') class MultiplexerIterDataPipe(IterDataPipe): r""" Yields one element at a time from each of the input Iterable DataPipes (functional name: ``mux``). As in, one element from the 1st input DataPipe, then one element from the 2nd DataPipe in the next iteration, and so on. It ends when the shortest input DataPipe is exhausted. Args: datapipes: Iterable DataPipes that will take turn to yield their elements, until the shortest DataPipe is exhausted Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> dp1, dp2, dp3 = IterableWrapper(range(3)), IterableWrapper(range(10, 15)), IterableWrapper(range(20, 25)) >>> list(dp1.mux(dp2, dp3)) [0, 10, 20, 1, 11, 21, 2, 12, 22] """ def __init__(self, *datapipes): self.datapipes = datapipes self.length: Optional[int] = None self.buffer: List = [] # Store values to be yielded only when every iterator provides one def __iter__(self): iterators = [iter(x) for x in self.datapipes] while len(iterators): for it in iterators: try: value = next(it) self.buffer.append(value) except StopIteration: self.buffer.clear() return for value in self.buffer: yield value self.buffer.clear() def __len__(self): if self.length is not None: if self.length == -1: raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) return self.length if all(isinstance(dp, Sized) for dp in self.datapipes): self.length = min(len(dp) for dp in self.datapipes) * len(self.datapipes) else: self.length = -1 return len(self) def reset(self) -> None: self.buffer = [] def __getstate__(self): if IterDataPipe.getstate_hook is not None: return IterDataPipe.getstate_hook(self) state = ( self.datapipes, self.length, ) return state def __setstate__(self, state): ( self.datapipes, self.length, ) = state self.buffer = [] def __del__(self): self.buffer.clear() @functional_datapipe('zip') class ZipperIterDataPipe(IterDataPipe[Tuple[T_co]]): r""" Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``). The output is stopped as soon as the shortest input DataPipe is exhausted. Args: *datapipes: Iterable DataPipes being aggregated Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> dp1, dp2, dp3 = IterableWrapper(range(5)), IterableWrapper(range(10, 15)), IterableWrapper(range(20, 25)) >>> list(dp1.zip(dp2, dp3)) [(0, 10, 20), (1, 11, 21), (2, 12, 22), (3, 13, 23), (4, 14, 24)] """ datapipes: Tuple[IterDataPipe] length: Optional[int] def __init__(self, *datapipes: IterDataPipe): if not all(isinstance(dp, IterDataPipe) for dp in datapipes): raise TypeError("All inputs are required to be `IterDataPipe` " "for `ZipIterDataPipe`.") super().__init__() self.datapipes = datapipes # type: ignore[assignment] self.length = None def __iter__(self) -> Iterator[Tuple[T_co]]: for data in zip(*self.datapipes): yield data def __len__(self) -> int: if self.length is not None: if self.length == -1: raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) return self.length if all(isinstance(dp, Sized) for dp in self.datapipes): self.length = min(len(dp) for dp in self.datapipes) else: self.length = -1 return len(self) msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/filelister.py000066400000000000000000000047141436713554100263660ustar00rootroot00000000000000from typing import Iterator, List, Sequence, Union from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.iter import IterableWrapper from ms_adapter.pytorch.utils.data.datapipes.utils.common import get_file_pathnames_from_root __all__ = ["FileListerIterDataPipe", ] @functional_datapipe("list_files") class FileListerIterDataPipe(IterDataPipe[str]): r""" Given path(s) to the root directory, yields file pathname(s) (path + filename) of files within the root directory. Multiple root directories can be provided. Args: root: Root directory or a sequence of root directories masks: Unix style filter string or string list for filtering file name(s) recursive: Whether to return pathname from nested directories or not abspath: Whether to return relative pathname or absolute pathname non_deterministic: Whether to return pathname in sorted order or not. If ``False``, the results yielded from each root directory will be sorted length: Nominal length of the datapipe Example: >>> from torchdata.datapipes.iter import FileLister >>> dp = FileLister(root=".", recursive=True) >>> list(dp) ['example.py', './data/data.tar'] """ def __init__( self, root: Union[str, Sequence[str], IterDataPipe] = '.', masks: Union[str, List[str]] = '', *, recursive: bool = False, abspath: bool = False, non_deterministic: bool = False, length: int = -1 ) -> None: super().__init__() if isinstance(root, str): root = [root, ] if not isinstance(root, IterDataPipe): root = IterableWrapper(root) self.datapipe: IterDataPipe = root self.masks: Union[str, List[str]] = masks self.recursive: bool = recursive self.abspath: bool = abspath self.non_deterministic: bool = non_deterministic self.length: int = length def __iter__(self) -> Iterator[str] : for path in self.datapipe: yield from get_file_pathnames_from_root(path, self.masks, self.recursive, self.abspath, self.non_deterministic) def __len__(self): if self.length == -1: raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) return self.length msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/fileopener.py000066400000000000000000000064111436713554100263500ustar00rootroot00000000000000from io import IOBase from typing import Iterable, Tuple, Optional from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.utils.common import get_file_binaries_from_pathnames, _deprecation_warning __all__ = [ "FileOpenerIterDataPipe", "FileLoaderIterDataPipe", ] @functional_datapipe("open_files") class FileOpenerIterDataPipe(IterDataPipe[Tuple[str, IOBase]]): r""" Given pathnames, opens files and yield pathname and file stream in a tuple (functional name: ``open_files``). Args: datapipe: Iterable datapipe that provides pathnames mode: An optional string that specifies the mode in which the file is opened by ``open()``. It defaults to ``b`` which means open for reading in binary mode. Another option is to use ``t`` for text mode encoding: An optional string that specifies the encoding of the underlying file. It defaults to ``None`` to match the default encoding of ``open``. length: Nominal length of the datapipe Note: The opened file handles will be closed by Python's GC periodically. Users can choose to close them explicitly. Example: >>> from torchdata.datapipes.iter import FileLister, FileOpener, StreamReader >>> dp = FileLister(root=".").filter(lambda fname: fname.endswith('.txt')) >>> dp = FileOpener(dp) >>> dp = StreamReader(dp) >>> list(dp) [('./abc.txt', 'abc')] """ def __init__( self, datapipe: Iterable[str], mode: str = 'r', encoding: Optional[str] = None, length: int = -1): super().__init__() self.datapipe: Iterable = datapipe self.mode: str = mode self.encoding: Optional[str] = encoding if self.mode not in ('b', 't', 'rb', 'rt', 'r'): raise ValueError("Invalid mode {}".format(mode)) # TODO: enforce typing for each instance based on mode, otherwise # `argument_validation` with this DataPipe may be potentially broken if 'b' in mode and encoding is not None: raise ValueError("binary mode doesn't take an encoding argument") self.length: int = length # Remove annotation due to 'IOBase' is a general type and true type # is determined at runtime based on mode. Some `DataPipe` requiring # a subtype would cause mypy error. def __iter__(self): yield from get_file_binaries_from_pathnames(self.datapipe, self.mode, self.encoding) def __len__(self): if self.length == -1: raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) return self.length class FileLoaderIterDataPipe(IterDataPipe[Tuple[str, IOBase]]): def __new__( cls, datapipe: Iterable[str], mode: str = 'b', length: int = -1): _deprecation_warning( cls.__name__, deprecation_version="1.12", removal_version="1.13", new_class_name="FileOpener", ) return FileOpenerIterDataPipe(datapipe=datapipe, mode=mode, length=length) msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/grouping.py000066400000000000000000000305101436713554100260470ustar00rootroot00000000000000from collections import defaultdict from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe, DataChunk from ms_adapter.pytorch.utils.data.datapipes.utils.common import _check_lambda_fn from typing import Any, Callable, DefaultDict, Iterator, List, Optional, Sized, TypeVar __all__ = [ "BatcherIterDataPipe", "GrouperIterDataPipe", "ShardingFilterIterDataPipe", "UnBatcherIterDataPipe", ] T_co = TypeVar('T_co', covariant=True) @functional_datapipe('sharding_filter') class ShardingFilterIterDataPipe(IterDataPipe): r""" Wrapper that allows DataPipe to be sharded (functional name: ``sharding_filter``). After ``apply_sharding`` is called, each instance of the DataPipe (on different workers) will have every `n`-th element of the original DataPipe, where `n` equals to the number of instances. Args: source_datapipe: Iterable DataPipe that will be sharded """ def __init__(self, source_datapipe: IterDataPipe): self.source_datapipe = source_datapipe self.num_of_instances = 1 self.instance_id = 0 def is_shardable(self): return True def apply_sharding(self, num_of_instances, instance_id): self.num_of_instances = num_of_instances self.instance_id = instance_id def __iter__(self): for i, item in enumerate(self.source_datapipe): if i % self.num_of_instances == self.instance_id: yield item def __len__(self): if isinstance(self.source_datapipe, Sized): return len(self.source_datapipe) // self.num_of_instances +\ (1 if (self.instance_id < len(self.source_datapipe) % self.num_of_instances) else 0) raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) @functional_datapipe('batch') class BatcherIterDataPipe(IterDataPipe[DataChunk]): r""" Creates mini-batches of data (functional name: ``batch``). An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``, or ``length % batch_size`` for the last batch if ``drop_last`` is set to ``False``. Args: datapipe: Iterable DataPipe being batched batch_size: The size of each batch drop_last: Option to drop the last batch if it's not full wrapper_class: wrapper to apply onto each batch (type ``List``) before yielding, defaults to ``DataChunk`` Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> dp = IterableWrapper(range(10)) >>> dp = dp.batch(batch_size=3, drop_last=True) >>> list(dp) [[0, 1, 2], [3, 4, 5], [6, 7, 8]] """ datapipe: IterDataPipe batch_size: int drop_last: bool length: Optional[int] def __init__(self, datapipe: IterDataPipe, batch_size: int, drop_last: bool = False, wrapper_class=DataChunk, ) -> None: assert batch_size > 0, "Batch size is required to be larger than 0!" super().__init__() self.datapipe = datapipe self.batch_size = batch_size self.drop_last = drop_last self.length = None self.wrapper_class = wrapper_class def __iter__(self) -> Iterator[DataChunk]: batch: List = [] for x in self.datapipe: batch.append(x) if len(batch) == self.batch_size: yield self.wrapper_class(batch) batch = [] if len(batch) > 0: if not self.drop_last: yield self.wrapper_class(batch) def __len__(self) -> int: if self.length is not None: return self.length if isinstance(self.datapipe, Sized): if self.drop_last: self.length = len(self.datapipe) // self.batch_size else: self.length = (len(self.datapipe) + self.batch_size - 1) // self.batch_size return self.length raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) @functional_datapipe('unbatch') class UnBatcherIterDataPipe(IterDataPipe): r""" Undoes batching of data (functional name: ``unbatch``). In other words, it flattens the data up to the specified level within a batched DataPipe. Args: datapipe: Iterable DataPipe being un-batched unbatch_level: Defaults to ``1`` (only flattening the top level). If set to ``2``, it will flatten the top two levels, and ``-1`` will flatten the entire DataPipe. Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> source_dp = IterableWrapper([[[0, 1], [2]], [[3, 4], [5]], [[6]]]) >>> dp1 = source_dp.unbatch() >>> list(dp1) [[0, 1], [2], [3, 4], [5], [6]] >>> dp2 = source_dp.unbatch(unbatch_level=2) >>> list(dp2) [0, 1, 2, 3, 4, 5, 6] """ def __init__(self, datapipe: IterDataPipe, unbatch_level: int = 1): self.datapipe = datapipe self.unbatch_level = unbatch_level def __iter__(self): for element in self.datapipe: for i in self._dive(element, unbatch_level=self.unbatch_level): yield i def _dive(self, element, unbatch_level): if unbatch_level < -1: raise ValueError("unbatch_level must be -1 or >= 0") if unbatch_level == -1: if isinstance(element, list) or isinstance(element, DataChunk): for item in element: for i in self._dive(item, unbatch_level=-1): yield i else: yield element elif unbatch_level == 0: yield element else: if isinstance(element, list) or isinstance(element, DataChunk): for item in element: for i in self._dive(item, unbatch_level=unbatch_level - 1): yield i else: raise IndexError(f"unbatch_level {self.unbatch_level} exceeds the depth of the DataPipe") @functional_datapipe('groupby') class GrouperIterDataPipe(IterDataPipe[DataChunk]): r""" Groups data from input IterDataPipe by keys which are generated from ``group_key_fn``, and yields a ``DataChunk`` with batch size up to ``group_size`` if defined (functional name: ``groupby``). The samples are read sequentially from the source ``datapipe``, and a batch of samples belonging to the same group will be yielded as soon as the size of the batch reaches ``group_size``. When the buffer is full, the DataPipe will yield the largest batch with the same key, provided that its size is larger than ``guaranteed_group_size``. If its size is smaller, it will be dropped if ``drop_remaining=True``. After iterating through the entirety of source ``datapipe``, everything not dropped due to the buffer capacity will be yielded from the buffer, even if the group sizes are smaller than ``guaranteed_group_size``. Args: datapipe: Iterable datapipe to be grouped group_key_fn: Function used to generate group key from the data of the source datapipe buffer_size: The size of buffer for ungrouped data group_size: The max size of each group, a batch is yielded as soon as it reaches this size guaranteed_group_size: The guaranteed minimum group size to be yielded in case the buffer is full drop_remaining: Specifies if the group smaller than ``guaranteed_group_size`` will be dropped from buffer when the buffer is full Example: >>> import os >>> from torchdata.datapipes.iter import IterableWrapper >>> def group_fn(file): ... return os.path.basename(file).split(".")[0] >>> source_dp = IterableWrapper(["a.png", "b.png", "a.json", "b.json", "a.jpg", "c.json"]) >>> dp0 = source_dp.groupby(group_key_fn=group_fn) >>> list(dp0) [['a.png', 'a.json', 'a.jpg'], ['b.png', 'b.json'], ['c.json']] >>> # A group is yielded as soon as its size equals to `group_size` >>> dp1 = source_dp.groupby(group_key_fn=group_fn, group_size=2) >>> list(dp1) [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']] >>> # Scenario where `buffer` is full, and group 'a' needs to be yielded since its size > `guaranteed_group_size` >>> dp2 = source_dp.groupby(group_key_fn=group_fn, buffer_size=3, group_size=3, guaranteed_group_size=2) >>> list(dp2) [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']] """ def __init__(self, datapipe: IterDataPipe[T_co], group_key_fn: Callable, *, buffer_size: int = 10000, group_size: Optional[int] = None, guaranteed_group_size: Optional[int] = None, drop_remaining: bool = False): _check_lambda_fn(group_key_fn) self.datapipe = datapipe self.group_key_fn = group_key_fn self.max_buffer_size = buffer_size self.buffer_elements: DefaultDict[Any, List] = defaultdict(list) self.curr_buffer_size = 0 self.group_size = group_size self.guaranteed_group_size = None if group_size is not None and buffer_size is not None: assert 0 < group_size <= buffer_size self.guaranteed_group_size = group_size if guaranteed_group_size is not None: assert group_size is not None and 0 < guaranteed_group_size <= group_size self.guaranteed_group_size = guaranteed_group_size self.drop_remaining = drop_remaining self.wrapper_class = DataChunk def _remove_biggest_key(self): biggest_key = None biggest_size = 0 result_to_yield = None for findkey in self.buffer_elements.keys(): if len(self.buffer_elements[findkey]) > biggest_size: biggest_size = len(self.buffer_elements[findkey]) biggest_key = findkey if self.guaranteed_group_size is not None and biggest_size < self.guaranteed_group_size and not self.drop_remaining: raise RuntimeError('Failed to group items', str(self.buffer_elements[biggest_key])) if self.guaranteed_group_size is None or biggest_size >= self.guaranteed_group_size: result_to_yield = self.buffer_elements[biggest_key] self.curr_buffer_size -= biggest_size del self.buffer_elements[biggest_key] return result_to_yield def __iter__(self): for x in self.datapipe: key = self.group_key_fn(x) self.buffer_elements[key].append(x) self.curr_buffer_size += 1 if self.group_size is not None and self.group_size == len(self.buffer_elements[key]): yield self.wrapper_class(self.buffer_elements[key]) self.curr_buffer_size -= len(self.buffer_elements[key]) del self.buffer_elements[key] if self.curr_buffer_size == self.max_buffer_size: result_to_yield = self._remove_biggest_key() if result_to_yield is not None: yield self.wrapper_class(result_to_yield) for key in tuple(self.buffer_elements.keys()): res = self.buffer_elements.pop(key) self.curr_buffer_size -= len(res) yield self.wrapper_class(res) def reset(self) -> None: self.curr_buffer_size = 0 self.buffer_elements = defaultdict(list) def __getstate__(self): if IterDataPipe.getstate_hook is not None: return IterDataPipe.getstate_hook(self) state = ( self.datapipe, self.group_key_fn, self.max_buffer_size, self.group_size, self.guaranteed_group_size, self.drop_remaining, self.wrapper_class, ) return state def __setstate__(self, state): ( self.datapipe, self.group_key_fn, self.max_buffer_size, self.group_size, self.guaranteed_group_size, self.drop_remaining, self.wrapper_class, ) = state self.curr_buffer_size = 0 self.buffer_elements = defaultdict(list) def __del__(self): self.buffer_elements.clear() msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/routeddecoder.py000066400000000000000000000053361436713554100270550ustar00rootroot00000000000000from io import BufferedIOBase from typing import Any, Callable, Iterable, Iterator, Sized, Tuple from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.utils.common import _deprecation_warning from ms_adapter.pytorch.utils.data.datapipes.utils.decoder import ( Decoder, basichandlers as decoder_basichandlers, imagehandler as decoder_imagehandler, extension_extract_fn ) __all__ = ["RoutedDecoderIterDataPipe", ] @functional_datapipe('routed_decode') class RoutedDecoderIterDataPipe(IterDataPipe[Tuple[str, Any]]): r""" Decodes binary streams from input DataPipe, yields pathname and decoded data in a tuple (functional name: ``routed_decode``). Args: datapipe: Iterable datapipe that provides pathname and binary stream in tuples handlers: Optional user defined decoder handlers. If ``None``, basic and image decoder handlers will be set as default. If multiple handles are provided, the priority order follows the order of handlers (the first handler has the top priority) key_fn: Function for decoder to extract key from pathname to dispatch handlers. Default is set to extract file extension from pathname Note: When ``key_fn`` is specified returning anything other than extension, the default handler will not work and users need to specify custom handler. Custom handler could use regex to determine the eligibility to handle data. """ def __init__(self, datapipe: Iterable[Tuple[str, BufferedIOBase]], *handlers: Callable, key_fn: Callable = extension_extract_fn) -> None: super().__init__() self.datapipe: Iterable[Tuple[str, BufferedIOBase]] = datapipe if not handlers: handlers = (decoder_basichandlers, decoder_imagehandler('torch')) self.decoder = Decoder(*handlers, key_fn=key_fn) _deprecation_warning( type(self).__name__, deprecation_version="1.12", removal_version="1.13", old_functional_name="routed_decode", ) def add_handler(self, *handler: Callable) -> None: self.decoder.add_handler(*handler) def __iter__(self) -> Iterator[Tuple[str, Any]]: for data in self.datapipe: pathname = data[0] result = self.decoder(data) yield (pathname, result[pathname]) def __len__(self) -> int: if isinstance(self.datapipe, Sized): return len(self.datapipe) raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/selecting.py000066400000000000000000000073621436713554100262030ustar00rootroot00000000000000from typing import Callable, Iterator, Optional, TypeVar from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe from ms_adapter.pytorch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper from ms_adapter.pytorch.utils.data.datapipes.utils.common import _check_lambda_fn, _deprecation_warning __all__ = ["FilterIterDataPipe", ] T_co = TypeVar('T_co', covariant=True) @functional_datapipe('filter') class FilterIterDataPipe(IterDataPipe[T_co]): r""" Filters out elements from the source datapipe according to input ``filter_fn`` (functional name: ``filter``). Args: datapipe: Iterable DataPipe being filtered filter_fn: Customized function mapping an element to a boolean. drop_empty_batches (Deprecated): By default, drops a batch if it is empty after filtering instead of keeping an empty list input_col: Index or indices of data which ``filter_fn`` is applied, such as: - ``None`` as default to apply ``filter_fn`` to the data directly. - Integer(s) is used for list/tuple. - Key(s) is used for dict. Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> def is_even(n): ... return n % 2 == 0 >>> dp = IterableWrapper(range(5)) >>> filter_dp = dp.filter(filter_fn=is_even) >>> list(filter_dp) [0, 2, 4] """ datapipe: IterDataPipe filter_fn: Callable drop_empty_batches: bool def __init__( self, datapipe: IterDataPipe, filter_fn: Callable, drop_empty_batches: Optional[bool] = None, input_col=None, ) -> None: super().__init__() self.datapipe = datapipe _check_lambda_fn(filter_fn) self.filter_fn = filter_fn # type: ignore[assignment] if drop_empty_batches is None: drop_empty_batches = True else: _deprecation_warning( type(self).__name__, deprecation_version="1.12", removal_version="1.14", old_argument_name="drop_empty_batches", ) self.drop_empty_batches = drop_empty_batches self.input_col = input_col def _apply_filter_fn(self, data) -> bool: if self.input_col is None: return self.filter_fn(data) elif isinstance(self.input_col, (list, tuple)): args = tuple(data[col] for col in self.input_col) return self.filter_fn(*args) else: return self.filter_fn(data[self.input_col]) def __iter__(self) -> Iterator[T_co]: for data in self.datapipe: filtered = self._returnIfTrue(data) if self._isNonEmpty(filtered): yield filtered def _returnIfTrue(self, data): condition = self._apply_filter_fn(data) if df_wrapper.is_column(condition): # We are operating on DataFrames filter here result = [] for idx, mask in enumerate(df_wrapper.iterate(condition)): if mask: result.append(df_wrapper.get_item(data, idx)) if len(result): return df_wrapper.concat(result) else: return None if not isinstance(condition, bool): raise ValueError("Boolean output is required for `filter_fn` of FilterIterDataPipe, got", type(condition)) if condition: return data def _isNonEmpty(self, data): if df_wrapper.is_dataframe(data): return True r = data is not None and \ not (isinstance(data, list) and len(data) == 0 and self.drop_empty_batches) return r msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/streamreader.py000066400000000000000000000025221436713554100266750ustar00rootroot00000000000000from typing import Tuple from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe __all__ = ["StreamReaderIterDataPipe", ] @functional_datapipe('read_from_stream') class StreamReaderIterDataPipe(IterDataPipe[Tuple[str, bytes]]): r""" Given IO streams and their label names, yields bytes with label name in a tuple (functional name: ``read_from_stream``). Args: datapipe: Iterable DataPipe provides label/URL and byte stream chunk: Number of bytes to be read from stream per iteration. If ``None``, all bytes will be read util the EOF. Example: >>> from torchdata.datapipes.iter import IterableWrapper, StreamReader >>> from io import StringIO >>> dp = IterableWrapper([("alphabet", StringIO("abcde"))]) >>> list(StreamReader(dp, chunk=1)) [('alphabet', 'a'), ('alphabet', 'b'), ('alphabet', 'c'), ('alphabet', 'd'), ('alphabet', 'e')] """ def __init__(self, datapipe, chunk=None): self.datapipe = datapipe self.chunk = chunk def __iter__(self): for furl, stream in self.datapipe: while True: d = stream.read(self.chunk) if not d: break yield (furl, d) msadapter/ms_adapter/pytorch/utils/data/datapipes/iter/utils.py000066400000000000000000000033751436713554100253660ustar00rootroot00000000000000import copy import warnings from ms_adapter.pytorch.utils.data.datapipes.datapipe import IterDataPipe __all__ = ["IterableWrapperIterDataPipe", ] class IterableWrapperIterDataPipe(IterDataPipe): r""" Wraps an iterable object to create an IterDataPipe. Args: iterable: Iterable object to be wrapped into an IterDataPipe deepcopy: Option to deepcopy input iterable object for each iterator. The copy is made when the first element is read in ``iter()``. .. note:: If ``deepcopy`` is explicitly set to ``False``, users should ensure that the data pipeline doesn't contain any in-place operations over the iterable instance to prevent data inconsistency across iterations. Example: >>> from torchdata.datapipes.iter import IterableWrapper >>> dp = IterableWrapper(range(10)) >>> list(dp) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] """ def __init__(self, iterable, deepcopy=True): self.iterable = iterable self.deepcopy = deepcopy def __iter__(self): source_data = self.iterable if self.deepcopy: try: source_data = copy.deepcopy(self.iterable) # For the case that data cannot be deep-copied, # all in-place operations will affect iterable variable. # When this DataPipe is iterated second time, it will # yield modified items. except TypeError: warnings.warn( "The input iterable can not be deepcopied, " "please be aware of in-place modification would affect source data." ) for data in source_data: yield data def __len__(self): return len(self.iterable) msadapter/ms_adapter/pytorch/utils/data/datapipes/map/000077500000000000000000000000001436713554100234565ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/datapipes/map/__init__.py000066400000000000000000000013201436713554100255630ustar00rootroot00000000000000# Functional DataPipe from ms_adapter.pytorch.utils.data.datapipes.map.callable import MapperMapDataPipe as Mapper from ms_adapter.pytorch.utils.data.datapipes.map.combinatorics import ShufflerMapDataPipe as Shuffler from ms_adapter.pytorch.utils.data.datapipes.map.combining import ( ConcaterMapDataPipe as Concater, ZipperMapDataPipe as Zipper ) from ms_adapter.pytorch.utils.data.datapipes.map.grouping import ( BatcherMapDataPipe as Batcher ) from ms_adapter.pytorch.utils.data.datapipes.map.utils import SequenceWrapperMapDataPipe as SequenceWrapper __all__ = ['Batcher', 'Concater', 'Mapper', 'SequenceWrapper', 'Shuffler', 'Zipper'] # Please keep this list sorted assert __all__ == sorted(__all__) msadapter/ms_adapter/pytorch/utils/data/datapipes/map/callable.py000066400000000000000000000034411436713554100255710ustar00rootroot00000000000000from ms_adapter.pytorch.utils.data.datapipes.utils.common import _check_lambda_fn from typing import Callable, TypeVar from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import MapDataPipe __all__ = ["MapperMapDataPipe", "default_fn"] T_co = TypeVar('T_co', covariant=True) # Default function to return each item directly # In order to keep datapipe picklable, eliminates the usage # of python lambda function def default_fn(data): return data @functional_datapipe('map') class MapperMapDataPipe(MapDataPipe[T_co]): r""" Apply the input function over each item from the source DataPipe (functional name: ``map``). The function can be any regular Python function or partial object. Lambda function is not recommended as it is not supported by pickle. Args: datapipe: Source MapDataPipe fn: Function being applied to each item Example: >>> from torchdata.datapipes.map import SequenceWrapper, Mapper >>> def add_one(x): ... return x + 1 >>> dp = SequenceWrapper(range(10)) >>> map_dp_1 = dp.map(add_one) >>> list(map_dp_1) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] >>> map_dp_2 = Mapper(dp, lambda x: x + 1) >>> list(map_dp_2) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] """ datapipe: MapDataPipe fn: Callable def __init__( self, datapipe: MapDataPipe, fn: Callable = default_fn, ) -> None: super().__init__() self.datapipe = datapipe _check_lambda_fn(fn) self.fn = fn # type: ignore[assignment] def __len__(self) -> int: return len(self.datapipe) def __getitem__(self, index) -> T_co: return self.fn(self.datapipe[index]) msadapter/ms_adapter/pytorch/utils/data/datapipes/map/combinatorics.py000066400000000000000000000045321436713554100266700ustar00rootroot00000000000000import random from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import MapDataPipe from typing import Iterator, List, Optional, TypeVar __all__ = ["ShufflerMapDataPipe", ] T_co = TypeVar('T_co', covariant=True) @functional_datapipe('shuffle') class ShufflerMapDataPipe(MapDataPipe[T_co]): r""" Shuffle the input DataPipe via its indices (functional name: ``shuffle``). When it is used with :class:`~ms_adapter.pytorch.utils.data.DataLoader`, the methods to set up random seed are different based on :attr:`num_workers`. For single-process mode (:attr:`num_workers == 0`), the random seed is set before the :class:`~ms_adapter.pytorch.utils.data.DataLoader` in the main process. For multi-process mode (:attr:`num_worker > 0`), ``worker_init_fn`` is used to set up a random seed for each worker process. Args: datapipe: MapDataPipe being shuffled indices: a list of indices of the MapDataPipe. If not provided, we assume it uses 0-based indexing Example: >>> from torchdata.datapipes.map import SequenceWrapper >>> dp = SequenceWrapper(range(10)) >>> shuffle_dp = dp.shuffle() >>> list(shuffle_dp) [0, 4, 1, 6, 3, 2, 9, 5, 7, 8] """ datapipe: MapDataPipe[T_co] def __init__(self, datapipe: MapDataPipe[T_co], *, indices: Optional[List] = None, ) -> None: super().__init__() self.datapipe = datapipe self.indices = list(range(len(datapipe))) if indices is None else indices self.index_map = {index_name: num_index for num_index, index_name in enumerate(self.indices)} # We do not lazily shuffle because this way is significantly faster in terms of total time random.shuffle(self.indices) def __getitem__(self, index) -> T_co: old_numeric_index = self.index_map[index] new_index = self.indices[old_numeric_index] return self.datapipe[new_index] # Without __iter__ implemented, by default it tries to use 0-index, # which doesn't work when there is a custom index. def __iter__(self) -> Iterator[T_co]: for i in self.indices: yield self.datapipe[i] def __len__(self) -> int: return len(self.datapipe) msadapter/ms_adapter/pytorch/utils/data/datapipes/map/combining.py000066400000000000000000000072531436713554100260040ustar00rootroot00000000000000from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import MapDataPipe from typing import Sized, Tuple, TypeVar __all__ = ["ConcaterMapDataPipe", "ZipperMapDataPipe"] T_co = TypeVar('T_co', covariant=True) @functional_datapipe('concat') class ConcaterMapDataPipe(MapDataPipe): r""" Concatenate multiple Map DataPipes (functional name: ``concat``). The new index of is the cumulative sum of source DataPipes. For example, if there are 2 source DataPipes both with length 5, index 0 to 4 of the resulting `ConcatMapDataPipe` would refer to elements of the first DataPipe, and 5 to 9 would refer to elements of the second DataPipe. Args: datapipes: Map DataPipes being concatenated Example: >>> from torchdata.datapipes.map import SequenceWrapper >>> dp1 = SequenceWrapper(range(3)) >>> dp2 = SequenceWrapper(range(3)) >>> concat_dp = dp1.concat(dp2) >>> list(concat_dp) [0, 1, 2, 0, 1, 2] """ datapipes: Tuple[MapDataPipe] length: int def __init__(self, *datapipes: MapDataPipe): if len(datapipes) == 0: raise ValueError("Expected at least one DataPipe, but got nothing") if not all(isinstance(dp, MapDataPipe) for dp in datapipes): raise TypeError("Expected all inputs to be `MapDataPipe`") if not all(isinstance(dp, Sized) for dp in datapipes): raise TypeError("Expected all inputs to be `Sized`") self.datapipes = datapipes # type: ignore[assignment] self.length = -1 def __getitem__(self, index) -> T_co: offset = 0 for dp in self.datapipes: if index - offset < len(dp): return dp[index - offset] else: offset += len(dp) raise IndexError("Index {} is out of range.".format(index)) def __len__(self) -> int: if self.length == -1: self.length = sum(len(dp) for dp in self.datapipes) return self.length @functional_datapipe('zip') class ZipperMapDataPipe(MapDataPipe[Tuple[T_co, ...]]): r""" Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``). This MataPipe is out of bound as soon as the shortest input DataPipe is exhausted. Args: *datapipes: Map DataPipes being aggregated Example: >>> from torchdata.datapipes.map import SequenceWrapper >>> dp1 = SequenceWrapper(range(3)) >>> dp2 = SequenceWrapper(range(10, 13)) >>> zip_dp = dp1.zip(dp2) >>> list(zip_dp) [(0, 10), (1, 11), (2, 12)] """ datapipes: Tuple[MapDataPipe[T_co], ...] length: int def __init__(self, *datapipes: MapDataPipe[T_co]) -> None: if len(datapipes) == 0: raise ValueError("Expected at least one DataPipe, but got nothing") if not all(isinstance(dp, MapDataPipe) for dp in datapipes): raise TypeError("Expected all inputs to be `MapDataPipe`") if not all(isinstance(dp, Sized) for dp in datapipes): raise TypeError("Expected all inputs to be `Sized`") self.datapipes = datapipes self.length = -1 def __getitem__(self, index) -> Tuple[T_co, ...]: res = [] for dp in self.datapipes: try: res.append(dp[index]) except IndexError: raise IndexError(f"Index {index} is out of range for one of the input MapDataPipes {dp}.") return tuple(res) def __len__(self) -> int: if self.length == -1: self.length = min(len(dp) for dp in self.datapipes) return self.length msadapter/ms_adapter/pytorch/utils/data/datapipes/map/grouping.py000066400000000000000000000050361436713554100256660ustar00rootroot00000000000000from ms_adapter.pytorch.utils.data.datapipes._decorator import functional_datapipe from ms_adapter.pytorch.utils.data.datapipes.datapipe import MapDataPipe, DataChunk from typing import List, Optional, Sized, TypeVar __all__ = ["BatcherMapDataPipe", ] T = TypeVar('T') @functional_datapipe('batch') class BatcherMapDataPipe(MapDataPipe[DataChunk]): r""" Create mini-batches of data (functional name: ``batch``). An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``, or ``length % batch_size`` for the last batch if ``drop_last`` is set to ``False``. Args: datapipe: Iterable DataPipe being batched batch_size: The size of each batch drop_last: Option to drop the last batch if it's not full Example: >>> from torchdata.datapipes.map import SequenceWrapper >>> dp = SequenceWrapper(range(10)) >>> batch_dp = dp.batch(batch_size=2) >>> list(batch_dp) [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]] """ datapipe: MapDataPipe batch_size: int drop_last: bool length: Optional[int] def __init__(self, datapipe: MapDataPipe[T], batch_size: int, drop_last: bool = False, wrapper_class=DataChunk, ) -> None: assert batch_size > 0, "Batch size is required to be larger than 0!" super().__init__() self.datapipe = datapipe self.batch_size = batch_size self.drop_last = drop_last self.length = None self.wrapper_class = wrapper_class def __getitem__(self, index) -> DataChunk: batch: List = [] indices = range(index * self.batch_size, (index + 1) * self.batch_size) try: for i in indices: batch.append(self.datapipe[i]) return self.wrapper_class(batch) except IndexError: if not self.drop_last and len(batch) > 0: return self.wrapper_class(batch) else: raise IndexError(f"Index {index} is out of bound.") def __len__(self) -> int: if self.length is not None: return self.length if isinstance(self.datapipe, Sized): if self.drop_last: self.length = len(self.datapipe) // self.batch_size else: self.length = (len(self.datapipe) + self.batch_size - 1) // self.batch_size return self.length raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) msadapter/ms_adapter/pytorch/utils/data/datapipes/map/utils.py000066400000000000000000000027721436713554100252000ustar00rootroot00000000000000import copy import warnings from ms_adapter.pytorch.utils.data.datapipes.datapipe import MapDataPipe __all__ = ["SequenceWrapperMapDataPipe", ] class SequenceWrapperMapDataPipe(MapDataPipe): r""" Wraps a sequence object into a MapDataPipe. Args: sequence: Sequence object to be wrapped into an MapDataPipe deepcopy: Option to deepcopy input sequence object .. note:: If ``deepcopy`` is set to False explicitly, users should ensure that data pipeline doesn't contain any in-place operations over the iterable instance, in order to prevent data inconsistency across iterations. Example: >>> from torchdata.datapipes.map import SequenceWrapper >>> dp = SequenceWrapper(range(10)) >>> list(dp) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] >>> dp = SequenceWrapper({'a': 100, 'b': 200, 'c': 300, 'd': 400}) >>> dp['a'] 100 """ def __init__(self, sequence, deepcopy=True): if deepcopy: try: self.sequence = copy.deepcopy(sequence) except TypeError: warnings.warn( "The input sequence can not be deepcopied, " "please be aware of in-place modification would affect source data" ) self.sequence = sequence else: self.sequence = sequence def __getitem__(self, index): return self.sequence[index] def __len__(self): return len(self.sequence) msadapter/ms_adapter/pytorch/utils/data/datapipes/utils/000077500000000000000000000000001436713554100240415ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/datapipes/utils/__init__.py000066400000000000000000000000001436713554100261400ustar00rootroot00000000000000msadapter/ms_adapter/pytorch/utils/data/datapipes/utils/common.py000066400000000000000000000175551436713554100257200ustar00rootroot00000000000000import os import fnmatch import warnings from io import IOBase from typing import Dict, Iterable, List, Tuple, Union, Optional from ms_adapter.pytorch.utils.data._utils.serialization import DILL_AVAILABLE __all__ = [ "StreamWrapper", "get_file_binaries_from_pathnames", "get_file_pathnames_from_root", "match_masks", "validate_pathname_binary_tuple", ] def _check_lambda_fn(fn): # Partial object has no attribute '__name__', but can be pickled if hasattr(fn, "__name__") and fn.__name__ == "" and not DILL_AVAILABLE: warnings.warn( "Lambda function is not supported for pickle, please use " "regular python function or functools.partial instead." ) def match_masks(name : str, masks : Union[str, List[str]]) -> bool: # empty mask matches any input name if not masks: return True if isinstance(masks, str): return fnmatch.fnmatch(name, masks) for mask in masks: if fnmatch.fnmatch(name, mask): return True return False def get_file_pathnames_from_root( root: str, masks: Union[str, List[str]], recursive: bool = False, abspath: bool = False, non_deterministic: bool = False) -> Iterable[str]: # print out an error message and raise the error out def onerror(err : OSError): warnings.warn(err.filename + " : " + err.strerror) raise err if os.path.isfile(root): path = root if abspath: path = os.path.abspath(path) fname = os.path.basename(path) if match_masks(fname, masks): yield path else: for path, dirs, files in os.walk(root, onerror=onerror): if abspath: path = os.path.abspath(path) if not non_deterministic: files.sort() for f in files: if match_masks(f, masks): yield os.path.join(path, f) if not recursive: break if not non_deterministic: # Note that this is in-place modifying the internal list from `os.walk` # This only works because `os.walk` doesn't shallow copy before turn # https://github.com/python/cpython/blob/f4c03484da59049eb62a9bf7777b963e2267d187/Lib/os.py#L407 dirs.sort() def get_file_binaries_from_pathnames(pathnames: Iterable, mode: str, encoding: Optional[str] = None): if not isinstance(pathnames, Iterable): pathnames = [pathnames, ] if mode in ('b', 't'): mode = 'r' + mode for pathname in pathnames: if not isinstance(pathname, str): raise TypeError("Expected string type for pathname, but got {}" .format(type(pathname))) yield pathname, StreamWrapper(open(pathname, mode, encoding=encoding)) def validate_pathname_binary_tuple(data: Tuple[str, IOBase]): if not isinstance(data, tuple): raise TypeError(f"pathname binary data should be tuple type, but it is type {type(data)}") if len(data) != 2: raise TypeError(f"pathname binary stream tuple length should be 2, but got {len(data)}") if not isinstance(data[0], str): raise TypeError(f"pathname within the tuple should have string type pathname, but it is type {type(data[0])}") if not isinstance(data[1], IOBase) and not isinstance(data[1], StreamWrapper): raise TypeError( f"binary stream within the tuple should have IOBase or" f"its subclasses as type, but it is type {type(data[1])}" ) # Deprecated function names and its corresponding DataPipe type and kwargs for the `_deprecation_warning` function _iter_deprecated_functional_names: Dict[str, Dict] = {"open_file_by_fsspec": {"old_class_name": "FSSpecFileOpener", "deprecation_version": "0.4.0", "removal_version": "0.6.0", "old_functional_name": "open_file_by_fsspec", "new_functional_name": "open_files_by_fsspec", "deprecate_functional_name_only": True}, "open_file_by_iopath": {"old_class_name": "IoPathFileOpener", "deprecation_version": "0.4.0", "removal_version": "0.6.0", "old_functional_name": "open_file_by_iopath", "new_functional_name": "open_files_by_iopath", "deprecate_functional_name_only": True}} _map_deprecated_functional_names: Dict[str, Dict] = {} def _deprecation_warning( old_class_name: str, *, deprecation_version: str, removal_version: str, old_functional_name: str = "", old_argument_name: str = "", new_class_name: str = "", new_functional_name: str = "", new_argument_name: str = "", deprecate_functional_name_only: bool = False, ) -> None: if new_functional_name and not old_functional_name: raise ValueError("Old functional API needs to be specified for the deprecation warning.") if new_argument_name and not old_argument_name: raise ValueError("Old argument name needs to be specified for the deprecation warning.") if old_functional_name and old_argument_name: raise ValueError("Deprecating warning for functional API and argument should be separated.") msg = f"`{old_class_name}()`" if deprecate_functional_name_only and old_functional_name: msg = f"{msg}'s functional API `.{old_functional_name}()` is" elif old_functional_name: msg = f"{msg} and its functional API `.{old_functional_name}()` are" elif old_argument_name: msg = f"The argument `{old_argument_name}` of {msg} is" else: msg = f"{msg} is" msg = ( f"{msg} deprecated since {deprecation_version} and will be removed in {removal_version}." f"\nSee https://github.com/pytorch/data/issues/163 for details." ) if new_class_name or new_functional_name: msg = f"{msg}\nPlease use" if new_class_name: msg = f"{msg} `{new_class_name}()`" if new_class_name and new_functional_name: msg = f"{msg} or" if new_functional_name: msg = f"{msg} `.{new_functional_name}()`" msg = f"{msg} instead." if new_argument_name: msg = f"{msg}\nPlease use `{old_class_name}({new_argument_name}=)` instead." warnings.warn(msg, FutureWarning) class StreamWrapper: ''' StreamWrapper is introduced to wrap file handler generated by DataPipe operation like `FileOpener`. StreamWrapper would guarantee the wrapped file handler is closed when it's out of scope. ''' def __init__(self, file_obj): self.file_obj = file_obj def __getattr__(self, name): file_obj = self.__dict__['file_obj'] return getattr(file_obj, name) def __dir__(self): attrs = list(self.__dict__.keys()) + list(StreamWrapper.__dict__.keys()) attrs += dir(self.file_obj) return list(set(list(attrs))) def __del__(self): try: self.file_obj.close() except AttributeError: pass def __iter__(self): for line in self.file_obj: yield line def __next__(self): return next(self.file_obj) def __repr__(self): return f"StreamWrapper<{self.file_obj!r}>" def __getstate__(self): return self.file_obj def __setstate__(self, obj): self.file_obj = obj msadapter/ms_adapter/pytorch/utils/data/datapipes/utils/decoder.py000066400000000000000000000254231436713554100260260ustar00rootroot00000000000000# This file takes partial of the implementation from NVIDIA's webdataset at here: # https://github.com/tmbdev/webdataset/blob/master/webdataset/autodecode.py import io import json import os.path import pickle import tempfile import torch from ms_adapter.pytorch.utils.data.datapipes.utils.common import StreamWrapper __all__ = [ "Decoder", "ImageHandler", "MatHandler", "audiohandler", "basichandlers", "extension_extract_fn", "handle_extension", "imagehandler", "mathandler", "videohandler", ] ################################################################ # handle basic datatypes ################################################################ def basichandlers(extension, data): if extension in "txt text transcript": return data.decode("utf-8") if extension in "cls cls2 class count index inx id".split(): try: return int(data) except ValueError: return None if extension in "json jsn": return json.loads(data) if extension in "pyd pickle".split(): return pickle.loads(data) if extension in "pt".split(): stream = io.BytesIO(data) return torch.load(stream) # if extension in "ten tb".split(): # from . import tenbin # return tenbin.decode_buffer(data) # if extension in "mp msgpack msg".split(): # import msgpack # return msgpack.unpackb(data) return None ################################################################ # handle images ################################################################ imagespecs = { "l8": ("numpy", "uint8", "l"), "rgb8": ("numpy", "uint8", "rgb"), "rgba8": ("numpy", "uint8", "rgba"), "l": ("numpy", "float", "l"), "rgb": ("numpy", "float", "rgb"), "rgba": ("numpy", "float", "rgba"), "torchl8": ("torch", "uint8", "l"), "torchrgb8": ("torch", "uint8", "rgb"), "torchrgba8": ("torch", "uint8", "rgba"), "torchl": ("torch", "float", "l"), "torchrgb": ("torch", "float", "rgb"), "torch": ("torch", "float", "rgb"), "torchrgba": ("torch", "float", "rgba"), "pill": ("pil", None, "l"), "pil": ("pil", None, "rgb"), "pilrgb": ("pil", None, "rgb"), "pilrgba": ("pil", None, "rgba"), } def handle_extension(extensions, f): """ Returns a decoder handler function for the list of extensions. Extensions can be a space separated list of extensions. Extensions can contain dots, in which case the corresponding number of extension components must be present in the key given to f. Comparisons are case insensitive. Examples: handle_extension("jpg jpeg", my_decode_jpg) # invoked for any file.jpg handle_extension("seg.jpg", special_case_jpg) # invoked only for file.seg.jpg """ extensions = extensions.lower().split() def g(key, data): extension = key.lower().split(".") for target in extensions: target = target.split(".") if len(target) > len(extension): continue if extension[-len(target):] == target: return f(data) return None return g class ImageHandler: """ Decode image data using the given `imagespec`. The `imagespec` specifies whether the image is decoded to numpy/torch/pi, decoded to uint8/float, and decoded to l/rgb/rgba: - l8: numpy uint8 l - rgb8: numpy uint8 rgb - rgba8: numpy uint8 rgba - l: numpy float l - rgb: numpy float rgb - rgba: numpy float rgba - torchl8: torch uint8 l - torchrgb8: torch uint8 rgb - torchrgba8: torch uint8 rgba - torchl: torch float l - torchrgb: torch float rgb - torch: torch float rgb - torchrgba: torch float rgba - pill: pil None l - pil: pil None rgb - pilrgb: pil None rgb - pilrgba: pil None rgba """ def __init__(self, imagespec): assert imagespec in list(imagespecs.keys()), "unknown image specification: {}".format(imagespec) self.imagespec = imagespec.lower() def __call__(self, extension, data): if extension.lower() not in "jpg jpeg png ppm pgm pbm pnm".split(): return None try: import numpy as np except ImportError as e: raise ModuleNotFoundError("Package `numpy` is required to be installed for default image decoder." "Please use `pip install numpy` to install the package") try: import PIL.Image except ImportError as e: raise ModuleNotFoundError("Package `PIL` is required to be installed for default image decoder." "Please use `pip install Pillow` to install the package") imagespec = self.imagespec atype, etype, mode = imagespecs[imagespec] with io.BytesIO(data) as stream: img = PIL.Image.open(stream) img.load() img = img.convert(mode.upper()) if atype == "pil": return img elif atype == "numpy": result = np.asarray(img) assert result.dtype == np.uint8, "numpy image array should be type uint8, but got {}".format(result.dtype) if etype == "uint8": return result else: return result.astype("f") / 255.0 elif atype == "torch": result = np.asarray(img) assert result.dtype == np.uint8, "numpy image array should be type uint8, but got {}".format(result.dtype) if etype == "uint8": result = np.array(result.transpose(2, 0, 1)) return torch.tensor(result) else: result = np.array(result.transpose(2, 0, 1)) return torch.tensor(result) / 255.0 return None def imagehandler(imagespec): return ImageHandler(imagespec) ################################################################ # torch video ################################################################ def videohandler(extension, data): if extension not in "mp4 ogv mjpeg avi mov h264 mpg webm wmv".split(): return None try: import torchvision.io except ImportError as e: raise ModuleNotFoundError("Package `torchvision` is required to be installed for default video file loader." "Please use `pip install torchvision` or `conda install torchvision -c pytorch`" "to install the package") with tempfile.TemporaryDirectory() as dirname: fname = os.path.join(dirname, f"file.{extension}") with open(fname, "wb") as stream: stream.write(data) return torchvision.io.read_video(fname) ################################################################ # torchaudio ################################################################ def audiohandler(extension, data): if extension not in ["flac", "mp3", "sox", "wav", "m4a", "ogg", "wma"]: return None try: import torchaudio # type: ignore[import] except ImportError as e: raise ModuleNotFoundError("Package `torchaudio` is required to be installed for default audio file loader." "Please use `pip install torchaudio` or `conda install torchaudio -c pytorch`" "to install the package") with tempfile.TemporaryDirectory() as dirname: fname = os.path.join(dirname, f"file.{extension}") with open(fname, "wb") as stream: stream.write(data) return torchaudio.load(fname) ################################################################ # mat ################################################################ class MatHandler: def __init__(self, **loadmat_kwargs) -> None: try: import scipy.io as sio except ImportError as e: raise ModuleNotFoundError("Package `scipy` is required to be installed for mat file." "Please use `pip install scipy` or `conda install scipy`" "to install the package") self.sio = sio self.loadmat_kwargs = loadmat_kwargs def __call__(self, extension, data): if extension != 'mat': return None with io.BytesIO(data) as stream: return self.sio.loadmat(stream, **self.loadmat_kwargs) def mathandler(**loadmat_kwargs): return MatHandler(**loadmat_kwargs) ################################################################ # a sample decoder ################################################################ # Extract extension from pathname def extension_extract_fn(pathname): ext = os.path.splitext(pathname)[1] # Remove dot if ext: ext = ext[1:] return ext class Decoder: """ Decode key/data sets using a list of handlers. For each key/data item, this iterates through the list of handlers until some handler returns something other than None. """ def __init__(self, *handler, key_fn=extension_extract_fn): self.handlers = list(handler) if handler else [] self.key_fn = key_fn # Insert new handler from the beginning of handlers list to make sure the new # handler having the highest priority def add_handler(self, *handler): if not handler: return self.handlers = list(handler) + self.handlers @staticmethod def _is_stream_handle(data): obj_to_check = data.file_obj if isinstance(data, StreamWrapper) else data return isinstance(obj_to_check, io.BufferedIOBase) or isinstance(obj_to_check, io.RawIOBase) def decode1(self, key, data): if not data: return data # if data is a stream handle, we need to read all the content before decoding if Decoder._is_stream_handle(data): ds = data # The behavior of .read can differ between streams (e.g. HTTPResponse), hence this is used instead data = b"".join(data) ds.close() for f in self.handlers: result = f(key, data) if result is not None: return result return data def decode(self, data): result = {} # single data tuple(pathname, data stream) if isinstance(data, tuple): data = [data] if data is not None: for k, v in data: # TODO: xinyu, figure out why Nvidia do this? if k[0] == "_": if isinstance(v, bytes): v = v.decode("utf-8") result[k] = v continue result[k] = self.decode1(self.key_fn(k), v) return result def __call__(self, data): return self.decode(data) msadapter/ms_adapter/pytorch/utils/data/dataset.py000066400000000000000000000322621436713554100227330ustar00rootroot00000000000000import bisect import warnings from typing import ( Generic, Iterable, Iterator, List, Optional, Sequence, Tuple, TypeVar, ) import math # No 'default_generator' in torch/__init__.pyi # from ms_adapter.pytorch import default_generator, randperm # from ms_adapter.pytorch import Tensor import ms_adapter.pytorch as torch def _accumulate(iterable, fn=lambda x, y: x + y): 'Return running totals' # _accumulate([1,2,3,4,5]) --> 1 3 6 10 15 # _accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120 it = iter(iterable) try: total = next(it) except StopIteration: return yield total for element in it: total = fn(total, element) yield total __all__ = [ "Dataset", "IterableDataset", "TensorDataset", "ConcatDataset", "ChainDataset", "Subset", "random_split", ] T_co = TypeVar('T_co', covariant=True) T = TypeVar('T') class Dataset(Generic[T_co]): r"""An abstract class representing a :class:`Dataset`. All datasets that represent a map from keys to data samples should subclass it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a data sample for a given key. Subclasses could also optionally overwrite :meth:`__len__`, which is expected to return the size of the dataset by many :class:`~torch.utils.data.Sampler` implementations and the default options of :class:`~torch.utils.data.DataLoader`. .. note:: :class:`~torch.utils.data.DataLoader` by default constructs a index sampler that yields integral indices. To make it work with a map-style dataset with non-integral indices/keys, a custom sampler must be provided. """ def __getitem__(self, index) -> T_co: raise NotImplementedError def __add__(self, other: 'Dataset[T_co]') -> 'ConcatDataset[T_co]': return ConcatDataset([self, other]) # No `def __len__(self)` default? # See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] # in pytorch/torch/utils/data/sampler.py class IterableDataset(Dataset[T_co]): r"""An iterable Dataset. All datasets that represent an iterable of data samples should subclass it. Such form of datasets is particularly useful when data come from a stream. All subclasses should overwrite :meth:`__iter__`, which would return an iterator of samples in this dataset. When a subclass is used with :class:`~torch.utils.data.DataLoader`, each item in the dataset will be yielded from the :class:`~torch.utils.data.DataLoader` iterator. When :attr:`num_workers > 0`, each worker process will have a different copy of the dataset object, so it is often desired to configure each copy independently to avoid having duplicate data returned from the workers. :func:`~torch.utils.data.get_worker_info`, when called in a worker process, returns information about the worker. It can be used in either the dataset's :meth:`__iter__` method or the :class:`~torch.utils.data.DataLoader` 's :attr:`worker_init_fn` option to modify each copy's behavior. Example 1: splitting workload across all workers in :meth:`__iter__`:: >>> class MyIterableDataset(torch.utils.data.IterableDataset): ... def __init__(self, start, end): ... super(MyIterableDataset).__init__() ... assert end > start, "this example code only works with end >= start" ... self.start = start ... self.end = end ... ... def __iter__(self): ... worker_info = torch.utils.data.get_worker_info() ... if worker_info is None: # single-process data loading, return the full iterator ... iter_start = self.start ... iter_end = self.end ... else: # in a worker process ... # split workload ... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers))) ... worker_id = worker_info.id ... iter_start = self.start + worker_id * per_worker ... iter_end = min(iter_start + per_worker, self.end) ... return iter(range(iter_start, iter_end)) ... >>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6]. >>> ds = MyIterableDataset(start=3, end=7) >>> # Single-process loading >>> print(list(torch.utils.data.DataLoader(ds, num_workers=0))) [3, 4, 5, 6] >>> # Mult-process loading with two worker processes >>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6]. >>> print(list(torch.utils.data.DataLoader(ds, num_workers=2))) [3, 5, 4, 6] >>> # With even more workers >>> print(list(torch.utils.data.DataLoader(ds, num_workers=20))) [3, 4, 5, 6] Example 2: splitting workload across all workers using :attr:`worker_init_fn`:: >>> class MyIterableDataset(torch.utils.data.IterableDataset): ... def __init__(self, start, end): ... super(MyIterableDataset).__init__() ... assert end > start, "this example code only works with end >= start" ... self.start = start ... self.end = end ... ... def __iter__(self): ... return iter(range(self.start, self.end)) ... >>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6]. >>> ds = MyIterableDataset(start=3, end=7) >>> # Single-process loading >>> print(list(torch.utils.data.DataLoader(ds, num_workers=0))) [3, 4, 5, 6] >>> >>> # Directly doing multi-process loading yields duplicate data >>> print(list(torch.utils.data.DataLoader(ds, num_workers=2))) [3, 3, 4, 4, 5, 5, 6, 6] >>> # Define a `worker_init_fn` that configures each dataset copy differently >>> def worker_init_fn(worker_id): ... worker_info = torch.utils.data.get_worker_info() ... dataset = worker_info.dataset # the dataset copy in this worker process ... overall_start = dataset.start ... overall_end = dataset.end ... # configure the dataset to only process the split workload ... per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers))) ... worker_id = worker_info.id ... dataset.start = overall_start + worker_id * per_worker ... dataset.end = min(dataset.start + per_worker, overall_end) ... >>> # Mult-process loading with the custom `worker_init_fn` >>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6]. >>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn))) [3, 5, 4, 6] >>> # With even more workers >>> print(list(torch.utils.data.DataLoader(ds, num_workers=20, worker_init_fn=worker_init_fn))) [3, 4, 5, 6] """ def __iter__(self) -> Iterator[T_co]: raise NotImplementedError def __add__(self, other: Dataset[T_co]): return ChainDataset([self, other]) # No `def __len__(self)` default? Subclasses raise `TypeError` when needed. # See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] class TensorDataset(Dataset[Tuple[Tensor, ...]]): r"""Dataset wrapping tensors. Each sample will be retrieved by indexing tensors along the first dimension. Args: *tensors (Tensor): tensors that have the same size of the first dimension. """ tensors: Tuple[Tensor, ...] def __init__(self, *tensors: Tensor) -> None: assert all(tensors[0].shape[0] == tensor.shape[0] for tensor in tensors), "Size mismatch between tensors" self.tensors = tensors def __getitem__(self, index): return tuple(tensor[index] for tensor in self.tensors) def __len__(self): return self.tensors[0].shape[0] class ConcatDataset(Dataset[T_co]): r"""Dataset as a concatenation of multiple datasets. This class is useful to assemble different existing datasets. Args: datasets (sequence): List of datasets to be concatenated """ datasets: List[Dataset[T_co]] cumulative_sizes: List[int] @staticmethod def cumsum(sequence): r, s = [], 0 for e in sequence: l = len(e) r.append(l + s) s += l return r def __init__(self, datasets: Iterable[Dataset]) -> None: super(ConcatDataset, self).__init__() self.datasets = list(datasets) assert len(self.datasets) > 0, 'datasets should not be an empty iterable' # type: ignore[arg-type] for d in self.datasets: assert not isinstance(d, IterableDataset), "ConcatDataset does not support IterableDataset" self.cumulative_sizes = self.cumsum(self.datasets) def __len__(self): return self.cumulative_sizes[-1] def __getitem__(self, idx): if idx < 0: if -idx > len(self): raise ValueError("absolute value of index should not exceed dataset length") idx = len(self) + idx dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) if dataset_idx == 0: sample_idx = idx else: sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] return self.datasets[dataset_idx][sample_idx] @property def cummulative_sizes(self): warnings.warn("cummulative_sizes attribute is renamed to " "cumulative_sizes", DeprecationWarning, stacklevel=2) return self.cumulative_sizes class ChainDataset(IterableDataset): r"""Dataset for chaining multiple :class:`IterableDataset` s. This class is useful to assemble different existing dataset streams. The chaining operation is done on-the-fly, so concatenating large-scale datasets with this class will be efficient. Args: datasets (iterable of IterableDataset): datasets to be chained together """ def __init__(self, datasets: Iterable[Dataset]) -> None: super(ChainDataset, self).__init__() self.datasets = datasets def __iter__(self): for d in self.datasets: assert isinstance(d, IterableDataset), "ChainDataset only supports IterableDataset" for x in d: yield x def __len__(self): total = 0 for d in self.datasets: assert isinstance(d, IterableDataset), "ChainDataset only supports IterableDataset" total += len(d) # type: ignore[arg-type] return total class Subset(Dataset[T_co]): r""" Subset of a dataset at specified indices. Args: dataset (Dataset): The whole Dataset indices (sequence): Indices in the whole set selected for subset """ dataset: Dataset[T_co] indices: Sequence[int] def __init__(self, dataset: Dataset[T_co], indices: Sequence[int]) -> None: self.dataset = dataset self.indices = indices def __getitem__(self, idx): if isinstance(idx, list): return self.dataset[[self.indices[i] for i in idx]] return self.dataset[self.indices[idx]] def __len__(self): return len(self.indices) def random_split(dataset: Dataset[T], lengths: Sequence[int], generator = None) -> List[Subset[T]]: r""" Randomly split a dataset into non-overlapping new datasets of given lengths. Optionally fix the generator for reproducible results, e.g.: >>> random_split(range(10), [3, 7], generator=torch.Generator().manual_seed(42)) Args: dataset (Dataset): Dataset to be split lengths (sequence): lengths of splits to be produced generator (Generator): Generator used for the random permutation. """ # Cannot verify that dataset is Sized if math.isclose(sum(lengths), 1) and sum(lengths) <= 1: subset_lengths: List[int] = [] for i, frac in enumerate(lengths): if frac < 0 or frac > 1: raise ValueError(f"Fraction at index {i} is not between 0 and 1") n_items_in_split = int( math.floor(len(dataset) * frac) # type: ignore[arg-type] ) subset_lengths.append(n_items_in_split) remainder = len(dataset) - sum(subset_lengths) # type: ignore[arg-type] # add 1 to all the lengths in round-robin fashion until the remainder is 0 for i in range(remainder): idx_to_add_at = i % len(subset_lengths) subset_lengths[idx_to_add_at] += 1 lengths = subset_lengths for i, length in enumerate(lengths): if length == 0: warnings.warn(f"Length of split at index {i} is 0. " f"This might result in an empty dataset.") if sum(lengths) != len(dataset): # type: ignore[arg-type] raise ValueError("Sum of input lengths does not equal the length of the input dataset!") if generator is not None: warnings.warn("Don't support generator now.") indices = torch.randperm(sum(lengths)).asnumpy().tolist() return [Subset(dataset, indices[offset - length : offset]) for offset, length in zip(_accumulate(lengths), lengths)] msadapter/ms_adapter/pytorch/utils/data/distributed.py000066400000000000000000000142611436713554100236270ustar00rootroot00000000000000import math import numpy as np from typing import TypeVar, Optional, Iterator import ms_adapter.pytorch as torch from . import Sampler, Dataset # import torch.distributed as dist TODO:dist import mindspore.communication as dist __all__ = ["DistributedSampler", ] T_co = TypeVar('T_co', covariant=True) class DistributedSampler(Sampler[T_co]): r"""Sampler that restricts data loading to a subset of the dataset. It is especially useful in conjunction with :class:`torch.nn.parallel.DistributedDataParallel`. In such a case, each process can pass a :class:`~torch.utils.data.DistributedSampler` instance as a :class:`~torch.utils.data.DataLoader` sampler, and load a subset of the original dataset that is exclusive to it. .. note:: Dataset is assumed to be of constant size and that any instance of it always returns the same elements in the same order. Args: dataset: Dataset used for sampling. num_replicas (int, optional): Number of processes participating in distributed training. By default, :attr:`world_size` is retrieved from the current distributed group. rank (int, optional): Rank of the current process within :attr:`num_replicas`. By default, :attr:`rank` is retrieved from the current distributed group. shuffle (bool, optional): If ``True`` (default), sampler will shuffle the indices. seed (int, optional): random seed used to shuffle the sampler if :attr:`shuffle=True`. This number should be identical across all processes in the distributed group. Default: ``0``. drop_last (bool, optional): if ``True``, then the sampler will drop the tail of the data to make it evenly divisible across the number of replicas. If ``False``, the sampler will add extra indices to make the data evenly divisible across the replicas. Default: ``False``. .. warning:: In distributed mode, calling the :meth:`set_epoch` method at the beginning of each epoch **before** creating the :class:`DataLoader` iterator is necessary to make shuffling work properly across multiple epochs. Otherwise, the same ordering will be always used. Example:: >>> sampler = DistributedSampler(dataset) if is_distributed else None >>> loader = DataLoader(dataset, shuffle=(sampler is None), ... sampler=sampler) >>> for epoch in range(start_epoch, n_epochs): ... if is_distributed: ... sampler.set_epoch(epoch) ... train(loader) """ def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None, rank: Optional[int] = None, shuffle: bool = True, seed: int = 0, drop_last: bool = False) -> None: if num_replicas is None: # if not dist.is_available(): TODO:dist # raise RuntimeError("Requires distributed package to be available") # num_replicas = dist.get_world_size() TODO:dist num_replicas = dist.get_group_size() if rank is None: # if not dist.is_available(): TODO:dist # raise RuntimeError("Requires distributed package to be available") rank = dist.get_rank() if rank >= num_replicas or rank < 0: raise ValueError( "Invalid rank {}, rank should be in the interval" " [0, {}]".format(rank, num_replicas - 1)) self.dataset = dataset self.num_replicas = num_replicas self.rank = rank self.epoch = 0 self.drop_last = drop_last # If the dataset length is evenly divisible by # of replicas, then there # is no need to drop any data, since the dataset will be split equally. if self.drop_last and len(self.dataset) % self.num_replicas != 0: # type: ignore[arg-type] # Split to nearest available length that is evenly divisible. # This is to ensure each rank receives the same amount of data when # using this Sampler. self.num_samples = math.ceil( (len(self.dataset) - self.num_replicas) / self.num_replicas # type: ignore[arg-type] ) else: self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) # type: ignore[arg-type] self.total_size = self.num_samples * self.num_replicas self.shuffle = shuffle self.seed = seed def __iter__(self) -> Iterator[T_co]: if self.shuffle: # deterministically shuffle based on epoch and seed # TODO:random # g = torch.Generator() # g.manual_seed(self.seed + self.epoch) # indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] rnd = np.random.RandomState(self.seed + self.epoch) indices = rnd.choice(np.arange(len(self.dataset)), len(self.dataset), replace=False).tolist() else: indices = list(range(len(self.dataset))) # type: ignore[arg-type] if not self.drop_last: # add extra samples to make it evenly divisible padding_size = self.total_size - len(indices) if padding_size <= len(indices): indices += indices[:padding_size] else: indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size] else: # remove tail of data to make it evenly divisible. indices = indices[:self.total_size] assert len(indices) == self.total_size # subsample indices = indices[self.rank:self.total_size:self.num_replicas] assert len(indices) == self.num_samples return iter(indices) def __len__(self) -> int: return self.num_samples def set_epoch(self, epoch: int) -> None: r""" Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas use a different random ordering for each epoch. Otherwise, the next iteration of this sampler will yield the same ordering. Args: epoch (int): Epoch number. """ self.epoch = epoch msadapter/ms_adapter/pytorch/utils/data/graph.py000066400000000000000000000061101436713554100224000ustar00rootroot00000000000000# from __future__ import absolute_import # # import io # import pickle # # from . import IterDataPipe,MapDataPipe # # from ms_adapter.pytorch.utils.data._utils.serialization import DILL_AVAILABLE # # from typing import Any, Dict, Set, Tuple, Type, Union # # __all__ = ["traverse", ] # # DataPipe = Union[IterDataPipe, MapDataPipe] # reduce_ex_hook = None # # # def _stub_unpickler(): # return "STUB" # # # # TODO(VitalyFedyunin): Make sure it works without dill module installed # def _list_connected_datapipes(scan_obj, only_datapipe, cache): # f = io.BytesIO() # p = pickle.Pickler(f) # Not going to work for lambdas, but dill infinite loops on typing and can't be used as is # if DILL_AVAILABLE: # from dill import Pickler as dill_Pickler # d = dill_Pickler(f) # else: # d = None # # def stub_pickler(obj): # return _stub_unpickler, () # # captured_connections = [] # # def getstate_hook(obj): # state = {} # for k, v in obj.__dict__.items(): # if isinstance(v, (IterDataPipe, MapDataPipe, tuple)): # state[k] = v # return state # # def reduce_hook(obj): # if obj == scan_obj or obj in cache: # raise NotImplementedError # else: # captured_connections.append(obj) # return _stub_unpickler, () # # datapipe_classes: Tuple[Type[DataPipe]] = (IterDataPipe, MapDataPipe) # type: ignore[assignment] # # try: # for cls in datapipe_classes: # cls.set_reduce_ex_hook(reduce_hook) # if only_datapipe: # cls.set_getstate_hook(getstate_hook) # try: # p.dump(scan_obj) # except (pickle.PickleError, AttributeError, TypeError): # if DILL_AVAILABLE: # d.dump(scan_obj) # else: # raise # finally: # for cls in datapipe_classes: # cls.set_reduce_ex_hook(None) # if only_datapipe: # cls.set_getstate_hook(None) # if DILL_AVAILABLE: # from dill import extend as dill_extend # dill_extend(False) # Undo change to dispatch table # return captured_connections # # # def traverse(datapipe, only_datapipe=False): # cache: Set[DataPipe] = set() # return _traverse_helper(datapipe, only_datapipe, cache) # # # # Add cache here to prevent infinite recursion on DataPipe # def _traverse_helper(datapipe, only_datapipe, cache): # if not isinstance(datapipe, (IterDataPipe, MapDataPipe)): # raise RuntimeError("Expected `IterDataPipe` or `MapDataPipe`, but {} is found".format(type(datapipe))) # # cache.add(datapipe) # items = _list_connected_datapipes(datapipe, only_datapipe, cache) # d: Dict[DataPipe, Any] = {datapipe: {}} # for item in items: # # Using cache.copy() here is to prevent recursion on a single path rather than global graph # # Single DataPipe can present multiple times in different paths in graph # d[datapipe].update(_traverse_helper(item, only_datapipe, cache.copy())) # return d msadapter/ms_adapter/pytorch/utils/data/graph_settings.py000066400000000000000000000051631436713554100243270ustar00rootroot00000000000000# from __future__ import absolute_import # # import ms_adapter.pytorch as torch # import ms_adapter.pytorch.utils.data.graph # from ms_adapter.pytorch.utils.data.datapipes.iter import Shuffler # import warnings # # __all__ = [ # "apply_sharding", # "apply_shuffle_seed", # "apply_shuffle_settings", # "get_all_graph_pipes", # ] # # # def get_all_graph_pipes(graph): # results = set() # for datapipe, sub_graph in graph.items(): # results.add(datapipe) # sub_items = get_all_graph_pipes(sub_graph) # for item in sub_items: # results.add(item) # return results # # # def apply_sharding(datapipe, num_of_instances, instance_id): # graph = ms_adapter.pytorch.utils.data.graph.traverse(datapipe, only_datapipe=True) # all_pipes = get_all_graph_pipes(graph) # already_applied_to = None # for pipe in all_pipes: # if hasattr(pipe, 'is_shardable'): # if pipe.is_shardable(): # if hasattr(pipe, 'apply_sharding'): # if already_applied_to is not None: # raise RuntimeError('This implementation of sharding can be only applied once per instance of DataPipeline.', # 'Already applied to', already_applied_to, 'while trying to apply to', pipe) # pipe.apply_sharding(num_of_instances, instance_id) # already_applied_to = pipe # # # def apply_shuffle_settings(datapipe, shuffle): # if shuffle is None: # return datapipe # # graph = ms_adapter.pytorch.utils.data.graph.traverse(datapipe, only_datapipe=True) # all_pipes = get_all_graph_pipes(graph) # shufflers = {pipe for pipe in all_pipes if isinstance(pipe, Shuffler)} # if not shufflers and shuffle: # warnings.warn( # "`shuffle=True` was set, but the datapipe does not contain a `Shuffler`. Adding one at the end. " # "Be aware that the default buffer size might not be sufficient for your task." # ) # datapipe = datapipe.shuffle() # shufflers = {datapipe} # # for shuffler in shufflers: # shuffler.set_shuffle(shuffle) # # return datapipe # # # def apply_shuffle_seed(datapipe, rng): # graph = ms_adapter.pytorch.utils.data.graph.traverse(datapipe, only_datapipe=True) # all_pipes = get_all_graph_pipes(graph) # shufflers = {pipe for pipe in all_pipes if isinstance(pipe, Shuffler)} # # for shuffler in shufflers: # shuffle_seed = int(torch.empty((), dtype=torch.int64).random_(generator=rng).item()) # shuffler.set_seed(shuffle_seed) # # return datapipe msadapter/ms_adapter/pytorch/utils/data/sampler.py000066400000000000000000000247311436713554100227530ustar00rootroot00000000000000import warnings import ms_adapter.pytorch as torch from ms_adapter.pytorch.tensor import Tensor from typing import Iterator, Iterable, Optional, Sequence, List, TypeVar, Generic, Sized, Union import mindspore as ms __all__ = [ "BatchSampler", "RandomSampler", "Sampler", "SequentialSampler", "SubsetRandomSampler", "WeightedRandomSampler", ] T_co = TypeVar('T_co', covariant=True) class Sampler(Generic[T_co]): r"""Base class for all Samplers. Every Sampler subclass has to provide an :meth:`__iter__` method, providing a way to iterate over indices of dataset elements, and a :meth:`__len__` method that returns the length of the returned iterators. .. note:: The :meth:`__len__` method isn't strictly required by :class:`~torch.utils.data.DataLoader`, but is expected in any calculation involving the length of a :class:`~torch.utils.data.DataLoader`. """ def __init__(self, data_source: Optional[Sized]) -> None: pass def __iter__(self) -> Iterator[T_co]: raise NotImplementedError # NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] # # Many times we have an abstract class representing a collection/iterable of # data, e.g., `torch.utils.data.Sampler`, with its subclasses optionally # implementing a `__len__` method. In such cases, we must make sure to not # provide a default implementation, because both straightforward default # implementations have their issues: # # + `return NotImplemented`: # Calling `len(subclass_instance)` raises: # TypeError: 'NotImplementedType' object cannot be interpreted as an integer # # + `raise NotImplementedError()`: # This prevents triggering some fallback behavior. E.g., the built-in # `list(X)` tries to call `len(X)` first, and executes a different code # path if the method is not found or `NotImplemented` is returned, while # raising an `NotImplementedError` will propagate and and make the call # fail where it could have use `__iter__` to complete the call. # # Thus, the only two sensible things to do are # # + **not** provide a default `__len__`. # # + raise a `TypeError` instead, which is what Python uses when users call # a method that is not defined on an object. class SequentialSampler(Sampler[int]): r"""Samples elements sequentially, always in the same order. Args: data_source (Dataset): dataset to sample from """ data_source: Sized def __init__(self, data_source: Sized) -> None: self.data_source = data_source def __iter__(self) -> Iterator[int]: return iter(range(len(self.data_source))) def __len__(self) -> int: return len(self.data_source) class RandomSampler(Sampler[int]): r"""Samples elements randomly. If without replacement, then sample from a shuffled dataset. If with replacement, then user can specify :attr:`num_samples` to draw. Args: data_source (Dataset): dataset to sample from replacement (bool): samples are drawn on-demand with replacement if ``True``, default=``False`` num_samples (int): number of samples to draw, default=`len(dataset)`. generator (Generator): Generator used in sampling. """ data_source: Sized replacement: bool def __init__(self, data_source: Sized, replacement: bool = False, num_samples: Optional[int] = None, generator=None) -> None: self.data_source = data_source self.replacement = replacement self._num_samples = num_samples self.generator = generator if generator is not None: warnings.warn("Don't support generator now.") if not isinstance(self.replacement, bool): raise TypeError("replacement should be a boolean value, but got " "replacement={}".format(self.replacement)) if not isinstance(self.num_samples, int) or self.num_samples <= 0: raise ValueError("num_samples should be a positive integer " "value, but got num_samples={}".format(self.num_samples)) @property def num_samples(self) -> int: # dataset size might change at runtime if self._num_samples is None: return len(self.data_source) return self._num_samples def __iter__(self) -> Iterator[int]: n = len(self.data_source) if self.replacement: for _ in range(self.num_samples // 32): yield from torch.randint(low=0, high=n, size=(32,), dtype=torch.int64).asnumpy().tolist() yield from torch.randint(low=0, high=n, size=(self.num_samples % 32,), dtype=torch.int64).asnumpy().tolist() else: for _ in range(self.num_samples // n): yield from torch.randperm(n).asnumpy().tolist() yield from torch.randperm(n).asnumpy().tolist()[:self.num_samples % n] def __len__(self) -> int: return self.num_samples class SubsetRandomSampler(Sampler[int]): r"""Samples elements randomly from a given list of indices, without replacement. Args: indices (sequence): a sequence of indices generator (Generator): Generator used in sampling. """ indices: Sequence[int] def __init__(self, indices: Sequence[int], generator=None) -> None: self.indices = indices self.generator = generator if generator is not None: warnings.warn("Don't support generator now.") def __iter__(self) -> Iterator[int]: index = torch.randperm(len(self.indices)) for i in index: yield self.indices[i] def __len__(self) -> int: return len(self.indices) class WeightedRandomSampler(Sampler[int]): r"""Samples elements from ``[0,..,len(weights)-1]`` with given probabilities (weights). Args: weights (sequence) : a sequence of weights, not necessary summing up to one num_samples (int): number of samples to draw replacement (bool): if ``True``, samples are drawn with replacement. If not, they are drawn without replacement, which means that when a sample index is drawn for a row, it cannot be drawn again for that row. generator (Generator): Generator used in sampling. Example: >>> list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True)) [4, 4, 1, 4, 5] >>> list(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False)) [0, 1, 4, 3, 2] """ weights: Tensor num_samples: int replacement: bool def __init__(self, weights: Sequence[float], num_samples: int, replacement: bool = True, generator=None) -> None: if not isinstance(num_samples, int) or isinstance(num_samples, bool) or \ num_samples <= 0: raise ValueError("num_samples should be a positive integer " "value, but got num_samples={}".format(num_samples)) if not isinstance(replacement, bool): raise ValueError("replacement should be a boolean value, but got " "replacement={}".format(replacement)) self.weights = Tensor(weights, dtype=ms.float32) self.num_samples = num_samples self.replacement = replacement self.generator = generator if generator is not None: warnings.warn("Don't support generator now.") def __iter__(self) -> Iterator[int]: rand_tensor = torch.multinomial(self.weights, self.num_samples, self.replacement).asnumpy().tolist() yield from iter(rand_tensor) def __len__(self) -> int: return self.num_samples class BatchSampler(Sampler[List[int]]): r"""Wraps another sampler to yield a mini-batch of indices. Args: sampler (Sampler or Iterable): Base sampler. Can be any iterable object batch_size (int): Size of mini-batch. drop_last (bool): If ``True``, the sampler will drop the last batch if its size would be less than ``batch_size`` Example: >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False)) [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True)) [[0, 1, 2], [3, 4, 5], [6, 7, 8]] """ def __init__(self, sampler: Union[Sampler[int], Iterable[int]], batch_size: int, drop_last: bool) -> None: # Since collections.abc.Iterable does not check for `__getitem__`, which # is one way for an object to be an iterable, we don't do an `isinstance` # check here. if not isinstance(batch_size, int) or isinstance(batch_size, bool) or \ batch_size <= 0: raise ValueError("batch_size should be a positive integer value, " "but got batch_size={}".format(batch_size)) if not isinstance(drop_last, bool): raise ValueError("drop_last should be a boolean value, but got " "drop_last={}".format(drop_last)) self.sampler = sampler self.batch_size = batch_size self.drop_last = drop_last def __iter__(self) -> Iterator[List[int]]: if self.drop_last: sampler_iter = iter(self.sampler) while True: try: batch = [next(sampler_iter) for _ in range(self.batch_size)] yield batch except StopIteration: break else: batch = [0] * self.batch_size idx_in_batch = 0 for idx in self.sampler: batch[idx_in_batch] = idx idx_in_batch += 1 if idx_in_batch == self.batch_size: yield batch idx_in_batch = 0 batch = [0] * self.batch_size if idx_in_batch > 0: yield batch[:idx_in_batch] def __len__(self) -> int: # Can only be called if self.sampler has __len__ implemented # We cannot enforce this condition, so we turn off typechecking for the # implementation below. # Somewhat related: see NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] if self.drop_last: return len(self.sampler) // self.batch_size # type: ignore[arg-type] else: return (len(self.sampler) + self.batch_size - 1) // self.batch_size # type: ignore[arg-type] msadapter/ms_adapter/torchvision/000077500000000000000000000000001436713554100175355ustar00rootroot00000000000000msadapter/ms_adapter/torchvision/__init__.py000066400000000000000000000060401436713554100216460ustar00rootroot00000000000000import os import warnings # import torch from ms_adapter.torchvision import datasets from ms_adapter.torchvision import io # from ms_adapter.torchvision import models # from ms_adapter.torchvision import ops from ms_adapter.torchvision import transforms # from ms_adapter.torchvision import utils from .extension import _HAS_OPS try: from .version import __version__ # noqa: F401 except ImportError: pass # Check if torchvision is being imported within the root folder if not _HAS_OPS and os.path.dirname(os.path.realpath(__file__)) == os.path.join( os.path.realpath(os.getcwd()), "torchvision" ): message = ( "You are importing torchvision within its own root folder ({}). " "This is not expected to work and may give errors. Please exit the " "torchvision project source and relaunch your python interpreter." ) warnings.warn(message.format(os.getcwd())) _image_backend = "PIL" _video_backend = "pyav" def set_image_backend(backend): """ Specifies the package used to load images. Args: backend (string): Name of the image backend. one of {'PIL', 'accimage'}. The :mod:`accimage` package uses the Intel IPP library. It is generally faster than PIL, but does not support as many operations. """ global _image_backend if backend not in ["PIL", "accimage"]: raise ValueError(f"Invalid backend '{backend}'. Options are 'PIL' and 'accimage'") _image_backend = backend def get_image_backend(): """ Gets the name of the package used to load images """ return _image_backend def set_video_backend(backend): """ Specifies the package used to decode videos. Args: backend (string): Name of the video backend. one of {'pyav', 'video_reader'}. The :mod:`pyav` package uses the 3rd party PyAv library. It is a Pythonic binding for the FFmpeg libraries. The :mod:`video_reader` package includes a native C++ implementation on top of FFMPEG libraries, and a python API of TorchScript custom operator. It generally decodes faster than :mod:`pyav`, but is perhaps less robust. .. note:: Building with FFMPEG is disabled by default in the latest `main`. If you want to use the 'video_reader' backend, please compile torchvision from source. """ global _video_backend if backend not in ["pyav", "video_reader"]: raise ValueError("Invalid video backend '%s'. Options are 'pyav' and 'video_reader'" % backend) if backend == "video_reader" and not io._HAS_VIDEO_OPT: message = "video_reader video backend is not available. Please compile torchvision from source and try again" warnings.warn(message) else: _video_backend = backend def get_video_backend(): """ Returns the currently active video backend used to decode videos. Returns: str: Name of the video backend. one of {'pyav', 'video_reader'}. """ return _video_backend # def _is_tracing(): # return torch._C._get_tracing_state() msadapter/ms_adapter/torchvision/_internally_replaced_utils.py000066400000000000000000000040171436713554100255100ustar00rootroot00000000000000import importlib.machinery import os ENV_TORCH_HOME = 'TORCH_HOME' ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME' DEFAULT_CACHE_DIR = '~/.cache' def _get_torch_home(): torch_home = os.path.expanduser( os.getenv(ENV_TORCH_HOME, os.path.join(os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'))) return torch_home _HOME = os.path.join(_get_torch_home(), "datasets", "vision") _USE_SHARDED_DATASETS = False def _download_file_from_remote_location(fpath: str, url: str) -> None: pass def _is_remote_location_available() -> bool: return False # try: # from torch.hub import load_state_dict_from_url # noqa: 401 # except ImportError: # from torch.utils.model_zoo import load_url as load_state_dict_from_url # noqa: 401 def _get_extension_path(lib_name): lib_dir = os.path.dirname(__file__) if os.name == "nt": # Register the main torchvision library location on the default DLL path import ctypes import sys kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) with_load_library_flags = hasattr(kernel32, "AddDllDirectory") prev_error_mode = kernel32.SetErrorMode(0x0001) if with_load_library_flags: kernel32.AddDllDirectory.restype = ctypes.c_void_p if sys.version_info >= (3, 8): os.add_dll_directory(lib_dir) elif with_load_library_flags: res = kernel32.AddDllDirectory(lib_dir) if res is None: err = ctypes.WinError(ctypes.get_last_error()) err.strerror += f' Error adding "{lib_dir}" to the DLL directories.' raise err kernel32.SetErrorMode(prev_error_mode) loader_details = (importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES) extfinder = importlib.machinery.FileFinder(lib_dir, loader_details) ext_specs = extfinder.find_spec(lib_name) if ext_specs is None: raise ImportError return ext_specs.origin msadapter/ms_adapter/torchvision/_utils.py000066400000000000000000000016461436713554100214150ustar00rootroot00000000000000import enum from typing import Sequence, TypeVar, Type T = TypeVar("T", bound=enum.Enum) class StrEnumMeta(enum.EnumMeta): auto = enum.auto def from_str(self: Type[T], member: str) -> T: # type: ignore[misc] try: return self[member] except KeyError: # TODO: use `add_suggestion` from torchvision.prototype.utils._internal to improve the error message as # soon as it is migrated. raise ValueError(f"Unknown value '{member}' for {self.__name__}.") from None class StrEnum(enum.Enum, metaclass=StrEnumMeta): pass def sequence_to_str(seq: Sequence, separate_last: str = "") -> str: if not seq: return "" if len(seq) == 1: return f"'{seq[0]}'" head = "'" + "', '".join([str(item) for item in seq[:-1]]) + "'" tail = f"{'' if separate_last and len(seq) == 2 else ','} {separate_last}'{seq[-1]}'" return head + tail msadapter/ms_adapter/torchvision/datasets/000077500000000000000000000000001436713554100213455ustar00rootroot00000000000000msadapter/ms_adapter/torchvision/datasets/__init__.py000066400000000000000000000047721436713554100234700ustar00rootroot00000000000000from ._optical_flow import KittiFlow, Sintel, FlyingChairs, FlyingThings3D, HD1K from .caltech import Caltech101, Caltech256 from .celeba import CelebA from .cifar import CIFAR10, CIFAR100 from .cityscapes import Cityscapes from .clevr import CLEVRClassification from .coco import CocoCaptions, CocoDetection from .country211 import Country211 from .dtd import DTD from .eurosat import EuroSAT from .fakedata import FakeData from .fer2013 import FER2013 from .fgvc_aircraft import FGVCAircraft from .flickr import Flickr8k, Flickr30k from .flowers102 import Flowers102 from .folder import ImageFolder, DatasetFolder from .food101 import Food101 from .gtsrb import GTSRB from .hmdb51 import HMDB51 from .imagenet import ImageNet from .inaturalist import INaturalist from .kinetics import Kinetics400, Kinetics from .kitti import Kitti from .lfw import LFWPeople, LFWPairs from .lsun import LSUN, LSUNClass from .mnist import MNIST, EMNIST, FashionMNIST, KMNIST, QMNIST from .omniglot import Omniglot from .oxford_iiit_pet import OxfordIIITPet from .pcam import PCAM from .phototour import PhotoTour from .places365 import Places365 from .rendered_sst2 import RenderedSST2 from .sbd import SBDataset from .sbu import SBU from .semeion import SEMEION from .stanford_cars import StanfordCars from .stl10 import STL10 from .sun397 import SUN397 from .svhn import SVHN from .ucf101 import UCF101 from .usps import USPS from .vision import VisionDataset from .voc import VOCSegmentation, VOCDetection from .widerface import WIDERFace __all__ = ( "LSUN", "LSUNClass", "ImageFolder", "DatasetFolder", "FakeData", "CocoCaptions", "CocoDetection", "CIFAR10", "CIFAR100", "EMNIST", "FashionMNIST", "QMNIST", "MNIST", "KMNIST", "StanfordCars", "STL10", "SUN397", "SVHN", "PhotoTour", "SEMEION", "Omniglot", "SBU", "Flickr8k", "Flickr30k", "Flowers102", "VOCSegmentation", "VOCDetection", "Cityscapes", "ImageNet", "Caltech101", "Caltech256", "CelebA", "WIDERFace", "SBDataset", "VisionDataset", "USPS", "Kinetics400", "Kinetics", "HMDB51", "UCF101", "Places365", "Kitti", "INaturalist", "LFWPeople", "LFWPairs", "KittiFlow", "Sintel", "FlyingChairs", "FlyingThings3D", "HD1K", "Food101", "DTD", "FER2013", "GTSRB", "CLEVRClassification", "OxfordIIITPet", "PCAM", "Country211", "FGVCAircraft", "EuroSAT", "RenderedSST2", ) msadapter/ms_adapter/torchvision/datasets/_optical_flow.py000066400000000000000000000446341436713554100245530ustar00rootroot00000000000000import itertools import os import re from abc import ABC, abstractmethod from glob import glob from pathlib import Path import numpy as np import ms_adapter.pytorch as torch from PIL import Image from ..io.image import _read_png_16 from .utils import verify_str_arg from .vision import VisionDataset __all__ = ( "KittiFlow", "Sintel", "FlyingThings3D", "FlyingChairs", "HD1K", ) class FlowDataset(ABC, VisionDataset): # Some datasets like Kitti have a built-in valid_flow_mask, indicating which flow values are valid # For those we return (img1, img2, flow, valid_flow_mask), and for the rest we return (img1, img2, flow), # and it's up to whatever consumes the dataset to decide what valid_flow_mask should be. def __init__(self, root, transforms=None): super().__init__(root=root) self.transforms = transforms self._has_builtin_flow_mask = False self._flow_list = [] self._image_list = [] def _read_img(self, file_name): img = Image.open(file_name) if img.mode != "RGB": img = img.convert("RGB") return img @abstractmethod def _read_flow(self, file_name): # Return the flow or a tuple with the flow and the valid_flow_mask if _has_builtin_flow_mask is True pass def __getitem__(self, index): img1 = self._read_img(self._image_list[index][0]) img2 = self._read_img(self._image_list[index][1]) if self._flow_list: # it will be empty for some dataset when split="test" flow = self._read_flow(self._flow_list[index]) if self._has_builtin_flow_mask: flow, valid_flow_mask = flow else: valid_flow_mask = None else: flow = valid_flow_mask = None if self.transforms is not None: img1, img2, flow, valid_flow_mask = self.transforms(img1, img2, flow, valid_flow_mask) if self._has_builtin_flow_mask or valid_flow_mask is not None: # The `or valid_flow_mask is not None` part is here because the mask can be generated within a transform return img1, img2, flow, valid_flow_mask else: return img1, img2, flow def __len__(self): return len(self._image_list) def __rmul__(self, v): return torch.utils.data.ConcatDataset([self] * v) class Sintel(FlowDataset): """`Sintel `_ Dataset for optical flow. The dataset is expected to have the following structure: :: root Sintel testing clean scene_1 scene_2 ... final scene_1 scene_2 ... training clean scene_1 scene_2 ... final scene_1 scene_2 ... flow scene_1 scene_2 ... Args: root (string): Root directory of the Sintel Dataset. split (string, optional): The dataset split, either "train" (default) or "test" pass_name (string, optional): The pass to use, either "clean" (default), "final", or "both". See link above for details on the different passes. transforms (callable, optional): A function/transform that takes in ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. """ def __init__(self, root, split="train", pass_name="clean", transforms=None): super().__init__(root=root, transforms=transforms) verify_str_arg(split, "split", valid_values=("train", "test")) verify_str_arg(pass_name, "pass_name", valid_values=("clean", "final", "both")) passes = ["clean", "final"] if pass_name == "both" else [pass_name] root = Path(root) / "Sintel" flow_root = root / "training" / "flow" for pass_name in passes: split_dir = "training" if split == "train" else split image_root = root / split_dir / pass_name for scene in os.listdir(image_root): image_list = sorted(glob(str(image_root / scene / "*.png"))) for i in range(len(image_list) - 1): self._image_list += [[image_list[i], image_list[i + 1]]] if split == "train": self._flow_list += sorted(glob(str(flow_root / scene / "*.flo"))) def __getitem__(self, index): """Return example at given index. Args: index(int): The index of the example to retrieve Returns: tuple: A 3-tuple with ``(img1, img2, flow)``. The flow is a numpy array of shape (2, H, W) and the images are PIL images. ``flow`` is None if ``split="test"``. If a valid flow mask is generated within the ``transforms`` parameter, a 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` is returned. """ return super().__getitem__(index) def _read_flow(self, file_name): return _read_flo(file_name) class KittiFlow(FlowDataset): """`KITTI `__ dataset for optical flow (2015). The dataset is expected to have the following structure: :: root KittiFlow testing image_2 training image_2 flow_occ Args: root (string): Root directory of the KittiFlow Dataset. split (string, optional): The dataset split, either "train" (default) or "test" transforms (callable, optional): A function/transform that takes in ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. """ def __init__(self, root, split="train", transforms=None): super().__init__(root=root, transforms=transforms) verify_str_arg(split, "split", valid_values=("train", "test")) self._has_builtin_flow_mask = True root = Path(root) / "KittiFlow" / (split + "ing") images1 = sorted(glob(str(root / "image_2" / "*_10.png"))) images2 = sorted(glob(str(root / "image_2" / "*_11.png"))) if not images1 or not images2: raise FileNotFoundError( "Could not find the Kitti flow images. Please make sure the directory structure is correct." ) for img1, img2 in zip(images1, images2): self._image_list += [[img1, img2]] if split == "train": self._flow_list = sorted(glob(str(root / "flow_occ" / "*_10.png"))) def __getitem__(self, index): """Return example at given index. Args: index(int): The index of the example to retrieve Returns: tuple: A 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` where ``valid_flow_mask`` is a numpy boolean mask of shape (H, W) indicating which flow values are valid. The flow is a numpy array of shape (2, H, W) and the images are PIL images. ``flow`` and ``valid_flow_mask`` are None if ``split="test"``. """ return super().__getitem__(index) def _read_flow(self, file_name): return _read_16bits_png_with_flow_and_valid_mask(file_name) class FlyingChairs(FlowDataset): """`FlyingChairs `_ Dataset for optical flow. You will also need to download the FlyingChairs_train_val.txt file from the dataset page. The dataset is expected to have the following structure: :: root FlyingChairs data 00001_flow.flo 00001_img1.ppm 00001_img2.ppm ... FlyingChairs_train_val.txt Args: root (string): Root directory of the FlyingChairs Dataset. split (string, optional): The dataset split, either "train" (default) or "val" transforms (callable, optional): A function/transform that takes in ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. """ def __init__(self, root, split="train", transforms=None): super().__init__(root=root, transforms=transforms) verify_str_arg(split, "split", valid_values=("train", "val")) root = Path(root) / "FlyingChairs" images = sorted(glob(str(root / "data" / "*.ppm"))) flows = sorted(glob(str(root / "data" / "*.flo"))) split_file_name = "FlyingChairs_train_val.txt" if not os.path.exists(root / split_file_name): raise FileNotFoundError( "The FlyingChairs_train_val.txt file was not found - please download it from the dataset page (see docstring)." ) split_list = np.loadtxt(str(root / split_file_name), dtype=np.int32) for i in range(len(flows)): split_id = split_list[i] if (split == "train" and split_id == 1) or (split == "val" and split_id == 2): self._flow_list += [flows[i]] self._image_list += [[images[2 * i], images[2 * i + 1]]] def __getitem__(self, index): """Return example at given index. Args: index(int): The index of the example to retrieve Returns: tuple: A 3-tuple with ``(img1, img2, flow)``. The flow is a numpy array of shape (2, H, W) and the images are PIL images. ``flow`` is None if ``split="val"``. If a valid flow mask is generated within the ``transforms`` parameter, a 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` is returned. """ return super().__getitem__(index) def _read_flow(self, file_name): return _read_flo(file_name) class FlyingThings3D(FlowDataset): """`FlyingThings3D `_ dataset for optical flow. The dataset is expected to have the following structure: :: root FlyingThings3D frames_cleanpass TEST TRAIN frames_finalpass TEST TRAIN optical_flow TEST TRAIN Args: root (string): Root directory of the intel FlyingThings3D Dataset. split (string, optional): The dataset split, either "train" (default) or "test" pass_name (string, optional): The pass to use, either "clean" (default) or "final" or "both". See link above for details on the different passes. camera (string, optional): Which camera to return images from. Can be either "left" (default) or "right" or "both". transforms (callable, optional): A function/transform that takes in ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. """ def __init__(self, root, split="train", pass_name="clean", camera="left", transforms=None): super().__init__(root=root, transforms=transforms) verify_str_arg(split, "split", valid_values=("train", "test")) split = split.upper() verify_str_arg(pass_name, "pass_name", valid_values=("clean", "final", "both")) passes = { "clean": ["frames_cleanpass"], "final": ["frames_finalpass"], "both": ["frames_cleanpass", "frames_finalpass"], }[pass_name] verify_str_arg(camera, "camera", valid_values=("left", "right", "both")) cameras = ["left", "right"] if camera == "both" else [camera] root = Path(root) / "FlyingThings3D" directions = ("into_future", "into_past") for pass_name, camera, direction in itertools.product(passes, cameras, directions): image_dirs = sorted(glob(str(root / pass_name / split / "*/*"))) image_dirs = sorted(Path(image_dir) / camera for image_dir in image_dirs) flow_dirs = sorted(glob(str(root / "optical_flow" / split / "*/*"))) flow_dirs = sorted(Path(flow_dir) / direction / camera for flow_dir in flow_dirs) if not image_dirs or not flow_dirs: raise FileNotFoundError( "Could not find the FlyingThings3D flow images. " "Please make sure the directory structure is correct." ) for image_dir, flow_dir in zip(image_dirs, flow_dirs): images = sorted(glob(str(image_dir / "*.png"))) flows = sorted(glob(str(flow_dir / "*.pfm"))) for i in range(len(flows) - 1): if direction == "into_future": self._image_list += [[images[i], images[i + 1]]] self._flow_list += [flows[i]] elif direction == "into_past": self._image_list += [[images[i + 1], images[i]]] self._flow_list += [flows[i + 1]] def __getitem__(self, index): """Return example at given index. Args: index(int): The index of the example to retrieve Returns: tuple: A 3-tuple with ``(img1, img2, flow)``. The flow is a numpy array of shape (2, H, W) and the images are PIL images. ``flow`` is None if ``split="test"``. If a valid flow mask is generated within the ``transforms`` parameter, a 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` is returned. """ return super().__getitem__(index) def _read_flow(self, file_name): return _read_pfm(file_name) class HD1K(FlowDataset): """`HD1K `__ dataset for optical flow. The dataset is expected to have the following structure: :: root hd1k hd1k_challenge image_2 hd1k_flow_gt flow_occ hd1k_input image_2 Args: root (string): Root directory of the HD1K Dataset. split (string, optional): The dataset split, either "train" (default) or "test" transforms (callable, optional): A function/transform that takes in ``img1, img2, flow, valid_flow_mask`` and returns a transformed version. """ def __init__(self, root, split="train", transforms=None): super().__init__(root=root, transforms=transforms) self._has_builtin_flow_mask = True verify_str_arg(split, "split", valid_values=("train", "test")) root = Path(root) / "hd1k" if split == "train": # There are 36 "sequences" and we don't want seq i to overlap with seq i + 1, so we need this for loop for seq_idx in range(36): flows = sorted(glob(str(root / "hd1k_flow_gt" / "flow_occ" / f"{seq_idx:06d}_*.png"))) images = sorted(glob(str(root / "hd1k_input" / "image_2" / f"{seq_idx:06d}_*.png"))) for i in range(len(flows) - 1): self._flow_list += [flows[i]] self._image_list += [[images[i], images[i + 1]]] else: images1 = sorted(glob(str(root / "hd1k_challenge" / "image_2" / "*10.png"))) images2 = sorted(glob(str(root / "hd1k_challenge" / "image_2" / "*11.png"))) for image1, image2 in zip(images1, images2): self._image_list += [[image1, image2]] if not self._image_list: raise FileNotFoundError( "Could not find the HD1K images. Please make sure the directory structure is correct." ) def _read_flow(self, file_name): return _read_16bits_png_with_flow_and_valid_mask(file_name) def __getitem__(self, index): """Return example at given index. Args: index(int): The index of the example to retrieve Returns: tuple: A 4-tuple with ``(img1, img2, flow, valid_flow_mask)`` where ``valid_flow_mask`` is a numpy boolean mask of shape (H, W) indicating which flow values are valid. The flow is a numpy array of shape (2, H, W) and the images are PIL images. ``flow`` and ``valid_flow_mask`` are None if ``split="test"``. """ return super().__getitem__(index) def _read_flo(file_name): """Read .flo file in Middlebury format""" # Code adapted from: # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy # Everything needs to be in little Endian according to # https://vision.middlebury.edu/flow/code/flow-code/README.txt with open(file_name, "rb") as f: magic = np.fromfile(f, "c", count=4).tobytes() if magic != b"PIEH": raise ValueError("Magic number incorrect. Invalid .flo file") w = int(np.fromfile(f, "`_ Dataset. .. warning:: This class needs `scipy `_ to load target files from `.mat` format. Args: root (string): Root directory of dataset where directory ``caltech101`` exists or will be saved to if download is set to True. target_type (string or list, optional): Type of target to use, ``category`` or ``annotation``. Can also be a list to output a tuple with all specified target types. ``category`` represents the target class, and ``annotation`` is a list of points from a hand-generated outline. Defaults to ``category``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ def __init__( self, root: str, target_type: Union[List[str], str] = "category", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(os.path.join(root, "caltech101"), transform=transform, target_transform=target_transform) os.makedirs(self.root, exist_ok=True) if isinstance(target_type, str): target_type = [target_type] self.target_type = [verify_str_arg(t, "target_type", ("category", "annotation")) for t in target_type] if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") self.categories = sorted(os.listdir(os.path.join(self.root, "101_ObjectCategories"))) self.categories.remove("BACKGROUND_Google") # this is not a real class # For some reason, the category names in "101_ObjectCategories" and # "Annotations" do not always match. This is a manual map between the # two. Defaults to using same name, since most names are fine. name_map = { "Faces": "Faces_2", "Faces_easy": "Faces_3", "Motorbikes": "Motorbikes_16", "airplanes": "Airplanes_Side_2", } self.annotation_categories = list(map(lambda x: name_map[x] if x in name_map else x, self.categories)) self.index: List[int] = [] self.y = [] for (i, c) in enumerate(self.categories): n = len(os.listdir(os.path.join(self.root, "101_ObjectCategories", c))) self.index.extend(range(1, n + 1)) self.y.extend(n * [i]) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where the type of target specified by target_type. """ import scipy.io img = Image.open( os.path.join( self.root, "101_ObjectCategories", self.categories[self.y[index]], f"image_{self.index[index]:04d}.jpg", ) ) target: Any = [] for t in self.target_type: if t == "category": target.append(self.y[index]) elif t == "annotation": data = scipy.io.loadmat( os.path.join( self.root, "Annotations", self.annotation_categories[self.y[index]], f"annotation_{self.index[index]:04d}.mat", ) ) target.append(data["obj_contour"]) target = tuple(target) if len(target) > 1 else target[0] if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def _check_integrity(self) -> bool: # can be more robust and check hash of files return os.path.exists(os.path.join(self.root, "101_ObjectCategories")) def __len__(self) -> int: return len(self.index) def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return warnings.warn("The dataset is stored on google drive, if you can't download it from google drive, " "please download it from the official website." "Caltech 101 ") download_and_extract_archive( "https://drive.google.com/file/d/137RyRjvTBkBiIfeYBNZBtViDHQ6_Ewsp", self.root, filename="101_ObjectCategories.tar.gz", md5="b224c7392d521a49829488ab0f1120d9", ) download_and_extract_archive( "https://drive.google.com/file/d/175kQy3UsZ0wUEHZjqkUDdNVssr7bgh_m", self.root, filename="Annotations.tar", md5="6f83eeb1f24d99cab4eb377263132c91", ) def extra_repr(self) -> str: return "Target type: {target_type}".format(**self.__dict__) class Caltech256(VisionDataset): """`Caltech 256 `_ Dataset. Args: root (string): Root directory of dataset where directory ``caltech256`` exists or will be saved to if download is set to True. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ def __init__( self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(os.path.join(root, "caltech256"), transform=transform, target_transform=target_transform) os.makedirs(self.root, exist_ok=True) if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") self.categories = sorted(os.listdir(os.path.join(self.root, "256_ObjectCategories"))) self.index: List[int] = [] self.y = [] for (i, c) in enumerate(self.categories): n = len( [ item for item in os.listdir(os.path.join(self.root, "256_ObjectCategories", c)) if item.endswith(".jpg") ] ) self.index.extend(range(1, n + 1)) self.y.extend(n * [i]) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target class. """ img = Image.open( os.path.join( self.root, "256_ObjectCategories", self.categories[self.y[index]], f"{self.y[index] + 1:03d}_{self.index[index]:04d}.jpg", ) ) target = self.y[index] if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def _check_integrity(self) -> bool: # can be more robust and check hash of files return os.path.exists(os.path.join(self.root, "256_ObjectCategories")) def __len__(self) -> int: return len(self.index) def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return # download_and_extract_archive( # "https://drive.google.com/file/d/1r6o0pSROcV1_VwT4oSjA2FBUSCWGuxLK", # self.root, # filename="256_ObjectCategories.tar", # md5="67b4f42ca05d46448c6bb8ecd2220f6d", # ) download_and_extract_archive( "https://data.caltech.edu/records/nyy15-4j048/files/256_ObjectCategories.tar?download=1", self.root, filename="256_ObjectCategories.tar", md5="67b4f42ca05d46448c6bb8ecd2220f6d", ) msadapter/ms_adapter/torchvision/datasets/celeba.py000066400000000000000000000226411436713554100231370ustar00rootroot00000000000000import csv import os import warnings from collections import namedtuple from typing import Any, Callable, List, Optional, Union, Tuple import PIL import ms_adapter.pytorch as torch import mindspore as ms from .utils import download_file_from_google_drive, check_integrity, verify_str_arg, extract_archive from .vision import VisionDataset CSV = namedtuple("CSV", ["header", "index", "data"]) class CelebA(VisionDataset): """`Large-scale CelebFaces Attributes (CelebA) Dataset `_ Dataset. Args: root (string): Root directory where images are downloaded to. split (string): One of {'train', 'valid', 'test', 'all'}. Accordingly dataset is selected. target_type (string or list, optional): Type of target to use, ``attr``, ``identity``, ``bbox``, or ``landmarks``. Can also be a list to output a tuple with all specified target types. The targets represent: - ``attr`` (np.array shape=(40,) dtype=int): binary (0, 1) labels for attributes - ``identity`` (int): label for each person (data points with the same identity are the same person) - ``bbox`` (np.array shape=(4,) dtype=int): bounding box (x, y, width, height) - ``landmarks`` (np.array shape=(10,) dtype=int): landmark points (lefteye_x, lefteye_y, righteye_x, righteye_y, nose_x, nose_y, leftmouth_x, leftmouth_y, rightmouth_x, rightmouth_y) Defaults to ``attr``. If empty, ``None`` will be returned as target. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ # There currently does not appear to be a easy way to extract 7z in python (without introducing additional # dependencies). The "in-the-wild" (not aligned+cropped) images are only in 7z, so they are not available # right now. def __init__( self, root: str, split: str = "train", target_type: Union[List[str], str] = "attr", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.file_list = [ # File ID MD5 Hash Filename ("0B7EVK8r0v71pZjFTYXZWM3FlRnM", "00d2c5bc6d35e252742224ab0c1e8fcb", "img_align_celeba.zip"), # ("0B7EVK8r0v71pbWNEUjJKdDQ3dGc","b6cd7e93bc7a96c2dc33f819aa3ac651", "img_align_celeba_png.7z"), # ("0B7EVK8r0v71peklHb0pGdDl6R28", "b6cd7e93bc7a96c2dc33f819aa3ac651", "img_celeba.7z"), ("0B7EVK8r0v71pblRyaVFSWGxPY0U", "75e246fa4810816ffd6ee81facbd244c", "list_attr_celeba.txt"), ("1_ee_0u7vcNLOfNLegJRHmolfH5ICW-XS", "32bd1bd63d3c78cd57e08160ec5ed1e2", "identity_CelebA.txt"), ("0B7EVK8r0v71pbThiMVRxWXZ4dU0", "00566efa6fedff7a56946cd1c10f1c16", "list_bbox_celeba.txt"), ("0B7EVK8r0v71pd0FJY3Blby1HUTQ", "cc24ecafdb5b50baae59b03474781f8c", "list_landmarks_align_celeba.txt"), # ("0B7EVK8r0v71pTzJIdlJWdHczRlU", "063ee6ddb681f96bc9ca28c6febb9d1a", "list_landmarks_celeba.txt"), ("0B7EVK8r0v71pY0NSMzRuSXJEVkk", "d32c9cbf5e040fd4025c592c306e6668", "list_eval_partition.txt"), ] self.base_folder = "celeba" self.split = split if isinstance(target_type, list): self.target_type = target_type else: self.target_type = [target_type] if not self.target_type and self.target_transform is not None: raise RuntimeError("target_transform is specified but target_type is empty") if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") split_map = { "train": 0, "valid": 1, "test": 2, "all": None, } split_ = split_map[verify_str_arg(split.lower(), "split", ("train", "valid", "test", "all"))] splits = self._load_csv("list_eval_partition.txt") identity = self._load_csv("identity_CelebA.txt") bbox = self._load_csv("list_bbox_celeba.txt", header=1) landmarks_align = self._load_csv("list_landmarks_align_celeba.txt", header=1) attr = self._load_csv("list_attr_celeba.txt", header=1) self.old_identity = identity self.old_attr = attr mask = slice(None) if split_ is None else (splits.data == split_).squeeze() if mask == slice(None): # if split == "all" self.filename = splits.index else: self.filename = [splits.index[i] for i in torch.squeeze(torch.nonzero(mask))] if mask != slice(None): self.identity = torch.tensor(ms.ops.masked_select(identity.data, ms.ops.expand_dims(mask, -1))).reshape(-1,identity.data.shape[1]) self.bbox = torch.tensor(ms.ops.masked_select(bbox.data,ms.ops.expand_dims(mask, -1))).reshape(-1, bbox.data.shape[1]) self.landmarks_align = torch.tensor(ms.ops.masked_select(landmarks_align.data, ms.ops.expand_dims(mask, -1))).reshape(-1,landmarks_align.data.shape[1]) self.attr = torch.tensor(ms.ops.masked_select(attr.data, ms.ops.expand_dims(mask, -1))).reshape(-1,attr.data.shape[1]) else: self.identity = torch.tensor(identity.data) self.bbox = torch.tensor(bbox.data) self.landmarks_align = torch.tensor(landmarks_align.data) self.attr = torch.tensor(attr.data) # self.identity = identity.data[mask] # self.bbox = bbox.data[mask] # self.landmarks_align = landmarks_align.data[mask] # self.attr = attr.data[mask] # map from {-1, 1} to {0, 1} self.attr = torch.div(self.attr + 1, 2, rounding_mode="floor") self.attr_names = attr.header def _load_csv( self, filename: str, header: Optional[int] = None, ) -> CSV: with open(os.path.join(self.root, self.base_folder, filename)) as csv_file: data = list(csv.reader(csv_file, delimiter=" ", skipinitialspace=True)) if header is not None: headers = data[header] data = data[header + 1 :] else: headers = [] indices = [row[0] for row in data] data = [row[1:] for row in data] data_int = [list(map(int, i)) for i in data] return CSV(headers, indices, torch.tensor(data_int)) def _check_integrity(self) -> bool: for (_, md5, filename) in self.file_list: fpath = os.path.join(self.root, self.base_folder, filename) _, ext = os.path.splitext(filename) # Allow original archive to be deleted (zip and 7z) # Only need the extracted images if ext not in [".zip", ".7z"] and not check_integrity(fpath, md5): return False # Should check a hash of the images return os.path.isdir(os.path.join(self.root, self.base_folder, "img_align_celeba")) def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return warnings.warn("The dataset is stored on google drive, if you can't download it from google drive, " "please download it from the official website." "Large-scale CelebFaces Attributes (CelebA) Dataset ") for (file_id, md5, filename) in self.file_list: download_file_from_google_drive(file_id, os.path.join(self.root, self.base_folder), filename, md5) extract_archive(os.path.join(self.root, self.base_folder, "img_align_celeba.zip")) def __getitem__(self, index: int) -> Tuple[Any, Any]: X = PIL.Image.open(os.path.join(self.root, self.base_folder, "img_align_celeba", self.filename[index])) target: Any = [] for t in self.target_type: if t == "attr": target.append(self.attr[index, :]) elif t == "identity": target.append(self.identity[index, 0]) elif t == "bbox": target.append(self.bbox[index, :]) elif t == "landmarks": target.append(self.landmarks_align[index, :]) else: # TODO: refactor with utils.verify_str_arg raise ValueError(f'Target type "{t}" is not recognized.') if self.transform is not None: X = self.transform(X) if target: target = tuple(target) if len(target) > 1 else target[0] if self.target_transform is not None: target = self.target_transform(target) else: target = None return X, target def __len__(self) -> int: return len(self.attr) def extra_repr(self) -> str: lines = ["Target type: {target_type}", "Split: {split}"] return "\n".join(lines).format(**self.__dict__) msadapter/ms_adapter/torchvision/datasets/cifar.py000066400000000000000000000133331436713554100230060ustar00rootroot00000000000000import os.path import pickle from typing import Any, Callable, Optional, Tuple import numpy as np from PIL import Image from .utils import check_integrity, download_and_extract_archive from .vision import VisionDataset class CIFAR10(VisionDataset): """`CIFAR10 `_ Dataset. Args: root (string): Root directory of dataset where directory ``cifar-10-batches-py`` exists or will be saved to if download is set to True. train (bool, optional): If True, creates dataset from training set, otherwise creates from test set. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ base_folder = "cifar-10-batches-py" url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz" filename = "cifar-10-python.tar.gz" tgz_md5 = "c58f30108f718f92721af3b95e74349a" train_list = [ ["data_batch_1", "c99cafc152244af753f735de768cd75f"], ["data_batch_2", "d4bba439e000b95fd0a9bffe97cbabec"], ["data_batch_3", "54ebc095f3ab1f0389bbae665268c751"], ["data_batch_4", "634d18415352ddfa80567beed471001a"], ["data_batch_5", "482c414d41f54cd18b22e5b47cb7c3cb"], ] test_list = [ ["test_batch", "40351d587109b95175f43aff81a1287e"], ] meta = { "filename": "batches.meta", "key": "label_names", "md5": "5ff9c542aee3614f3951f8cda6e48888", } def __init__( self, root: str, train: bool = True, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.train = train # training set or test set if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") if self.train: downloaded_list = self.train_list else: downloaded_list = self.test_list self.data: Any = [] self.targets = [] # now load the picked numpy arrays for file_name, checksum in downloaded_list: file_path = os.path.join(self.root, self.base_folder, file_name) with open(file_path, "rb") as f: entry = pickle.load(f, encoding="latin1") self.data.append(entry["data"]) if "labels" in entry: self.targets.extend(entry["labels"]) else: self.targets.extend(entry["fine_labels"]) self.data = np.vstack(self.data).reshape(-1, 3, 32, 32) self.data = self.data.transpose((0, 2, 3, 1)) # convert to HWC self._load_meta() def _load_meta(self) -> None: path = os.path.join(self.root, self.base_folder, self.meta["filename"]) if not check_integrity(path, self.meta["md5"]): raise RuntimeError("Dataset metadata file not found or corrupted. You can use download=True to download it") with open(path, "rb") as infile: data = pickle.load(infile, encoding="latin1") self.classes = data[self.meta["key"]] self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)} def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target class. """ img, target = self.data[index], self.targets[index] # doing this so that it is consistent with all other datasets # to return a PIL Image img = Image.fromarray(img) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.data) def _check_integrity(self) -> bool: root = self.root for fentry in self.train_list + self.test_list: filename, md5 = fentry[0], fentry[1] fpath = os.path.join(root, self.base_folder, filename) if not check_integrity(fpath, md5): return False return True def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5) def extra_repr(self) -> str: split = "Train" if self.train is True else "Test" return f"Split: {split}" class CIFAR100(CIFAR10): """`CIFAR100 `_ Dataset. This is a subclass of the `CIFAR10` Dataset. """ base_folder = "cifar-100-python" url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz" filename = "cifar-100-python.tar.gz" tgz_md5 = "eb9058c3a382ffc7106e4002c42a8d85" train_list = [ ["train", "16019d7e3df5f24257cddd939b257f8d"], ] test_list = [ ["test", "f0ef6b0ae62326f3e7ffdfab6717acfc"], ] meta = { "filename": "meta", "key": "fine_label_names", "md5": "7973b15100ade9c7d40fb424638fde48", } msadapter/ms_adapter/torchvision/datasets/cityscapes.py000066400000000000000000000237751436713554100241040ustar00rootroot00000000000000import json import os from collections import namedtuple from typing import Any, Callable, Dict, List, Optional, Union, Tuple from PIL import Image from .utils import extract_archive, verify_str_arg, iterable_to_str from .vision import VisionDataset class Cityscapes(VisionDataset): """`Cityscapes `_ Dataset. Args: root (string): Root directory of dataset where directory ``leftImg8bit`` and ``gtFine`` or ``gtCoarse`` are located. split (string, optional): The image split to use, ``train``, ``test`` or ``val`` if mode="fine" otherwise ``train``, ``train_extra`` or ``val`` mode (string, optional): The quality mode to use, ``fine`` or ``coarse`` target_type (string or list, optional): Type of target to use, ``instance``, ``semantic``, ``polygon`` or ``color``. Can also be a list to output a tuple with all specified target types. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. Examples: Get semantic segmentation target .. code-block:: python dataset = Cityscapes('./data/cityscapes', split='train', mode='fine', target_type='semantic') img, smnt = dataset[0] Get multiple targets .. code-block:: python dataset = Cityscapes('./data/cityscapes', split='train', mode='fine', target_type=['instance', 'color', 'polygon']) img, (inst, col, poly) = dataset[0] Validate on the "coarse" set .. code-block:: python dataset = Cityscapes('./data/cityscapes', split='val', mode='coarse', target_type='semantic') img, smnt = dataset[0] """ # Based on https://github.com/mcordts/cityscapesScripts CityscapesClass = namedtuple( "CityscapesClass", ["name", "id", "train_id", "category", "category_id", "has_instances", "ignore_in_eval", "color"], ) classes = [ CityscapesClass("unlabeled", 0, 255, "void", 0, False, True, (0, 0, 0)), CityscapesClass("ego vehicle", 1, 255, "void", 0, False, True, (0, 0, 0)), CityscapesClass("rectification border", 2, 255, "void", 0, False, True, (0, 0, 0)), CityscapesClass("out of roi", 3, 255, "void", 0, False, True, (0, 0, 0)), CityscapesClass("static", 4, 255, "void", 0, False, True, (0, 0, 0)), CityscapesClass("dynamic", 5, 255, "void", 0, False, True, (111, 74, 0)), CityscapesClass("ground", 6, 255, "void", 0, False, True, (81, 0, 81)), CityscapesClass("road", 7, 0, "flat", 1, False, False, (128, 64, 128)), CityscapesClass("sidewalk", 8, 1, "flat", 1, False, False, (244, 35, 232)), CityscapesClass("parking", 9, 255, "flat", 1, False, True, (250, 170, 160)), CityscapesClass("rail track", 10, 255, "flat", 1, False, True, (230, 150, 140)), CityscapesClass("building", 11, 2, "construction", 2, False, False, (70, 70, 70)), CityscapesClass("wall", 12, 3, "construction", 2, False, False, (102, 102, 156)), CityscapesClass("fence", 13, 4, "construction", 2, False, False, (190, 153, 153)), CityscapesClass("guard rail", 14, 255, "construction", 2, False, True, (180, 165, 180)), CityscapesClass("bridge", 15, 255, "construction", 2, False, True, (150, 100, 100)), CityscapesClass("tunnel", 16, 255, "construction", 2, False, True, (150, 120, 90)), CityscapesClass("pole", 17, 5, "object", 3, False, False, (153, 153, 153)), CityscapesClass("polegroup", 18, 255, "object", 3, False, True, (153, 153, 153)), CityscapesClass("traffic light", 19, 6, "object", 3, False, False, (250, 170, 30)), CityscapesClass("traffic sign", 20, 7, "object", 3, False, False, (220, 220, 0)), CityscapesClass("vegetation", 21, 8, "nature", 4, False, False, (107, 142, 35)), CityscapesClass("terrain", 22, 9, "nature", 4, False, False, (152, 251, 152)), CityscapesClass("sky", 23, 10, "sky", 5, False, False, (70, 130, 180)), CityscapesClass("person", 24, 11, "human", 6, True, False, (220, 20, 60)), CityscapesClass("rider", 25, 12, "human", 6, True, False, (255, 0, 0)), CityscapesClass("car", 26, 13, "vehicle", 7, True, False, (0, 0, 142)), CityscapesClass("truck", 27, 14, "vehicle", 7, True, False, (0, 0, 70)), CityscapesClass("bus", 28, 15, "vehicle", 7, True, False, (0, 60, 100)), CityscapesClass("caravan", 29, 255, "vehicle", 7, True, True, (0, 0, 90)), CityscapesClass("trailer", 30, 255, "vehicle", 7, True, True, (0, 0, 110)), CityscapesClass("train", 31, 16, "vehicle", 7, True, False, (0, 80, 100)), CityscapesClass("motorcycle", 32, 17, "vehicle", 7, True, False, (0, 0, 230)), CityscapesClass("bicycle", 33, 18, "vehicle", 7, True, False, (119, 11, 32)), CityscapesClass("license plate", -1, -1, "vehicle", 7, False, True, (0, 0, 142)), ] def __init__( self, root: str, split: str = "train", mode: str = "fine", target_type: Union[List[str], str] = "instance", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, transforms: Optional[Callable] = None, ) -> None: super().__init__(root, transforms, transform, target_transform) self.mode = "gtFine" if mode == "fine" else "gtCoarse" self.images_dir = os.path.join(self.root, "leftImg8bit", split) self.targets_dir = os.path.join(self.root, self.mode, split) self.target_type = target_type self.split = split self.images = [] self.targets = [] verify_str_arg(mode, "mode", ("fine", "coarse")) if mode == "fine": valid_modes = ("train", "test", "val") else: valid_modes = ("train", "train_extra", "val") msg = "Unknown value '{}' for argument split if mode is '{}'. Valid values are {{{}}}." msg = msg.format(split, mode, iterable_to_str(valid_modes)) verify_str_arg(split, "split", valid_modes, msg) if not isinstance(target_type, list): self.target_type = [target_type] [ verify_str_arg(value, "target_type", ("instance", "semantic", "polygon", "color")) for value in self.target_type ] if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir): if split == "train_extra": image_dir_zip = os.path.join(self.root, "leftImg8bit_trainextra.zip") else: image_dir_zip = os.path.join(self.root, "leftImg8bit_trainvaltest.zip") if self.mode == "gtFine": target_dir_zip = os.path.join(self.root, f"{self.mode}_trainvaltest.zip") elif self.mode == "gtCoarse": target_dir_zip = os.path.join(self.root, f"{self.mode}.zip") if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip): extract_archive(from_path=image_dir_zip, to_path=self.root) extract_archive(from_path=target_dir_zip, to_path=self.root) else: raise RuntimeError( "Dataset not found or incomplete. Please make sure all required folders for the" ' specified "split" and "mode" are inside the "root" directory' ) for city in os.listdir(self.images_dir): img_dir = os.path.join(self.images_dir, city) target_dir = os.path.join(self.targets_dir, city) for file_name in os.listdir(img_dir): target_types = [] for t in self.target_type: target_name = "{}_{}".format( file_name.split("_leftImg8bit")[0], self._get_target_suffix(self.mode, t) ) target_types.append(os.path.join(target_dir, target_name)) self.images.append(os.path.join(img_dir, file_name)) self.targets.append(target_types) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is a tuple of all target types if target_type is a list with more than one item. Otherwise target is a json object if target_type="polygon", else the image segmentation. """ image = Image.open(self.images[index]).convert("RGB") targets: Any = [] for i, t in enumerate(self.target_type): if t == "polygon": target = self._load_json(self.targets[index][i]) else: target = Image.open(self.targets[index][i]) targets.append(target) target = tuple(targets) if len(targets) > 1 else targets[0] if self.transforms is not None: image, target = self.transforms(image, target) return image, target def __len__(self) -> int: return len(self.images) def extra_repr(self) -> str: lines = ["Split: {split}", "Mode: {mode}", "Type: {target_type}"] return "\n".join(lines).format(**self.__dict__) def _load_json(self, path: str) -> Dict[str, Any]: with open(path) as file: data = json.load(file) return data def _get_target_suffix(self, mode: str, target_type: str) -> str: if target_type == "instance": return f"{mode}_instanceIds.png" elif target_type == "semantic": return f"{mode}_labelIds.png" elif target_type == "color": return f"{mode}_color.png" else: return f"{mode}_polygons.json" msadapter/ms_adapter/torchvision/datasets/clevr.py000066400000000000000000000065301436713554100230360ustar00rootroot00000000000000import json import pathlib from typing import Any, Callable, Optional, Tuple, List from urllib.parse import urlparse from PIL import Image from .utils import download_and_extract_archive, verify_str_arg from .vision import VisionDataset class CLEVRClassification(VisionDataset): """`CLEVR `_ classification dataset. The number of objects in a scene are used as label. Args: root (string): Root directory of dataset where directory ``root/clevr`` exists or will be saved to if download is set to True. split (string, optional): The dataset split, supports ``"train"`` (default), ``"val"``, or ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in them target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ _URL = "https://dl.fbaipublicfiles.com/clevr/CLEVR_v1.0.zip" _MD5 = "b11922020e72d0cd9154779b2d3d07d2" def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: self._split = verify_str_arg(split, "split", ("train", "val", "test")) super().__init__(root, transform=transform, target_transform=target_transform) self._base_folder = pathlib.Path(self.root) / "clevr" self._data_folder = self._base_folder / pathlib.Path(urlparse(self._URL).path).stem if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") self._image_files = sorted(self._data_folder.joinpath("images", self._split).glob("*")) self._labels: List[Optional[int]] if self._split != "test": with open(self._data_folder / "scenes" / f"CLEVR_{self._split}_scenes.json") as file: content = json.load(file) num_objects = {scene["image_filename"]: len(scene["objects"]) for scene in content["scenes"]} self._labels = [num_objects[image_file.name] for image_file in self._image_files] else: self._labels = [None] * len(self._image_files) def __len__(self) -> int: return len(self._image_files) def __getitem__(self, idx: int) -> Tuple[Any, Any]: image_file = self._image_files[idx] label = self._labels[idx] image = Image.open(image_file).convert("RGB") if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def _check_exists(self) -> bool: return self._data_folder.exists() and self._data_folder.is_dir() def _download(self) -> None: if self._check_exists(): return download_and_extract_archive(self._URL, str(self._base_folder), md5=self._MD5) def extra_repr(self) -> str: return f"split={self._split}" msadapter/ms_adapter/torchvision/datasets/coco.py000066400000000000000000000076041436713554100226510ustar00rootroot00000000000000import os.path from typing import Any, Callable, Optional, Tuple, List from PIL import Image from .vision import VisionDataset class CocoDetection(VisionDataset): """`MS Coco Detection `_ Dataset. It requires the `COCO API to be installed `_. Args: root (string): Root directory where images are downloaded to. annFile (string): Path to json annotation file. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. """ def __init__( self, root: str, annFile: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, transforms: Optional[Callable] = None, ) -> None: super().__init__(root, transforms, transform, target_transform) from pycocotools.coco import COCO self.coco = COCO(annFile) self.ids = list(sorted(self.coco.imgs.keys())) def _load_image(self, id: int) -> Image.Image: path = self.coco.loadImgs(id)[0]["file_name"] return Image.open(os.path.join(self.root, path)).convert("RGB") def _load_target(self, id: int) -> List[Any]: return self.coco.loadAnns(self.coco.getAnnIds(id)) def __getitem__(self, index: int) -> Tuple[Any, Any]: id = self.ids[index] image = self._load_image(id) target = self._load_target(id) if self.transforms is not None: image, target = self.transforms(image, target) return image, target def __len__(self) -> int: return len(self.ids) class CocoCaptions(CocoDetection): """`MS Coco Captions `_ Dataset. It requires the `COCO API to be installed `_. Args: root (string): Root directory where images are downloaded to. annFile (string): Path to json annotation file. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. Example: .. code:: python import torchvision.datasets as dset import torchvision.transforms as transforms cap = dset.CocoCaptions(root = 'dir where images are', annFile = 'json annotation file', transform=transforms.PILToTensor()) print('Number of samples: ', len(cap)) img, target = cap[3] # load 4th sample print("Image Size: ", img.size()) print(target) Output: :: Number of samples: 82783 Image Size: (3L, 427L, 640L) [u'A plane emitting smoke stream flying over a mountain.', u'A plane darts across a bright blue sky behind a mountain covered in snow', u'A plane leaves a contrail above the snowy mountain top.', u'A mountain that has a plane flying overheard in the distance.', u'A mountain view with a plume of smoke in the background'] """ def _load_target(self, id: int) -> List[str]: return [ann["caption"] for ann in super()._load_target(id)] msadapter/ms_adapter/torchvision/datasets/country211.py000066400000000000000000000045501436713554100236520ustar00rootroot00000000000000from pathlib import Path from typing import Callable, Optional from .folder import ImageFolder from .utils import verify_str_arg, download_and_extract_archive class Country211(ImageFolder): """`The Country211 Data Set `_ from OpenAI. This dataset was built by filtering the images from the YFCC100m dataset that have GPS coordinate corresponding to a ISO-3166 country code. The dataset is balanced by sampling 150 train images, 50 validation images, and 100 test images images for each country. Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"train"`` (default), ``"valid"`` and ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it into ``root/country211/``. If dataset is already downloaded, it is not downloaded again. """ _URL = "https://openaipublic.azureedge.net/clip/data/country211.tgz" _MD5 = "84988d7644798601126c29e9877aab6a" def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: self._split = verify_str_arg(split, "split", ("train", "valid", "test")) root = Path(root).expanduser() self.root = str(root) self._base_folder = root / "country211" if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") super().__init__(str(self._base_folder / self._split), transform=transform, target_transform=target_transform) self.root = str(root) def _check_exists(self) -> bool: return self._base_folder.exists() and self._base_folder.is_dir() def _download(self) -> None: if self._check_exists(): return download_and_extract_archive(self._URL, download_root=self.root, md5=self._MD5) msadapter/ms_adapter/torchvision/datasets/dtd.py000066400000000000000000000075431436713554100225030ustar00rootroot00000000000000import os import pathlib from typing import Optional, Callable import PIL.Image from .utils import verify_str_arg, download_and_extract_archive from .vision import VisionDataset class DTD(VisionDataset): """`Describable Textures Dataset (DTD) `_. Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"train"`` (default), ``"val"``, or ``"test"``. partition (int, optional): The dataset partition. Should be ``1 <= partition <= 10``. Defaults to ``1``. .. note:: The partition only changes which split each image belongs to. Thus, regardless of the selected partition, combining all splits will result in all images. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. Default is False. """ _URL = "https://www.robots.ox.ac.uk/~vgg/data/dtd/download/dtd-r1.0.1.tar.gz" _MD5 = "fff73e5086ae6bdbea199a49dfb8a4c1" def __init__( self, root: str, split: str = "train", partition: int = 1, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: self._split = verify_str_arg(split, "split", ("train", "val", "test")) if not isinstance(partition, int) and not (1 <= partition <= 10): raise ValueError( f"Parameter 'partition' should be an integer with `1 <= partition <= 10`, " f"but got {partition} instead" ) self._partition = partition super().__init__(root, transform=transform, target_transform=target_transform) self._base_folder = pathlib.Path(self.root) / type(self).__name__.lower() self._data_folder = self._base_folder / "dtd" self._meta_folder = self._data_folder / "labels" self._images_folder = self._data_folder / "images" if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") self._image_files = [] classes = [] with open(self._meta_folder / f"{self._split}{self._partition}.txt") as file: for line in file: cls, name = line.strip().split("/") self._image_files.append(self._images_folder.joinpath(cls, name)) classes.append(cls) self.classes = sorted(set(classes)) self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) self._labels = [self.class_to_idx[cls] for cls in classes] def __len__(self) -> int: return len(self._image_files) def __getitem__(self, idx): image_file, label = self._image_files[idx], self._labels[idx] image = PIL.Image.open(image_file).convert("RGB") if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def extra_repr(self) -> str: return f"split={self._split}, partition={self._partition}" def _check_exists(self) -> bool: return os.path.exists(self._data_folder) and os.path.isdir(self._data_folder) def _download(self) -> None: if self._check_exists(): return download_and_extract_archive(self._URL, download_root=str(self._base_folder), md5=self._MD5) msadapter/ms_adapter/torchvision/datasets/eurosat.py000066400000000000000000000040051436713554100234000ustar00rootroot00000000000000import os from typing import Callable, Optional from .folder import ImageFolder from .utils import download_and_extract_archive class EuroSAT(ImageFolder): """RGB version of the `EuroSAT `_ Dataset. Args: root (string): Root directory of dataset where ``root/eurosat`` exists. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. Default is False. """ def __init__( self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: self.root = os.path.expanduser(root) self._base_folder = os.path.join(self.root, "eurosat") self._data_folder = os.path.join(self._base_folder, "2750") if download: self.download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") super().__init__(self._data_folder, transform=transform, target_transform=target_transform) self.root = os.path.expanduser(root) def __len__(self) -> int: return len(self.samples) def _check_exists(self) -> bool: return os.path.exists(self._data_folder) def download(self) -> None: if self._check_exists(): return os.makedirs(self._base_folder, exist_ok=True) download_and_extract_archive( "https://madm.dfki.de/files/sentinel/EuroSAT.zip", download_root=self._base_folder, md5="c8fa014336c82ac7804f0398fcb19387", ) msadapter/ms_adapter/torchvision/datasets/fakedata.py000066400000000000000000000050621436713554100234620ustar00rootroot00000000000000from typing import Any, Callable, Optional, Tuple import numpy as np import ms_adapter.pytorch as torch from ms_adapter.torchvision import transforms from ms_adapter.torchvision.datasets.vision import VisionDataset class FakeData(VisionDataset): """A fake dataset that returns randomly generated images and returns them as PIL images Args: size (int, optional): Size of the dataset. Default: 1000 images image_size(tuple, optional): Size if the returned images. Default: (3, 224, 224) num_classes(int, optional): Number of classes in the dataset. Default: 10 transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. random_offset (int): Offsets the index-based random seed used to generate each image. Default: 0 """ def __init__( self, size: int = 1000, image_size: Tuple[int, int, int] = (224, 224, 3), num_classes: int = 10, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, random_offset: int = 0, ) -> None: super().__init__(None, transform=transform, target_transform=target_transform) # type: ignore[arg-type] self.size = size self.num_classes = num_classes self.image_size = image_size self.random_offset = random_offset def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is class_index of the target class. """ # create random image that is consistent with the index id if index >= len(self): raise IndexError(f"{self.__class__.__name__} index out of range") rng = np.random.RandomState() rng_state = rng.get_state() rng.seed(index + self.random_offset) img = np.asarray(rng.randn(*self.image_size), dtype=np.uint8) target = rng.randint(0, self.num_classes, size=(1,), dtype=np.long)[0] rng.set_state(rng_state) # convert to PIL Image img = transforms.ToPILImage()(img) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target.item() def __len__(self) -> int: return self.size msadapter/ms_adapter/torchvision/datasets/fer2013.py000066400000000000000000000053401436713554100230030ustar00rootroot00000000000000import csv import pathlib from typing import Any, Callable, Optional, Tuple import ms_adapter.pytorch as torch from PIL import Image from .utils import verify_str_arg, check_integrity from .vision import VisionDataset class FER2013(VisionDataset): """`FER2013 `_ Dataset. Args: root (string): Root directory of dataset where directory ``root/fer2013`` exists. split (string, optional): The dataset split, supports ``"train"`` (default), or ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ _RESOURCES = { "train": ("train.csv", "3f0dfb3d3fd99c811a1299cb947e3131"), "test": ("test.csv", "b02c2298636a634e8c2faabbf3ea9a23"), } def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, ) -> None: self._split = verify_str_arg(split, "split", self._RESOURCES.keys()) super().__init__(root, transform=transform, target_transform=target_transform) base_folder = pathlib.Path(self.root) / "fer2013" file_name, md5 = self._RESOURCES[self._split] data_file = base_folder / file_name if not check_integrity(str(data_file), md5=md5): raise RuntimeError( f"{file_name} not found in {base_folder} or corrupted. " f"You can download it from " f"https://www.kaggle.com/c/challenges-in-representation-learning-facial-expression-recognition-challenge" ) with open(data_file, "r", newline="") as file: self._samples = [ ( torch.tensor([int(idx) for idx in row["pixels"].split()], dtype=torch.uint8).reshape(48, 48), int(row["emotion"]) if "emotion" in row else None, ) for row in csv.DictReader(file) ] def __len__(self) -> int: return len(self._samples) def __getitem__(self, idx: int) -> Tuple[Any, Any]: image_tensor, target = self._samples[idx] image = Image.fromarray(image_tensor.numpy()) if self.transform is not None: image = self.transform(image) if self.target_transform is not None: target = self.target_transform(target) return image, target def extra_repr(self) -> str: return f"split={self._split}" msadapter/ms_adapter/torchvision/datasets/fgvc_aircraft.py000066400000000000000000000107211436713554100245200ustar00rootroot00000000000000from __future__ import annotations import os from typing import Any, Callable, Optional, Tuple import PIL.Image from .utils import download_and_extract_archive, verify_str_arg from .vision import VisionDataset class FGVCAircraft(VisionDataset): """`FGVC Aircraft `_ Dataset. The dataset contains 10,000 images of aircraft, with 100 images for each of 100 different aircraft model variants, most of which are airplanes. Aircraft models are organized in a three-levels hierarchy. The three levels, from finer to coarser, are: - ``variant``, e.g. Boeing 737-700. A variant collapses all the models that are visually indistinguishable into one class. The dataset comprises 100 different variants. - ``family``, e.g. Boeing 737. The dataset comprises 70 different families. - ``manufacturer``, e.g. Boeing. The dataset comprises 30 different manufacturers. Args: root (string): Root directory of the FGVC Aircraft dataset. split (string, optional): The dataset split, supports ``train``, ``val``, ``trainval`` and ``test``. annotation_level (str, optional): The annotation level, supports ``variant``, ``family`` and ``manufacturer``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ _URL = "https://www.robots.ox.ac.uk/~vgg/data/fgvc-aircraft/archives/fgvc-aircraft-2013b.tar.gz" def __init__( self, root: str, split: str = "trainval", annotation_level: str = "variant", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self._split = verify_str_arg(split, "split", ("train", "val", "trainval", "test")) self._annotation_level = verify_str_arg( annotation_level, "annotation_level", ("variant", "family", "manufacturer") ) self._data_path = os.path.join(self.root, "fgvc-aircraft-2013b") if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") annotation_file = os.path.join( self._data_path, "data", { "variant": "variants.txt", "family": "families.txt", "manufacturer": "manufacturers.txt", }[self._annotation_level], ) with open(annotation_file, "r") as f: self.classes = [line.strip() for line in f] self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) image_data_folder = os.path.join(self._data_path, "data", "images") labels_file = os.path.join(self._data_path, "data", f"images_{self._annotation_level}_{self._split}.txt") self._image_files = [] self._labels = [] with open(labels_file, "r") as f: for line in f: image_name, label_name = line.strip().split(" ", 1) self._image_files.append(os.path.join(image_data_folder, f"{image_name}.jpg")) self._labels.append(self.class_to_idx[label_name]) def __len__(self) -> int: return len(self._image_files) def __getitem__(self, idx) -> Tuple[Any, Any]: image_file, label = self._image_files[idx], self._labels[idx] image = PIL.Image.open(image_file).convert("RGB") if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def _download(self) -> None: """ Download the FGVC Aircraft dataset archive and extract it under root. """ if self._check_exists(): return download_and_extract_archive(self._URL, self.root) def _check_exists(self) -> bool: return os.path.exists(self._data_path) and os.path.isdir(self._data_path) msadapter/ms_adapter/torchvision/datasets/flickr.py000066400000000000000000000123331436713554100231730ustar00rootroot00000000000000import glob import os from collections import defaultdict from html.parser import HTMLParser from typing import Any, Callable, Dict, List, Optional, Tuple from PIL import Image from .vision import VisionDataset class Flickr8kParser(HTMLParser): """Parser for extracting captions from the Flickr8k dataset web page.""" def __init__(self, root: str) -> None: super().__init__() self.root = root # Data structure to store captions self.annotations: Dict[str, List[str]] = {} # State variables self.in_table = False self.current_tag: Optional[str] = None self.current_img: Optional[str] = None def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]) -> None: self.current_tag = tag if tag == "table": self.in_table = True def handle_endtag(self, tag: str) -> None: self.current_tag = None if tag == "table": self.in_table = False def handle_data(self, data: str) -> None: if self.in_table: if data == "Image Not Found": self.current_img = None elif self.current_tag == "a": img_id = data.split("/")[-2] img_id = os.path.join(self.root, img_id + "_*.jpg") img_id = glob.glob(img_id)[0] self.current_img = img_id self.annotations[img_id] = [] elif self.current_tag == "li" and self.current_img: img_id = self.current_img self.annotations[img_id].append(data.strip()) class Flickr8k(VisionDataset): """`Flickr8k Entities `_ Dataset. Args: root (string): Root directory where images are downloaded to. ann_file (string): Path to annotation file. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ def __init__( self, root: str, ann_file: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.ann_file = os.path.expanduser(ann_file) # Read annotations and store in a dict parser = Flickr8kParser(self.root) with open(self.ann_file) as fh: parser.feed(fh.read()) self.annotations = parser.annotations self.ids = list(sorted(self.annotations.keys())) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: Tuple (image, target). target is a list of captions for the image. """ img_id = self.ids[index] # Image img = Image.open(img_id).convert("RGB") if self.transform is not None: img = self.transform(img) # Captions target = self.annotations[img_id] if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.ids) class Flickr30k(VisionDataset): """`Flickr30k Entities `_ Dataset. Args: root (string): Root directory where images are downloaded to. ann_file (string): Path to annotation file. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ def __init__( self, root: str, ann_file: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.ann_file = os.path.expanduser(ann_file) # Read annotations and store in a dict self.annotations = defaultdict(list) with open(self.ann_file) as fh: for line in fh: img_id, caption = line.strip().split("\t") self.annotations[img_id[:-2]].append(caption) self.ids = list(sorted(self.annotations.keys())) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: Tuple (image, target). target is a list of captions for the image. """ img_id = self.ids[index] # Image filename = os.path.join(self.root, img_id) img = Image.open(filename).convert("RGB") if self.transform is not None: img = self.transform(img) # Captions target = self.annotations[img_id] if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.ids) msadapter/ms_adapter/torchvision/datasets/flowers102.py000066400000000000000000000107701436713554100236300ustar00rootroot00000000000000from pathlib import Path from typing import Any, Tuple, Callable, Optional import PIL.Image from .utils import check_integrity, download_and_extract_archive, download_url, verify_str_arg from .vision import VisionDataset class Flowers102(VisionDataset): """`Oxford 102 Flower `_ Dataset. .. warning:: This class needs `scipy `_ to load target files from `.mat` format. Oxford 102 Flower is an image classification dataset consisting of 102 flower categories. The flowers were chosen to be flowers commonly occurring in the United Kingdom. Each class consists of between 40 and 258 images. The images have large scale, pose and light variations. In addition, there are categories that have large variations within the category, and several very similar categories. Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"train"`` (default), ``"val"``, or ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ _download_url_prefix = "https://www.robots.ox.ac.uk/~vgg/data/flowers/102/" _file_dict = { # filename, md5 "image": ("102flowers.tgz", "52808999861908f626f3c1f4e79d11fa"), "label": ("imagelabels.mat", "e0620be6f572b9609742df49c70aed4d"), "setid": ("setid.mat", "a5357ecc9cb78c4bef273ce3793fc85c"), } _splits_map = {"train": "trnid", "val": "valid", "test": "tstid"} def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self._split = verify_str_arg(split, "split", ("train", "val", "test")) self._base_folder = Path(self.root) / "flowers-102" self._images_folder = self._base_folder / "jpg" if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") from scipy.io import loadmat set_ids = loadmat(self._base_folder / self._file_dict["setid"][0], squeeze_me=True) image_ids = set_ids[self._splits_map[self._split]].tolist() labels = loadmat(self._base_folder / self._file_dict["label"][0], squeeze_me=True) image_id_to_label = dict(enumerate((labels["labels"] - 1).tolist(), 1)) self._labels = [] self._image_files = [] for image_id in image_ids: self._labels.append(image_id_to_label[image_id]) self._image_files.append(self._images_folder / f"image_{image_id:05d}.jpg") def __len__(self) -> int: return len(self._image_files) def __getitem__(self, idx) -> Tuple[Any, Any]: image_file, label = self._image_files[idx], self._labels[idx] image = PIL.Image.open(image_file).convert("RGB") if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def extra_repr(self) -> str: return f"split={self._split}" def _check_integrity(self): if not (self._images_folder.exists() and self._images_folder.is_dir()): return False for id in ["label", "setid"]: filename, md5 = self._file_dict[id] if not check_integrity(str(self._base_folder / filename), md5): return False return True def download(self): if self._check_integrity(): return download_and_extract_archive( f"{self._download_url_prefix}{self._file_dict['image'][0]}", str(self._base_folder), md5=self._file_dict["image"][1], ) for id in ["label", "setid"]: filename, md5 = self._file_dict[id] download_url(self._download_url_prefix + filename, str(self._base_folder), md5=md5) msadapter/ms_adapter/torchvision/datasets/folder.py000066400000000000000000000272431436713554100232020ustar00rootroot00000000000000import os import os.path from typing import Any, Callable, cast, Dict, List, Optional, Tuple from typing import Union from PIL import Image from .vision import VisionDataset def has_file_allowed_extension(filename: str, extensions: Union[str, Tuple[str, ...]]) -> bool: """Checks if a file is an allowed extension. Args: filename (string): path to a file extensions (tuple of strings): extensions to consider (lowercase) Returns: bool: True if the filename ends with one of given extensions """ return filename.lower().endswith(extensions if isinstance(extensions, str) else tuple(extensions)) def is_image_file(filename: str) -> bool: """Checks if a file is an allowed image extension. Args: filename (string): path to a file Returns: bool: True if the filename ends with a known image extension """ return has_file_allowed_extension(filename, IMG_EXTENSIONS) def find_classes(directory: str) -> Tuple[List[str], Dict[str, int]]: """Finds the class folders in a dataset. See :class:`DatasetFolder` for details. """ classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir()) if not classes: raise FileNotFoundError(f"Couldn't find any class folder in {directory}.") class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)} return classes, class_to_idx def make_dataset( directory: str, class_to_idx: Optional[Dict[str, int]] = None, extensions: Optional[Union[str, Tuple[str, ...]]] = None, is_valid_file: Optional[Callable[[str], bool]] = None, ) -> List[Tuple[str, int]]: """Generates a list of samples of a form (path_to_sample, class). See :class:`DatasetFolder` for details. Note: The class_to_idx parameter is here optional and will use the logic of the ``find_classes`` function by default. """ directory = os.path.expanduser(directory) if class_to_idx is None: _, class_to_idx = find_classes(directory) elif not class_to_idx: raise ValueError("'class_to_index' must have at least one entry to collect any samples.") both_none = extensions is None and is_valid_file is None both_something = extensions is not None and is_valid_file is not None if both_none or both_something: raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time") if extensions is not None: def is_valid_file(x: str) -> bool: return has_file_allowed_extension(x, extensions) # type: ignore[arg-type] is_valid_file = cast(Callable[[str], bool], is_valid_file) instances = [] available_classes = set() for target_class in sorted(class_to_idx.keys()): class_index = class_to_idx[target_class] target_dir = os.path.join(directory, target_class) if not os.path.isdir(target_dir): continue for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)): for fname in sorted(fnames): path = os.path.join(root, fname) if is_valid_file(path): item = path, class_index instances.append(item) if target_class not in available_classes: available_classes.add(target_class) empty_classes = set(class_to_idx.keys()) - available_classes if empty_classes: msg = f"Found no valid file for the classes {', '.join(sorted(empty_classes))}. " if extensions is not None: msg += f"Supported extensions are: {extensions if isinstance(extensions, str) else ', '.join(extensions)}" raise FileNotFoundError(msg) return instances class DatasetFolder(VisionDataset): """A generic data loader. This default directory structure can be customized by overriding the :meth:`find_classes` method. Args: root (string): Root directory path. loader (callable): A function to load a sample given its path. extensions (tuple[string]): A list of allowed extensions. both extensions and is_valid_file should not be passed. transform (callable, optional): A function/transform that takes in a sample and returns a transformed version. E.g, ``transforms.RandomCrop`` for images. target_transform (callable, optional): A function/transform that takes in the target and transforms it. is_valid_file (callable, optional): A function that takes path of a file and check if the file is a valid file (used to check of corrupt files) both extensions and is_valid_file should not be passed. Attributes: classes (list): List of the class names sorted alphabetically. class_to_idx (dict): Dict with items (class_name, class_index). samples (list): List of (sample path, class_index) tuples targets (list): The class_index value for each image in the dataset """ def __init__( self, root: str, loader: Callable[[str], Any], extensions: Optional[Tuple[str, ...]] = None, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, is_valid_file: Optional[Callable[[str], bool]] = None, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) classes, class_to_idx = self.find_classes(self.root) samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file) self.loader = loader self.extensions = extensions self.classes = classes self.class_to_idx = class_to_idx self.samples = samples self.targets = [s[1] for s in samples] @staticmethod def make_dataset( directory: str, class_to_idx: Dict[str, int], extensions: Optional[Tuple[str, ...]] = None, is_valid_file: Optional[Callable[[str], bool]] = None, ) -> List[Tuple[str, int]]: """Generates a list of samples of a form (path_to_sample, class). This can be overridden to e.g. read files from a compressed zip file instead of from the disk. Args: directory (str): root dataset directory, corresponding to ``self.root``. class_to_idx (Dict[str, int]): Dictionary mapping class name to class index. extensions (optional): A list of allowed extensions. Either extensions or is_valid_file should be passed. Defaults to None. is_valid_file (optional): A function that takes path of a file and checks if the file is a valid file (used to check of corrupt files) both extensions and is_valid_file should not be passed. Defaults to None. Raises: ValueError: In case ``class_to_idx`` is empty. ValueError: In case ``extensions`` and ``is_valid_file`` are None or both are not None. FileNotFoundError: In case no valid file was found for any class. Returns: List[Tuple[str, int]]: samples of a form (path_to_sample, class) """ if class_to_idx is None: # prevent potential bug since make_dataset() would use the class_to_idx logic of the # find_classes() function, instead of using that of the find_classes() method, which # is potentially overridden and thus could have a different logic. raise ValueError("The class_to_idx parameter cannot be None.") return make_dataset(directory, class_to_idx, extensions=extensions, is_valid_file=is_valid_file) def find_classes(self, directory: str) -> Tuple[List[str], Dict[str, int]]: """Find the class folders in a dataset structured as follows:: directory/ ├── class_x │ ├── xxx.ext │ ├── xxy.ext │ └── ... │ └── xxz.ext └── class_y ├── 123.ext ├── nsdf3.ext └── ... └── asd932_.ext This method can be overridden to only consider a subset of classes, or to adapt to a different dataset directory structure. Args: directory(str): Root directory path, corresponding to ``self.root`` Raises: FileNotFoundError: If ``dir`` has no class folders. Returns: (Tuple[List[str], Dict[str, int]]): List of all classes and dictionary mapping each class to an index. """ return find_classes(directory) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (sample, target) where target is class_index of the target class. """ path, target = self.samples[index] sample = self.loader(path) if self.transform is not None: sample = self.transform(sample) if self.target_transform is not None: target = self.target_transform(target) return sample, target def __len__(self) -> int: return len(self.samples) IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".pgm", ".tif", ".tiff", ".webp") def pil_loader(path: str) -> Image.Image: # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) with open(path, "rb") as f: img = Image.open(f) return img.convert("RGB") # TODO: specify the return type def accimage_loader(path: str) -> Any: import accimage try: return accimage.Image(path) except OSError: # Potentially a decoding problem, fall back to PIL.Image return pil_loader(path) def default_loader(path: str) -> Any: from ms_adapter.torchvision import get_image_backend if get_image_backend() == "accimage": return accimage_loader(path) else: return pil_loader(path) class ImageFolder(DatasetFolder): """A generic data loader where the images are arranged in this way by default: :: root/dog/xxx.png root/dog/xxy.png root/dog/[...]/xxz.png root/cat/123.png root/cat/nsdf3.png root/cat/[...]/asd932_.png This class inherits from :class:`~torchvision.datasets.DatasetFolder` so the same methods can be overridden to customize the dataset. Args: root (string): Root directory path. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. loader (callable, optional): A function to load an image given its path. is_valid_file (callable, optional): A function that takes path of an Image file and check if the file is a valid file (used to check of corrupt files) Attributes: classes (list): List of the class names sorted alphabetically. class_to_idx (dict): Dict with items (class_name, class_index). imgs (list): List of (image path, class_index) tuples """ def __init__( self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, loader: Callable[[str], Any] = default_loader, is_valid_file: Optional[Callable[[str], bool]] = None, ): super().__init__( root, loader, IMG_EXTENSIONS if is_valid_file is None else None, transform=transform, target_transform=target_transform, is_valid_file=is_valid_file, ) self.imgs = self.samples msadapter/ms_adapter/torchvision/datasets/food101.py000066400000000000000000000072011436713554100230700ustar00rootroot00000000000000import json from pathlib import Path from typing import Any, Tuple, Callable, Optional import PIL.Image from .utils import verify_str_arg, download_and_extract_archive from .vision import VisionDataset class Food101(VisionDataset): """`The Food-101 Data Set `_. The Food-101 is a challenging data set of 101 food categories, with 101'000 images. For each class, 250 manually reviewed test images are provided as well as 750 training images. On purpose, the training images were not cleaned, and thus still contain some amount of noise. This comes mostly in the form of intense colors and sometimes wrong labels. All images were rescaled to have a maximum side length of 512 pixels. Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"train"`` (default) and ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. Default is False. """ _URL = "http://data.vision.ee.ethz.ch/cvl/food-101.tar.gz" _MD5 = "85eeb15f3717b99a5da872d97d918f87" def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self._split = verify_str_arg(split, "split", ("train", "test")) self._base_folder = Path(self.root) / "food-101" self._meta_folder = self._base_folder / "meta" self._images_folder = self._base_folder / "images" if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") self._labels = [] self._image_files = [] with open(self._meta_folder / f"{split}.json") as f: metadata = json.loads(f.read()) self.classes = sorted(metadata.keys()) self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) for class_label, im_rel_paths in metadata.items(): self._labels += [self.class_to_idx[class_label]] * len(im_rel_paths) self._image_files += [ self._images_folder.joinpath(*f"{im_rel_path}.jpg".split("/")) for im_rel_path in im_rel_paths ] def __len__(self) -> int: return len(self._image_files) def __getitem__(self, idx) -> Tuple[Any, Any]: image_file, label = self._image_files[idx], self._labels[idx] image = PIL.Image.open(image_file).convert("RGB") if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def extra_repr(self) -> str: return f"split={self._split}" def _check_exists(self) -> bool: return all(folder.exists() and folder.is_dir() for folder in (self._meta_folder, self._images_folder)) def _download(self) -> None: if self._check_exists(): return download_and_extract_archive(self._URL, download_root=self.root, md5=self._MD5) msadapter/ms_adapter/torchvision/datasets/gtsrb.py000066400000000000000000000072361436713554100230500ustar00rootroot00000000000000import csv import pathlib from typing import Any, Callable, Optional, Tuple import PIL from .folder import make_dataset from .utils import download_and_extract_archive, verify_str_arg from .vision import VisionDataset class GTSRB(VisionDataset): """`German Traffic Sign Recognition Benchmark (GTSRB) `_ Dataset. Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"train"`` (default), or ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self._split = verify_str_arg(split, "split", ("train", "test")) self._base_folder = pathlib.Path(root) / "gtsrb" self._target_folder = ( self._base_folder / "GTSRB" / ("Training" if self._split == "train" else "Final_Test/Images") ) if download: self.download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") if self._split == "train": samples = make_dataset(str(self._target_folder), extensions=(".ppm",)) else: with open(self._base_folder / "GT-final_test.csv") as csv_file: samples = [ (str(self._target_folder / row["Filename"]), int(row["ClassId"])) for row in csv.DictReader(csv_file, delimiter=";", skipinitialspace=True) ] self._samples = samples self.transform = transform self.target_transform = target_transform def __len__(self) -> int: return len(self._samples) def __getitem__(self, index: int) -> Tuple[Any, Any]: path, target = self._samples[index] sample = PIL.Image.open(path).convert("RGB") if self.transform is not None: sample = self.transform(sample) if self.target_transform is not None: target = self.target_transform(target) return sample, target def _check_exists(self) -> bool: return self._target_folder.is_dir() def download(self) -> None: if self._check_exists(): return base_url = "https://sid.erda.dk/public/archives/daaeac0d7ce1152aea9b61d9f1e19370/" if self._split == "train": download_and_extract_archive( f"{base_url}GTSRB-Training_fixed.zip", download_root=str(self._base_folder), md5="513f3c79a4c5141765e10e952eaa2478", ) else: download_and_extract_archive( f"{base_url}GTSRB_Final_Test_Images.zip", download_root=str(self._base_folder), md5="c7e4e6327067d32654124b0fe9e82185", ) download_and_extract_archive( f"{base_url}GTSRB_Final_Test_GT.zip", download_root=str(self._base_folder), md5="fe31e9c9270bbcd7b84b7f21a9d9d9e5", ) msadapter/ms_adapter/torchvision/datasets/hmdb51.py000066400000000000000000000133371436713554100230060ustar00rootroot00000000000000import glob import os from typing import Optional, Callable, Tuple, Dict, Any, List from .folder import find_classes, make_dataset from .video_utils import VideoClips from .vision import VisionDataset class HMDB51(VisionDataset): """ `HMDB51 `_ dataset. HMDB51 is an action recognition video dataset. This dataset consider every video as a collection of video clips of fixed size, specified by ``frames_per_clip``, where the step in frames between each clip is given by ``step_between_clips``. To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5`` and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two elements will come from video 1, and the next three elements from video 2. Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all frames in a video might be present. Internally, it uses a VideoClips object to handle clip creation. Args: root (string): Root directory of the HMDB51 Dataset. annotation_path (str): Path to the folder containing the split files. frames_per_clip (int): Number of frames in a clip. step_between_clips (int): Number of frames between each clip. fold (int, optional): Which fold to use. Should be between 1 and 3. train (bool, optional): If ``True``, creates a dataset from the train split, otherwise from the ``test`` split. transform (callable, optional): A function/transform that takes in a TxHxWxC video and returns a transformed version. output_format (str, optional): The format of the output video tensors (before transforms). Can be either "THWC" (default) or "TCHW". Returns: tuple: A 3-tuple with the following entries: - video (Tensor[T, H, W, C] or Tensor[T, C, H, W]): The `T` video frames - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels and `L` is the number of points - label (int): class of the video clip """ data_url = "https://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/hmdb51_org.rar" splits = { "url": "https://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/test_train_splits.rar", "md5": "15e67781e70dcfbdce2d7dbb9b3344b5", } TRAIN_TAG = 1 TEST_TAG = 2 def __init__( self, root: str, annotation_path: str, frames_per_clip: int, step_between_clips: int = 1, frame_rate: Optional[int] = None, fold: int = 1, train: bool = True, transform: Optional[Callable] = None, _precomputed_metadata: Optional[Dict[str, Any]] = None, num_workers: int = 1, _video_width: int = 0, _video_height: int = 0, _video_min_dimension: int = 0, _audio_samples: int = 0, output_format: str = "THWC", ) -> None: super().__init__(root) if fold not in (1, 2, 3): raise ValueError(f"fold should be between 1 and 3, got {fold}") extensions = ("avi",) self.classes, class_to_idx = find_classes(self.root) self.samples = make_dataset( self.root, class_to_idx, extensions, ) video_paths = [path for (path, _) in self.samples] video_clips = VideoClips( video_paths, frames_per_clip, step_between_clips, frame_rate, _precomputed_metadata, num_workers=num_workers, _video_width=_video_width, _video_height=_video_height, _video_min_dimension=_video_min_dimension, _audio_samples=_audio_samples, output_format=output_format, ) # we bookkeep the full version of video clips because we want to be able # to return the meta data of full version rather than the subset version of # video clips self.full_video_clips = video_clips self.fold = fold self.train = train self.indices = self._select_fold(video_paths, annotation_path, fold, train) self.video_clips = video_clips.subset(self.indices) self.transform = transform @property def metadata(self) -> Dict[str, Any]: return self.full_video_clips.metadata def _select_fold(self, video_list: List[str], annotations_dir: str, fold: int, train: bool) -> List[int]: target_tag = self.TRAIN_TAG if train else self.TEST_TAG split_pattern_name = f"*test_split{fold}.txt" split_pattern_path = os.path.join(annotations_dir, split_pattern_name) annotation_paths = glob.glob(split_pattern_path) selected_files = set() for filepath in annotation_paths: with open(filepath) as fid: lines = fid.readlines() for line in lines: video_filename, tag_string = line.split() tag = int(tag_string) if tag == target_tag: selected_files.add(video_filename) indices = [] for video_index, video_path in enumerate(video_list): if os.path.basename(video_path) in selected_files: indices.append(video_index) return indices def __len__(self) -> int: return self.video_clips.num_clips() def __getitem__(self, idx: int): video, audio, _, video_idx = self.video_clips.get_clip(idx) sample_index = self.indices[video_idx] _, class_index = self.samples[sample_index] if self.transform is not None: video = self.transform(video) return video, audio, class_index msadapter/ms_adapter/torchvision/datasets/imagenet.py000066400000000000000000000204011436713554100235050ustar00rootroot00000000000000import os import shutil import tempfile import pickle from contextlib import contextmanager from typing import Any, Dict, List, Iterator, Optional, Tuple from .folder import ImageFolder from .utils import check_integrity, extract_archive, verify_str_arg ARCHIVE_META = { "train": ("ILSVRC2012_img_train.tar", "1d675b47d978889d74fa0da5fadfb00e"), "val": ("ILSVRC2012_img_val.tar", "29b22e2961454d5413ddabcf34fc5622"), "devkit": ("ILSVRC2012_devkit_t12.tar.gz", "fa75699e90414af021442c21a62c3abf"), } META_FILE = "meta.bin" class ImageNet(ImageFolder): """`ImageNet `_ 2012 Classification Dataset. Args: root (string): Root directory of the ImageNet Dataset. split (string, optional): The dataset split, supports ``train``, or ``val``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. loader (callable, optional): A function to load an image given its path. Attributes: classes (list): List of the class name tuples. class_to_idx (dict): Dict with items (class_name, class_index). wnids (list): List of the WordNet IDs. wnid_to_idx (dict): Dict with items (wordnet_id, class_index). imgs (list): List of (image path, class_index) tuples targets (list): The class_index value for each image in the dataset """ def __init__(self, root: str, split: str = "train", **kwargs: Any) -> None: root = self.root = os.path.expanduser(root) self.split = verify_str_arg(split, "split", ("train", "val")) self.parse_archives() #TODO save dataset into MindRecord wnid_to_classes = load_meta_file(self.root)[0] # wnid_to_classes, self.val_wnid= parse_devkit_archive(self.root) super().__init__(self.split_folder, **kwargs) self.root = root self.wnids = self.classes self.wnid_to_idx = self.class_to_idx self.classes = [wnid_to_classes[wnid] for wnid in self.wnids] self.class_to_idx = {cls: idx for idx, clss in enumerate(self.classes) for cls in clss} def parse_archives(self) -> None: if not check_integrity(os.path.join(self.root, META_FILE)): parse_devkit_archive(self.root) if not os.path.isdir(self.split_folder): if self.split == "train": parse_train_archive(self.root) elif self.split == "val": parse_val_archive(self.root) @property def split_folder(self) -> str: return os.path.join(self.root, self.split) def extra_repr(self) -> str: return "Split: {split}".format(**self.__dict__) def load_meta_file(root: str, file: Optional[str] = None) -> Tuple[Dict[str, str], List[str]]: if file is None: file = META_FILE file = os.path.join(root, file) if check_integrity(file): with open(file, "rb") as f: data = pickle.load(f) return data else: msg = ( "The meta file {} is not present in the root directory or is corrupted. " "This file is automatically created by the ImageNet dataset." ) raise RuntimeError(msg.format(file, root)) def _verify_archive(root: str, file: str, md5: str) -> None: if not check_integrity(os.path.join(root, file), md5): msg = ( "The archive {} is not present in the root directory or is corrupted. " "You need to download it externally and place it in {}." ) raise RuntimeError(msg.format(file, root)) def parse_devkit_archive(root: str, file: Optional[str] = None): """Parse the devkit archive of the ImageNet2012 classification dataset and save the meta information in a binary file. Args: root (str): Root directory containing the devkit archive file (str, optional): Name of devkit archive. Defaults to 'ILSVRC2012_devkit_t12.tar.gz' """ import scipy.io as sio def parse_meta_mat(devkit_root: str) -> Tuple[Dict[int, str], Dict[str, Tuple[str, ...]]]: metafile = os.path.join(devkit_root, "data", "meta.mat") meta = sio.loadmat(metafile, squeeze_me=True)["synsets"] nums_children = list(zip(*meta))[4] meta = [meta[idx] for idx, num_children in enumerate(nums_children) if num_children == 0] idcs, wnids, classes = list(zip(*meta))[:3] classes = [tuple(clss.split(", ")) for clss in classes] idx_to_wnid = {idx: wnid for idx, wnid in zip(idcs, wnids)} wnid_to_classes = {wnid: clss for wnid, clss in zip(wnids, classes)} return idx_to_wnid, wnid_to_classes def parse_val_groundtruth_txt(devkit_root: str) -> List[int]: file = os.path.join(devkit_root, "data", "ILSVRC2012_validation_ground_truth.txt") with open(file) as txtfh: val_idcs = txtfh.readlines() return [int(val_idx) for val_idx in val_idcs] @contextmanager def get_tmp_dir() -> Iterator[str]: tmp_dir = tempfile.mkdtemp() try: yield tmp_dir finally: shutil.rmtree(tmp_dir) archive_meta = ARCHIVE_META["devkit"] if file is None: file = archive_meta[0] md5 = archive_meta[1] _verify_archive(root, file, md5) with get_tmp_dir() as tmp_dir: extract_archive(os.path.join(root, file), tmp_dir) devkit_root = os.path.join(tmp_dir, "ILSVRC2012_devkit_t12") idx_to_wnid, wnid_to_classes = parse_meta_mat(devkit_root) val_idcs = parse_val_groundtruth_txt(devkit_root) val_wnids = [idx_to_wnid[idx] for idx in val_idcs] with open(os.path.join(root, META_FILE), "wb") as f: pickle.dump((wnid_to_classes, val_wnids), f) # torch.save((wnid_to_classes, val_wnids), os.path.join(root, META_FILE)) # return wnid_to_classes, val_wnids def parse_train_archive(root: str, file: Optional[str] = None, folder: str = "train") -> None: """Parse the train images archive of the ImageNet2012 classification dataset and prepare it for usage with the ImageNet dataset. Args: root (str): Root directory containing the train images archive file (str, optional): Name of train images archive. Defaults to 'ILSVRC2012_img_train.tar' folder (str, optional): Optional name for train images folder. Defaults to 'train' """ archive_meta = ARCHIVE_META["train"] if file is None: file = archive_meta[0] md5 = archive_meta[1] _verify_archive(root, file, md5) train_root = os.path.join(root, folder) extract_archive(os.path.join(root, file), train_root) archives = [os.path.join(train_root, archive) for archive in os.listdir(train_root)] for archive in archives: extract_archive(archive, os.path.splitext(archive)[0], remove_finished=True) def parse_val_archive( root: str, file: Optional[str] = None, wnids: Optional[List[str]] = None, folder: str = "val" ) -> None: """Parse the validation images archive of the ImageNet2012 classification dataset and prepare it for usage with the ImageNet dataset. Args: root (str): Root directory containing the validation images archive file (str, optional): Name of validation images archive. Defaults to 'ILSVRC2012_img_val.tar' wnids (list, optional): List of WordNet IDs of the validation images. If None is given, the IDs are loaded from the meta file in the root directory folder (str, optional): Optional name for validation images folder. Defaults to 'val' """ archive_meta = ARCHIVE_META["val"] if file is None: file = archive_meta[0] md5 = archive_meta[1] if wnids is None: wnids = load_meta_file(root)[1] _verify_archive(root, file, md5) val_root = os.path.join(root, folder) extract_archive(os.path.join(root, file), val_root) images = sorted(os.path.join(val_root, image) for image in os.listdir(val_root)) for wnid in set(wnids): os.mkdir(os.path.join(val_root, wnid)) for wnid, img_file in zip(wnids, images): shutil.move(img_file, os.path.join(val_root, wnid, os.path.basename(img_file))) msadapter/ms_adapter/torchvision/datasets/inaturalist.py000066400000000000000000000235731436713554100242700ustar00rootroot00000000000000import os import os.path from typing import Any, Callable, Dict, List, Optional, Union, Tuple from PIL import Image from .utils import download_and_extract_archive, verify_str_arg from .vision import VisionDataset CATEGORIES_2021 = ["kingdom", "phylum", "class", "order", "family", "genus"] DATASET_URLS = { "2017": "https://ml-inat-competition-datasets.s3.amazonaws.com/2017/train_val_images.tar.gz", "2018": "https://ml-inat-competition-datasets.s3.amazonaws.com/2018/train_val2018.tar.gz", "2019": "https://ml-inat-competition-datasets.s3.amazonaws.com/2019/train_val2019.tar.gz", "2021_train": "https://ml-inat-competition-datasets.s3.amazonaws.com/2021/train.tar.gz", "2021_train_mini": "https://ml-inat-competition-datasets.s3.amazonaws.com/2021/train_mini.tar.gz", "2021_valid": "https://ml-inat-competition-datasets.s3.amazonaws.com/2021/val.tar.gz", } DATASET_MD5 = { "2017": "7c784ea5e424efaec655bd392f87301f", "2018": "b1c6952ce38f31868cc50ea72d066cc3", "2019": "c60a6e2962c9b8ccbd458d12c8582644", "2021_train": "e0526d53c7f7b2e3167b2b43bb2690ed", "2021_train_mini": "db6ed8330e634445efc8fec83ae81442", "2021_valid": "f6f6e0e242e3d4c9569ba56400938afc", } class INaturalist(VisionDataset): """`iNaturalist `_ Dataset. Args: root (string): Root directory of dataset where the image files are stored. This class does not require/use annotation files. version (string, optional): Which version of the dataset to download/use. One of '2017', '2018', '2019', '2021_train', '2021_train_mini', '2021_valid'. Default: `2021_train`. target_type (string or list, optional): Type of target to use, for 2021 versions, one of: - ``full``: the full category (species) - ``kingdom``: e.g. "Animalia" - ``phylum``: e.g. "Arthropoda" - ``class``: e.g. "Insecta" - ``order``: e.g. "Coleoptera" - ``family``: e.g. "Cleridae" - ``genus``: e.g. "Trichodes" for 2017-2019 versions, one of: - ``full``: the full (numeric) category - ``super``: the super category, e.g. "Amphibians" Can also be a list to output a tuple with all specified target types. Defaults to ``full``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ def __init__( self, root: str, version: str = "2021_train", target_type: Union[List[str], str] = "full", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: self.version = verify_str_arg(version, "version", DATASET_URLS.keys()) super().__init__(os.path.join(root, version), transform=transform, target_transform=target_transform) os.makedirs(root, exist_ok=True) if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") self.all_categories: List[str] = [] # map: category type -> name of category -> index self.categories_index: Dict[str, Dict[str, int]] = {} # list indexed by category id, containing mapping from category type -> index self.categories_map: List[Dict[str, int]] = [] if not isinstance(target_type, list): target_type = [target_type] if self.version[:4] == "2021": self.target_type = [verify_str_arg(t, "target_type", ("full", *CATEGORIES_2021)) for t in target_type] self._init_2021() else: self.target_type = [verify_str_arg(t, "target_type", ("full", "super")) for t in target_type] self._init_pre2021() # index of all files: (full category id, filename) self.index: List[Tuple[int, str]] = [] for dir_index, dir_name in enumerate(self.all_categories): files = os.listdir(os.path.join(self.root, dir_name)) for fname in files: self.index.append((dir_index, fname)) def _init_2021(self) -> None: """Initialize based on 2021 layout""" self.all_categories = sorted(os.listdir(self.root)) # map: category type -> name of category -> index self.categories_index = {k: {} for k in CATEGORIES_2021} for dir_index, dir_name in enumerate(self.all_categories): pieces = dir_name.split("_") if len(pieces) != 8: raise RuntimeError(f"Unexpected category name {dir_name}, wrong number of pieces") if pieces[0] != f"{dir_index:05d}": raise RuntimeError(f"Unexpected category id {pieces[0]}, expecting {dir_index:05d}") cat_map = {} for cat, name in zip(CATEGORIES_2021, pieces[1:7]): if name in self.categories_index[cat]: cat_id = self.categories_index[cat][name] else: cat_id = len(self.categories_index[cat]) self.categories_index[cat][name] = cat_id cat_map[cat] = cat_id self.categories_map.append(cat_map) def _init_pre2021(self) -> None: """Initialize based on 2017-2019 layout""" # map: category type -> name of category -> index self.categories_index = {"super": {}} cat_index = 0 super_categories = sorted(os.listdir(self.root)) for sindex, scat in enumerate(super_categories): self.categories_index["super"][scat] = sindex subcategories = sorted(os.listdir(os.path.join(self.root, scat))) for subcat in subcategories: if self.version == "2017": # this version does not use ids as directory names subcat_i = cat_index cat_index += 1 else: try: subcat_i = int(subcat) except ValueError: raise RuntimeError(f"Unexpected non-numeric dir name: {subcat}") if subcat_i >= len(self.categories_map): old_len = len(self.categories_map) self.categories_map.extend([{}] * (subcat_i - old_len + 1)) self.all_categories.extend([""] * (subcat_i - old_len + 1)) if self.categories_map[subcat_i]: raise RuntimeError(f"Duplicate category {subcat}") self.categories_map[subcat_i] = {"super": sindex} self.all_categories[subcat_i] = os.path.join(scat, subcat) # validate the dictionary for cindex, c in enumerate(self.categories_map): if not c: raise RuntimeError(f"Missing category {cindex}") def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where the type of target specified by target_type. """ cat_id, fname = self.index[index] img = Image.open(os.path.join(self.root, self.all_categories[cat_id], fname)) target: Any = [] for t in self.target_type: if t == "full": target.append(cat_id) else: target.append(self.categories_map[cat_id][t]) target = tuple(target) if len(target) > 1 else target[0] if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.index) def category_name(self, category_type: str, category_id: int) -> str: """ Args: category_type(str): one of "full", "kingdom", "phylum", "class", "order", "family", "genus" or "super" category_id(int): an index (class id) from this category Returns: the name of the category """ if category_type == "full": return self.all_categories[category_id] else: if category_type not in self.categories_index: raise ValueError(f"Invalid category type '{category_type}'") else: for name, id in self.categories_index[category_type].items(): if id == category_id: return name raise ValueError(f"Invalid category id {category_id} for {category_type}") def _check_integrity(self) -> bool: return os.path.exists(self.root) and len(os.listdir(self.root)) > 0 def download(self) -> None: if self._check_integrity(): raise RuntimeError( f"The directory {self.root} already exists. " f"If you want to re-download or re-extract the images, delete the directory." ) base_root = os.path.dirname(self.root) download_and_extract_archive( DATASET_URLS[self.version], base_root, filename=f"{self.version}.tgz", md5=DATASET_MD5[self.version] ) orig_dir_name = os.path.join(base_root, os.path.basename(DATASET_URLS[self.version]).rstrip(".tar.gz")) if not os.path.exists(orig_dir_name): raise RuntimeError(f"Unable to find downloaded files at {orig_dir_name}") os.rename(orig_dir_name, self.root) print(f"Dataset version '{self.version}' has been downloaded and prepared for use") msadapter/ms_adapter/torchvision/datasets/kinetics.py000066400000000000000000000321361436713554100235350ustar00rootroot00000000000000import csv import os import time import urllib import warnings from functools import partial from multiprocessing import Pool from os import path from typing import Any, Callable, Dict, Optional, Tuple from .folder import find_classes, make_dataset from .utils import download_and_extract_archive, download_url, verify_str_arg, check_integrity from .video_utils import VideoClips from .vision import VisionDataset def _dl_wrap(tarpath: str, videopath: str, line: str) -> None: download_and_extract_archive(line, tarpath, videopath) class Kinetics(VisionDataset): """`Generic Kinetics `_ dataset. Kinetics-400/600/700 are action recognition video datasets. This dataset consider every video as a collection of video clips of fixed size, specified by ``frames_per_clip``, where the step in frames between each clip is given by ``step_between_clips``. To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5`` and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two elements will come from video 1, and the next three elements from video 2. Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all frames in a video might be present. Args: root (string): Root directory of the Kinetics Dataset. Directory should be structured as follows: .. code:: root/ ├── split │ ├── class1 │ │ ├── clip1.mp4 │ │ ├── clip2.mp4 │ │ ├── clip3.mp4 │ │ ├── ... │ ├── class2 │ │ ├── clipx.mp4 │ │ └── ... Note: split is appended automatically using the split argument. frames_per_clip (int): number of frames in a clip num_classes (int): select between Kinetics-400 (default), Kinetics-600, and Kinetics-700 split (str): split of the dataset to consider; supports ``"train"`` (default) ``"val"`` ``"test"`` frame_rate (float): If omitted, interpolate different frame rate for each clip. step_between_clips (int): number of frames between each clip transform (callable, optional): A function/transform that takes in a TxHxWxC video and returns a transformed version. download (bool): Download the official version of the dataset to root folder. num_workers (int): Use multiple workers for VideoClips creation num_download_workers (int): Use multiprocessing in order to speed up download. output_format (str, optional): The format of the output video tensors (before transforms). Can be either "THWC" or "TCHW" (default). Note that in most other utils and datasets, the default is actually "THWC". Returns: tuple: A 3-tuple with the following entries: - video (Tensor[T, C, H, W] or Tensor[T, H, W, C]): the `T` video frames in torch.uint8 tensor - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels and `L` is the number of points in torch.float tensor - label (int): class of the video clip Raises: RuntimeError: If ``download is True`` and the video archives are already extracted. """ _TAR_URLS = { "400": "https://s3.amazonaws.com/kinetics/400/{split}/k400_{split}_path.txt", "600": "https://s3.amazonaws.com/kinetics/600/{split}/k600_{split}_path.txt", "700": "https://s3.amazonaws.com/kinetics/700_2020/{split}/k700_2020_{split}_path.txt", } _ANNOTATION_URLS = { "400": "https://s3.amazonaws.com/kinetics/400/annotations/{split}.csv", "600": "https://s3.amazonaws.com/kinetics/600/annotations/{split}.csv", "700": "https://s3.amazonaws.com/kinetics/700_2020/annotations/{split}.csv", } def __init__( self, root: str, frames_per_clip: int, num_classes: str = "400", split: str = "train", frame_rate: Optional[int] = None, step_between_clips: int = 1, transform: Optional[Callable] = None, extensions: Tuple[str, ...] = ("avi", "mp4"), download: bool = False, num_download_workers: int = 1, num_workers: int = 1, _precomputed_metadata: Optional[Dict[str, Any]] = None, _video_width: int = 0, _video_height: int = 0, _video_min_dimension: int = 0, _audio_samples: int = 0, _audio_channels: int = 0, _legacy: bool = False, output_format: str = "TCHW", ) -> None: # TODO: support test self.num_classes = verify_str_arg(num_classes, arg="num_classes", valid_values=["400", "600", "700"]) self.extensions = extensions self.num_download_workers = num_download_workers self.root = root self._legacy = _legacy if _legacy: print("Using legacy structure") self.split_folder = root self.split = "unknown" output_format = "THWC" if download: raise ValueError("Cannot download the videos using legacy_structure.") else: self.split_folder = path.join(root, split) self.split = verify_str_arg(split, arg="split", valid_values=["train", "val", "test"]) if download: self.download_and_process_videos() super().__init__(self.root) self.classes, class_to_idx = find_classes(self.split_folder) self.samples = make_dataset(self.split_folder, class_to_idx, extensions, is_valid_file=None) video_list = [x[0] for x in self.samples] self.video_clips = VideoClips( video_list, frames_per_clip, step_between_clips, frame_rate, _precomputed_metadata, num_workers=num_workers, _video_width=_video_width, _video_height=_video_height, _video_min_dimension=_video_min_dimension, _audio_samples=_audio_samples, _audio_channels=_audio_channels, output_format=output_format, ) self.transform = transform def download_and_process_videos(self) -> None: """Downloads all the videos to the _root_ folder in the expected format.""" tic = time.time() self._download_videos() toc = time.time() print("Elapsed time for downloading in mins ", (toc - tic) / 60) self._make_ds_structure() toc2 = time.time() print("Elapsed time for processing in mins ", (toc2 - toc) / 60) print("Elapsed time overall in mins ", (toc2 - tic) / 60) def _download_videos(self) -> None: """download tarballs containing the video to "tars" folder and extract them into the _split_ folder where split is one of the official dataset splits. Raises: RuntimeError: if download folder exists, break to prevent downloading entire dataset again. """ if path.exists(self.split_folder): raise RuntimeError( f"The directory {self.split_folder} already exists. " f"If you want to re-download or re-extract the images, delete the directory." ) tar_path = path.join(self.root, "tars") file_list_path = path.join(self.root, "files") split_url = self._TAR_URLS[self.num_classes].format(split=self.split) split_url_filepath = path.join(file_list_path, path.basename(split_url)) if not check_integrity(split_url_filepath): download_url(split_url, file_list_path) with open(split_url_filepath) as file: list_video_urls = [urllib.parse.quote(line, safe="/,:") for line in file.read().splitlines()] if self.num_download_workers == 1: for line in list_video_urls: download_and_extract_archive(line, tar_path, self.split_folder) else: part = partial(_dl_wrap, tar_path, self.split_folder) poolproc = Pool(self.num_download_workers) poolproc.map(part, list_video_urls) def _make_ds_structure(self) -> None: """move videos from split_folder/ ├── clip1.avi ├── clip2.avi to the correct format as described below: split_folder/ ├── class1 │ ├── clip1.avi """ annotation_path = path.join(self.root, "annotations") if not check_integrity(path.join(annotation_path, f"{self.split}.csv")): download_url(self._ANNOTATION_URLS[self.num_classes].format(split=self.split), annotation_path) annotations = path.join(annotation_path, f"{self.split}.csv") file_fmtstr = "{ytid}_{start:06}_{end:06}.mp4" with open(annotations) as csvfile: reader = csv.DictReader(csvfile) for row in reader: f = file_fmtstr.format( ytid=row["youtube_id"], start=int(row["time_start"]), end=int(row["time_end"]), ) label = row["label"].replace(" ", "_").replace("'", "").replace("(", "").replace(")", "") os.makedirs(path.join(self.split_folder, label), exist_ok=True) downloaded_file = path.join(self.split_folder, f) if path.isfile(downloaded_file): os.replace( downloaded_file, path.join(self.split_folder, label, f), ) @property def metadata(self) -> Dict[str, Any]: return self.video_clips.metadata def __len__(self) -> int: return self.video_clips.num_clips() def __getitem__(self, idx: int): video, audio, info, video_idx = self.video_clips.get_clip(idx) label = self.samples[video_idx][1] if self.transform is not None: video = self.transform(video) return video, audio, label class Kinetics400(Kinetics): """ `Kinetics-400 `_ dataset. .. warning:: This class was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``Kinetics(..., num_classes='400')`` instead. Kinetics-400 is an action recognition video dataset. This dataset consider every video as a collection of video clips of fixed size, specified by ``frames_per_clip``, where the step in frames between each clip is given by ``step_between_clips``. To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5`` and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two elements will come from video 1, and the next three elements from video 2. Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all frames in a video might be present. Internally, it uses a VideoClips object to handle clip creation. Args: root (string): Root directory of the Kinetics-400 Dataset. Should be structured as follows: .. code:: root/ ├── class1 │ ├── clip1.avi │ ├── clip2.avi │ ├── clip3.mp4 │ └── ... └── class2 ├── clipx.avi └── ... frames_per_clip (int): number of frames in a clip step_between_clips (int): number of frames between each clip transform (callable, optional): A function/transform that takes in a TxHxWxC video and returns a transformed version. Returns: tuple: A 3-tuple with the following entries: - video (Tensor[T, H, W, C]): the `T` video frames - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels and `L` is the number of points - label (int): class of the video clip """ def __init__( self, root: str, frames_per_clip: int, num_classes: Any = None, split: Any = None, download: Any = None, num_download_workers: Any = None, **kwargs: Any, ) -> None: warnings.warn( "The Kinetics400 class is deprecated since 0.12 and will be removed in 0.14." "Please use Kinetics(..., num_classes='400') instead." "Note that Kinetics(..., num_classes='400') returns video in a Tensor[T, C, H, W] format." ) if any(value is not None for value in (num_classes, split, download, num_download_workers)): raise RuntimeError( "Usage of 'num_classes', 'split', 'download', or 'num_download_workers' is not supported in " "Kinetics400. Please use Kinetics instead." ) super().__init__( root=root, frames_per_clip=frames_per_clip, _legacy=True, **kwargs, ) msadapter/ms_adapter/torchvision/datasets/kitti.py000066400000000000000000000127401436713554100230470ustar00rootroot00000000000000import csv import os from typing import Any, Callable, List, Optional, Tuple from PIL import Image from .utils import download_and_extract_archive from .vision import VisionDataset class Kitti(VisionDataset): """`KITTI `_ Dataset. It corresponds to the "left color images of object" dataset, for object detection. Args: root (string): Root directory where images are downloaded to. Expects the following folder structure if download=False: .. code:: └── Kitti └─ raw ├── training | ├── image_2 | └── label_2 └── testing └── image_2 train (bool, optional): Use ``train`` split if true, else ``test`` split. Defaults to ``train``. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ data_url = "https://s3.eu-central-1.amazonaws.com/avg-kitti/" resources = [ "data_object_image_2.zip", "data_object_label_2.zip", ] image_dir_name = "image_2" labels_dir_name = "label_2" def __init__( self, root: str, train: bool = True, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, transforms: Optional[Callable] = None, download: bool = False, ): super().__init__( root, transform=transform, target_transform=target_transform, transforms=transforms, ) self.images = [] self.targets = [] self.root = root self.train = train self._location = "training" if self.train else "testing" if download: self.download() if not self._check_exists(): raise RuntimeError("Dataset not found. You may use download=True to download it.") image_dir = os.path.join(self._raw_folder, self._location, self.image_dir_name) if self.train: labels_dir = os.path.join(self._raw_folder, self._location, self.labels_dir_name) for img_file in os.listdir(image_dir): self.images.append(os.path.join(image_dir, img_file)) if self.train: self.targets.append(os.path.join(labels_dir, f"{img_file.split('.')[0]}.txt")) def __getitem__(self, index: int) -> Tuple[Any, Any]: """Get item at a given index. Args: index (int): Index Returns: tuple: (image, target), where target is a list of dictionaries with the following keys: - type: str - truncated: float - occluded: int - alpha: float - bbox: float[4] - dimensions: float[3] - locations: float[3] - rotation_y: float """ image = Image.open(self.images[index]) target = self._parse_target(index) if self.train else None if self.transforms: image, target = self.transforms(image, target) return image, target def _parse_target(self, index: int) -> List: target = [] with open(self.targets[index]) as inp: content = csv.reader(inp, delimiter=" ") for line in content: target.append( { "type": line[0], "truncated": float(line[1]), "occluded": int(line[2]), "alpha": float(line[3]), "bbox": [float(x) for x in line[4:8]], "dimensions": [float(x) for x in line[8:11]], "location": [float(x) for x in line[11:14]], "rotation_y": float(line[14]), } ) return target def __len__(self) -> int: return len(self.images) @property def _raw_folder(self) -> str: return os.path.join(self.root, self.__class__.__name__, "raw") def _check_exists(self) -> bool: """Check if the data directory exists.""" folders = [self.image_dir_name] if self.train: folders.append(self.labels_dir_name) return all(os.path.isdir(os.path.join(self._raw_folder, self._location, fname)) for fname in folders) def download(self) -> None: """Download the KITTI data if it doesn't exist already.""" if self._check_exists(): return os.makedirs(self._raw_folder, exist_ok=True) # download files for fname in self.resources: download_and_extract_archive( url=f"{self.data_url}{fname}", download_root=self._raw_folder, filename=fname, ) msadapter/ms_adapter/torchvision/datasets/lfw.py000066400000000000000000000240521436713554100225120ustar00rootroot00000000000000import os from typing import Any, Callable, List, Optional, Tuple from PIL import Image from .utils import check_integrity, download_and_extract_archive, download_url, verify_str_arg from .vision import VisionDataset class _LFW(VisionDataset): base_folder = "lfw-py" download_url_prefix = "http://vis-www.cs.umass.edu/lfw/" file_dict = { "original": ("lfw", "lfw.tgz", "a17d05bd522c52d84eca14327a23d494"), "funneled": ("lfw_funneled", "lfw-funneled.tgz", "1b42dfed7d15c9b2dd63d5e5840c86ad"), "deepfunneled": ("lfw-deepfunneled", "lfw-deepfunneled.tgz", "68331da3eb755a505a502b5aacb3c201"), } checksums = { "pairs.txt": "9f1ba174e4e1c508ff7cdf10ac338a7d", "pairsDevTest.txt": "5132f7440eb68cf58910c8a45a2ac10b", "pairsDevTrain.txt": "4f27cbf15b2da4a85c1907eb4181ad21", "people.txt": "450f0863dd89e85e73936a6d71a3474b", "peopleDevTest.txt": "e4bf5be0a43b5dcd9dc5ccfcb8fb19c5", "peopleDevTrain.txt": "54eaac34beb6d042ed3a7d883e247a21", "lfw-names.txt": "a6d0a479bd074669f656265a6e693f6d", } annot_file = {"10fold": "", "train": "DevTrain", "test": "DevTest"} names = "lfw-names.txt" def __init__( self, root: str, split: str, image_set: str, view: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ): super().__init__(os.path.join(root, self.base_folder), transform=transform, target_transform=target_transform) self.image_set = verify_str_arg(image_set.lower(), "image_set", self.file_dict.keys()) images_dir, self.filename, self.md5 = self.file_dict[self.image_set] self.view = verify_str_arg(view.lower(), "view", ["people", "pairs"]) self.split = verify_str_arg(split.lower(), "split", ["10fold", "train", "test"]) self.labels_file = f"{self.view}{self.annot_file[self.split]}.txt" self.data: List[Any] = [] if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") self.images_dir = os.path.join(self.root, images_dir) def _loader(self, path: str) -> Image.Image: with open(path, "rb") as f: img = Image.open(f) return img.convert("RGB") def _check_integrity(self): st1 = check_integrity(os.path.join(self.root, self.filename), self.md5) st2 = check_integrity(os.path.join(self.root, self.labels_file), self.checksums[self.labels_file]) if not st1 or not st2: return False if self.view == "people": return check_integrity(os.path.join(self.root, self.names), self.checksums[self.names]) return True def download(self): if self._check_integrity(): print("Files already downloaded and verified") return url = f"{self.download_url_prefix}{self.filename}" download_and_extract_archive(url, self.root, filename=self.filename, md5=self.md5) download_url(f"{self.download_url_prefix}{self.labels_file}", self.root) if self.view == "people": download_url(f"{self.download_url_prefix}{self.names}", self.root) def _get_path(self, identity, no): return os.path.join(self.images_dir, identity, f"{identity}_{int(no):04d}.jpg") def extra_repr(self) -> str: return f"Alignment: {self.image_set}\nSplit: {self.split}" def __len__(self): return len(self.data) class LFWPeople(_LFW): """`LFW `_ Dataset. Args: root (string): Root directory of dataset where directory ``lfw-py`` exists or will be saved to if download is set to True. split (string, optional): The image split to use. Can be one of ``train``, ``test``, ``10fold`` (default). image_set (str, optional): Type of image funneling to use, ``original``, ``funneled`` or ``deepfunneled``. Defaults to ``funneled``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomRotation`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ def __init__( self, root: str, split: str = "10fold", image_set: str = "funneled", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ): super().__init__(root, split, image_set, "people", transform, target_transform, download) self.class_to_idx = self._get_classes() self.data, self.targets = self._get_people() def _get_people(self): data, targets = [], [] with open(os.path.join(self.root, self.labels_file)) as f: lines = f.readlines() n_folds, s = (int(lines[0]), 1) if self.split == "10fold" else (1, 0) for fold in range(n_folds): n_lines = int(lines[s]) people = [line.strip().split("\t") for line in lines[s + 1 : s + n_lines + 1]] s += n_lines + 1 for i, (identity, num_imgs) in enumerate(people): for num in range(1, int(num_imgs) + 1): img = self._get_path(identity, num) data.append(img) targets.append(self.class_to_idx[identity]) return data, targets def _get_classes(self): with open(os.path.join(self.root, self.names)) as f: lines = f.readlines() names = [line.strip().split()[0] for line in lines] class_to_idx = {name: i for i, name in enumerate(names)} return class_to_idx def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: Tuple (image, target) where target is the identity of the person. """ img = self._loader(self.data[index]) target = self.targets[index] if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def extra_repr(self) -> str: return super().extra_repr() + f"\nClasses (identities): {len(self.class_to_idx)}" class LFWPairs(_LFW): """`LFW `_ Dataset. Args: root (string): Root directory of dataset where directory ``lfw-py`` exists or will be saved to if download is set to True. split (string, optional): The image split to use. Can be one of ``train``, ``test``, ``10fold``. Defaults to ``10fold``. image_set (str, optional): Type of image funneling to use, ``original``, ``funneled`` or ``deepfunneled``. Defaults to ``funneled``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomRotation`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ def __init__( self, root: str, split: str = "10fold", image_set: str = "funneled", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ): super().__init__(root, split, image_set, "pairs", transform, target_transform, download) self.pair_names, self.data, self.targets = self._get_pairs(self.images_dir) def _get_pairs(self, images_dir): pair_names, data, targets = [], [], [] with open(os.path.join(self.root, self.labels_file)) as f: lines = f.readlines() if self.split == "10fold": n_folds, n_pairs = lines[0].split("\t") n_folds, n_pairs = int(n_folds), int(n_pairs) else: n_folds, n_pairs = 1, int(lines[0]) s = 1 for fold in range(n_folds): matched_pairs = [line.strip().split("\t") for line in lines[s : s + n_pairs]] unmatched_pairs = [line.strip().split("\t") for line in lines[s + n_pairs : s + (2 * n_pairs)]] s += 2 * n_pairs for pair in matched_pairs: img1, img2, same = self._get_path(pair[0], pair[1]), self._get_path(pair[0], pair[2]), 1 pair_names.append((pair[0], pair[0])) data.append((img1, img2)) targets.append(same) for pair in unmatched_pairs: img1, img2, same = self._get_path(pair[0], pair[1]), self._get_path(pair[2], pair[3]), 0 pair_names.append((pair[0], pair[2])) data.append((img1, img2)) targets.append(same) return pair_names, data, targets def __getitem__(self, index: int) -> Tuple[Any, Any, int]: """ Args: index (int): Index Returns: tuple: (image1, image2, target) where target is `0` for different indentities and `1` for same identities. """ img1, img2 = self.data[index] img1, img2 = self._loader(img1), self._loader(img2) target = self.targets[index] if self.transform is not None: img1, img2 = self.transform(img1), self.transform(img2) if self.target_transform is not None: target = self.target_transform(target) return img1, img2, target msadapter/ms_adapter/torchvision/datasets/lsun.py000066400000000000000000000130531436713554100227020ustar00rootroot00000000000000import io import os.path import pickle import string from collections.abc import Iterable from typing import Any, Callable, cast, List, Optional, Tuple, Union from PIL import Image from .utils import verify_str_arg, iterable_to_str from .vision import VisionDataset class LSUNClass(VisionDataset): def __init__( self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None ) -> None: import lmdb super().__init__(root, transform=transform, target_transform=target_transform) self.env = lmdb.open(root, max_readers=1, readonly=True, lock=False, readahead=False, meminit=False) with self.env.begin(write=False) as txn: self.length = txn.stat()["entries"] cache_file = "_cache_" + "".join(c for c in root if c in string.ascii_letters) if os.path.isfile(cache_file): self.keys = pickle.load(open(cache_file, "rb")) else: with self.env.begin(write=False) as txn: self.keys = [key for key in txn.cursor().iternext(keys=True, values=False)] pickle.dump(self.keys, open(cache_file, "wb")) def __getitem__(self, index: int) -> Tuple[Any, Any]: img, target = None, None env = self.env with env.begin(write=False) as txn: imgbuf = txn.get(self.keys[index]) buf = io.BytesIO() buf.write(imgbuf) buf.seek(0) img = Image.open(buf).convert("RGB") if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return self.length class LSUN(VisionDataset): """`LSUN `_ dataset. You will need to install the ``lmdb`` package to use this dataset: run ``pip install lmdb`` Args: root (string): Root directory for the database files. classes (string or list): One of {'train', 'val', 'test'} or a list of categories to load. e,g. ['bedroom_train', 'church_outdoor_train']. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ def __init__( self, root: str, classes: Union[str, List[str]] = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.classes = self._verify_classes(classes) # for each class, create an LSUNClassDataset self.dbs = [] for c in self.classes: self.dbs.append(LSUNClass(root=os.path.join(root, f"{c}_lmdb"), transform=transform)) self.indices = [] count = 0 for db in self.dbs: count += len(db) self.indices.append(count) self.length = count def _verify_classes(self, classes: Union[str, List[str]]) -> List[str]: categories = [ "bedroom", "bridge", "church_outdoor", "classroom", "conference_room", "dining_room", "kitchen", "living_room", "restaurant", "tower", ] dset_opts = ["train", "val", "test"] try: classes = cast(str, classes) verify_str_arg(classes, "classes", dset_opts) if classes == "test": classes = [classes] else: classes = [c + "_" + classes for c in categories] except ValueError: if not isinstance(classes, Iterable): msg = "Expected type str or Iterable for argument classes, but got type {}." raise ValueError(msg.format(type(classes))) classes = list(classes) msg_fmtstr_type = "Expected type str for elements in argument classes, but got type {}." for c in classes: verify_str_arg(c, custom_msg=msg_fmtstr_type.format(type(c))) c_short = c.split("_") category, dset_opt = "_".join(c_short[:-1]), c_short[-1] msg_fmtstr = "Unknown value '{}' for {}. Valid values are {{{}}}." msg = msg_fmtstr.format(category, "LSUN class", iterable_to_str(categories)) verify_str_arg(category, valid_values=categories, custom_msg=msg) msg = msg_fmtstr.format(dset_opt, "postfix", iterable_to_str(dset_opts)) verify_str_arg(dset_opt, valid_values=dset_opts, custom_msg=msg) return classes def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: Tuple (image, target) where target is the index of the target category. """ target = 0 sub = 0 for ind in self.indices: if index < ind: break target += 1 sub = ind db = self.dbs[target] index = index - sub if self.target_transform is not None: target = self.target_transform(target) img, _ = db[index] return img, target def __len__(self) -> int: return self.length def extra_repr(self) -> str: return "Classes: {classes}".format(**self.__dict__) msadapter/ms_adapter/torchvision/datasets/mnist.py000066400000000000000000000515151436713554100230600ustar00rootroot00000000000000import codecs import os import os.path import shutil import string import sys import warnings from typing import Any, Callable, Dict, List, Optional, Tuple from urllib.error import URLError import numpy as np import ms_adapter.pytorch as torch from PIL import Image from .utils import download_and_extract_archive, extract_archive, verify_str_arg, check_integrity from .vision import VisionDataset class MNIST(VisionDataset): """`MNIST `_ Dataset. Args: root (string): Root directory of dataset where ``MNIST/raw/train-images-idx3-ubyte`` and ``MNIST/raw/t10k-images-idx3-ubyte`` exist. train (bool, optional): If True, creates dataset from ``train-images-idx3-ubyte``, otherwise from ``t10k-images-idx3-ubyte``. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ mirrors = [ "http://yann.lecun.com/exdb/mnist/", "https://ossci-datasets.s3.amazonaws.com/mnist/", ] resources = [ ("train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"), ("train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"), ("t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"), ("t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c"), ] training_file = "training.pt" test_file = "test.pt" classes = [ "0 - zero", "1 - one", "2 - two", "3 - three", "4 - four", "5 - five", "6 - six", "7 - seven", "8 - eight", "9 - nine", ] @property def train_labels(self): warnings.warn("train_labels has been renamed targets") return self.targets @property def test_labels(self): warnings.warn("test_labels has been renamed targets") return self.targets @property def train_data(self): warnings.warn("train_data has been renamed data") return self.data @property def test_data(self): warnings.warn("test_data has been renamed data") return self.data def __init__( self, root: str, train: bool = True, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.train = train # training set or test set # if self._check_legacy_exist(): # self.data, self.targets = self._load_legacy_data() # return if download: self.download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") self.data, self.targets = self._load_data() # def _check_legacy_exist(self): # processed_folder_exists = os.path.exists(self.processed_folder) # if not processed_folder_exists: # return False # # return all( # check_integrity(os.path.join(self.processed_folder, file)) for file in (self.training_file, self.test_file) # ) # def _load_legacy_data(self): # # This is for BC only. We no longer cache the data in a custom binary, but simply read from the raw data # # directly. # data_file = self.training_file if self.train else self.test_file # return torch.load(os.path.join(self.processed_folder, data_file)) def _load_data(self): image_file = f"{'train' if self.train else 't10k'}-images-idx3-ubyte" data = read_image_file(os.path.join(self.raw_folder, image_file)) label_file = f"{'train' if self.train else 't10k'}-labels-idx1-ubyte" targets = read_label_file(os.path.join(self.raw_folder, label_file)) return data, targets def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target class. """ img, target = self.data[index], int(self.targets[index]) # doing this so that it is consistent with all other datasets # to return a PIL Image img = Image.fromarray(img.numpy(), mode="L") if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.data) @property def raw_folder(self) -> str: return os.path.join(self.root, self.__class__.__name__, "raw") @property def processed_folder(self) -> str: return os.path.join(self.root, self.__class__.__name__, "processed") @property def class_to_idx(self) -> Dict[str, int]: return {_class: i for i, _class in enumerate(self.classes)} def _check_exists(self) -> bool: return all( check_integrity(os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0])) for url, _ in self.resources ) def download(self) -> None: """Download the MNIST data if it doesn't exist already.""" if self._check_exists(): return os.makedirs(self.raw_folder, exist_ok=True) # download files for filename, md5 in self.resources: for mirror in self.mirrors: url = f"{mirror}{filename}" try: print(f"Downloading {url}") download_and_extract_archive(url, download_root=self.raw_folder, filename=filename, md5=md5) except URLError as error: print(f"Failed to download (trying next):\n{error}") continue finally: print() break else: raise RuntimeError(f"Error downloading {filename}") def extra_repr(self) -> str: split = "Train" if self.train is True else "Test" return f"Split: {split}" class FashionMNIST(MNIST): """`Fashion-MNIST `_ Dataset. Args: root (string): Root directory of dataset where ``FashionMNIST/raw/train-images-idx3-ubyte`` and ``FashionMNIST/raw/t10k-images-idx3-ubyte`` exist. train (bool, optional): If True, creates dataset from ``train-images-idx3-ubyte``, otherwise from ``t10k-images-idx3-ubyte``. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ mirrors = ["http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/"] resources = [ ("train-images-idx3-ubyte.gz", "8d4fb7e6c68d591d4c3dfef9ec88bf0d"), ("train-labels-idx1-ubyte.gz", "25c81989df183df01b3e8a0aad5dffbe"), ("t10k-images-idx3-ubyte.gz", "bef4ecab320f06d8554ea6380940ec79"), ("t10k-labels-idx1-ubyte.gz", "bb300cfdad3c16e7a12a480ee83cd310"), ] classes = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"] class KMNIST(MNIST): """`Kuzushiji-MNIST `_ Dataset. Args: root (string): Root directory of dataset where ``KMNIST/raw/train-images-idx3-ubyte`` and ``KMNIST/raw/t10k-images-idx3-ubyte`` exist. train (bool, optional): If True, creates dataset from ``train-images-idx3-ubyte``, otherwise from ``t10k-images-idx3-ubyte``. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ mirrors = ["http://codh.rois.ac.jp/kmnist/dataset/kmnist/"] resources = [ ("train-images-idx3-ubyte.gz", "bdb82020997e1d708af4cf47b453dcf7"), ("train-labels-idx1-ubyte.gz", "e144d726b3acfaa3e44228e80efcd344"), ("t10k-images-idx3-ubyte.gz", "5c965bf0a639b31b8f53240b1b52f4d7"), ("t10k-labels-idx1-ubyte.gz", "7320c461ea6c1c855c0b718fb2a4b134"), ] classes = ["o", "ki", "su", "tsu", "na", "ha", "ma", "ya", "re", "wo"] class EMNIST(MNIST): """`EMNIST `_ Dataset. Args: root (string): Root directory of dataset where ``EMNIST/raw/train-images-idx3-ubyte`` and ``EMNIST/raw/t10k-images-idx3-ubyte`` exist. split (string): The dataset has 6 different splits: ``byclass``, ``bymerge``, ``balanced``, ``letters``, ``digits`` and ``mnist``. This argument specifies which one to use. train (bool, optional): If True, creates dataset from ``training.pt``, otherwise from ``test.pt``. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. """ url = "https://www.itl.nist.gov/iaui/vip/cs_links/EMNIST/gzip.zip" md5 = "58c8d27c78d21e728a6bc7b3cc06412e" splits = ("byclass", "bymerge", "balanced", "letters", "digits", "mnist") # Merged Classes assumes Same structure for both uppercase and lowercase version _merged_classes = {"c", "i", "j", "k", "l", "m", "o", "p", "s", "u", "v", "w", "x", "y", "z"} _all_classes = set(string.digits + string.ascii_letters) classes_split_dict = { "byclass": sorted(list(_all_classes)), "bymerge": sorted(list(_all_classes - _merged_classes)), "balanced": sorted(list(_all_classes - _merged_classes)), "letters": ["N/A"] + list(string.ascii_lowercase), "digits": list(string.digits), "mnist": list(string.digits), } def __init__(self, root: str, split: str, **kwargs: Any) -> None: self.split = verify_str_arg(split, "split", self.splits) self.training_file = self._training_file(split) self.test_file = self._test_file(split) super().__init__(root, **kwargs) self.classes = self.classes_split_dict[self.split] @staticmethod def _training_file(split) -> str: return f"training_{split}.pt" @staticmethod def _test_file(split) -> str: return f"test_{split}.pt" @property def _file_prefix(self) -> str: return f"emnist-{self.split}-{'train' if self.train else 'test'}" @property def images_file(self) -> str: return os.path.join(self.raw_folder, f"{self._file_prefix}-images-idx3-ubyte") @property def labels_file(self) -> str: return os.path.join(self.raw_folder, f"{self._file_prefix}-labels-idx1-ubyte") def _load_data(self): return read_image_file(self.images_file), read_label_file(self.labels_file) def _check_exists(self) -> bool: return all(check_integrity(file) for file in (self.images_file, self.labels_file)) def download(self) -> None: """Download the EMNIST data if it doesn't exist already.""" if self._check_exists(): return os.makedirs(self.raw_folder, exist_ok=True) download_and_extract_archive(self.url, download_root=self.raw_folder, md5=self.md5) gzip_folder = os.path.join(self.raw_folder, "gzip") for gzip_file in os.listdir(gzip_folder): if gzip_file.endswith(".gz"): extract_archive(os.path.join(gzip_folder, gzip_file), self.raw_folder) shutil.rmtree(gzip_folder) class QMNIST(MNIST): """`QMNIST `_ Dataset. Args: root (string): Root directory of dataset whose ``raw`` subdir contains binary files of the datasets. what (string,optional): Can be 'train', 'test', 'test10k', 'test50k', or 'nist' for respectively the mnist compatible training set, the 60k qmnist testing set, the 10k qmnist examples that match the mnist testing set, the 50k remaining qmnist testing examples, or all the nist digits. The default is to select 'train' or 'test' according to the compatibility argument 'train'. compat (bool,optional): A boolean that says whether the target for each example is class number (for compatibility with the MNIST dataloader) or a torch vector containing the full qmnist information. Default=True. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. train (bool,optional,compatibility): When argument 'what' is not specified, this boolean decides whether to load the training set ot the testing set. Default: True. """ subsets = {"train": "train", "test": "test", "test10k": "test", "test50k": "test", "nist": "nist"} resources: Dict[str, List[Tuple[str, str]]] = { # type: ignore[assignment] "train": [ ( "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-images-idx3-ubyte.gz", "ed72d4157d28c017586c42bc6afe6370", ), ( "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-labels-idx2-int.gz", "0058f8dd561b90ffdd0f734c6a30e5e4", ), ], "test": [ ( "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-images-idx3-ubyte.gz", "1394631089c404de565df7b7aeaf9412", ), ( "https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-labels-idx2-int.gz", "5b5b05890a5e13444e108efe57b788aa", ), ], "nist": [ ( "https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-images-idx3-ubyte.xz", "7f124b3b8ab81486c9d8c2749c17f834", ), ( "https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-labels-idx2-int.xz", "5ed0e788978e45d4a8bd4b7caec3d79d", ), ], } classes = [ "0 - zero", "1 - one", "2 - two", "3 - three", "4 - four", "5 - five", "6 - six", "7 - seven", "8 - eight", "9 - nine", ] def __init__( self, root: str, what: Optional[str] = None, compat: bool = True, train: bool = True, **kwargs: Any ) -> None: if what is None: what = "train" if train else "test" self.what = verify_str_arg(what, "what", tuple(self.subsets.keys())) self.compat = compat self.data_file = what + ".pt" self.training_file = self.data_file self.test_file = self.data_file super().__init__(root, train, **kwargs) @property def images_file(self) -> str: (url, _), _ = self.resources[self.subsets[self.what]] return os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0]) @property def labels_file(self) -> str: _, (url, _) = self.resources[self.subsets[self.what]] return os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0]) def _check_exists(self) -> bool: return all(check_integrity(file) for file in (self.images_file, self.labels_file)) def _load_data(self): data = read_sn3_pascalvincent_tensor(self.images_file) if data.dtype != torch.uint8: raise TypeError(f"data should be of dtype torch.uint8 instead of {data.dtype}") if data.ndimension() != 3: raise ValueError("data should have 3 dimensions instead of {data.ndimension()}") targets = read_sn3_pascalvincent_tensor(self.labels_file).long() if targets.ndimension() != 2: raise ValueError(f"targets should have 2 dimensions instead of {targets.ndimension()}") if self.what == "test10k": data = data[0:10000, :, :].clone() targets = targets[0:10000, :].clone() elif self.what == "test50k": data = data[10000:, :, :].clone() targets = targets[10000:, :].clone() return data, targets def download(self) -> None: """Download the QMNIST data if it doesn't exist already. Note that we only download what has been asked for (argument 'what'). """ if self._check_exists(): return os.makedirs(self.raw_folder, exist_ok=True) split = self.resources[self.subsets[self.what]] for url, md5 in split: download_and_extract_archive(url, self.raw_folder, md5=md5) def __getitem__(self, index: int) -> Tuple[Any, Any]: # redefined to handle the compat flag img, target = self.data[index], self.targets[index] img = Image.fromarray(img.numpy(), mode="L") if self.transform is not None: img = self.transform(img) if self.compat: target = int(target[0]) if self.target_transform is not None: target = self.target_transform(target) return img, target def extra_repr(self) -> str: return f"Split: {self.what}" def get_int(b: bytes) -> int: return int(codecs.encode(b, "hex"), 16) SN3_PASCALVINCENT_TYPEMAP = { 8: torch.uint8, 9: torch.int8, 11: torch.int16, 12: torch.int32, 13: torch.float32, 14: torch.float64, } def read_sn3_pascalvincent_tensor(path: str, strict: bool = True) -> torch.Tensor: """Read a SN3 file in "Pascal Vincent" format (Lush file 'libidx/idx-io.lsh'). Argument may be a filename, compressed filename, or file object. """ # read with open(path, "rb") as f: data = f.read() # parse magic = get_int(data[0:4]) nd = magic % 256 ty = magic // 256 assert 1 <= nd <= 3 assert 8 <= ty <= 14 torch_type = SN3_PASCALVINCENT_TYPEMAP[ty] s = [get_int(data[4 * (i + 1) : 4 * (i + 2)]) for i in range(nd)] num_bytes_per_value = torch.iinfo(torch_type).bits // 8 # The MNIST format uses the big endian byte order. If the system uses little endian byte order by default, # we need to reverse the bytes before we can read them with torch.frombuffer(). needs_byte_reversal = sys.byteorder == "little" and num_bytes_per_value > 1 parsed = torch.frombuffer(bytearray(data), dtype=torch_type, offset=(4 * (nd + 1))) if needs_byte_reversal: parsed = parsed.flip(0) assert parsed.shape[0] == np.prod(s) or not strict return parsed.view(*s) def read_label_file(path: str) -> torch.Tensor: x = read_sn3_pascalvincent_tensor(path, strict=False) if x.dtype != torch.uint8: raise TypeError(f"x should be of dtype torch.uint8 instead of {x.dtype}") if x.ndimension() != 1: raise ValueError(f"x should have 1 dimension instead of {x.ndimension()}") return x.long() def read_image_file(path: str) -> torch.Tensor: x = read_sn3_pascalvincent_tensor(path, strict=False) if x.dtype != torch.uint8: raise TypeError(f"x should be of dtype torch.uint8 instead of {x.dtype}") if x.ndimension() != 3: raise ValueError(f"x should have 3 dimension instead of {x.ndimension()}") return x msadapter/ms_adapter/torchvision/datasets/omniglot.py000066400000000000000000000077731436713554100235650ustar00rootroot00000000000000from os.path import join from typing import Any, Callable, List, Optional, Tuple from PIL import Image from .utils import download_and_extract_archive, check_integrity, list_dir, list_files from .vision import VisionDataset class Omniglot(VisionDataset): """`Omniglot `_ Dataset. Args: root (string): Root directory of dataset where directory ``omniglot-py`` exists. background (bool, optional): If True, creates dataset from the "background" set, otherwise creates from the "evaluation" set. This terminology is defined by the authors. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset zip files from the internet and puts it in root directory. If the zip files are already downloaded, they are not downloaded again. """ folder = "omniglot-py" download_url_prefix = "https://raw.githubusercontent.com/brendenlake/omniglot/master/python" zips_md5 = { "images_background": "68d2efa1b9178cc56df9314c21c6e718", "images_evaluation": "6b91aef0f799c5bb55b94e3f2daec811", } def __init__( self, root: str, background: bool = True, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(join(root, self.folder), transform=transform, target_transform=target_transform) self.background = background if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") self.target_folder = join(self.root, self._get_target_folder()) self._alphabets = list_dir(self.target_folder) self._characters: List[str] = sum( ([join(a, c) for c in list_dir(join(self.target_folder, a))] for a in self._alphabets), [] ) self._character_images = [ [(image, idx) for image in list_files(join(self.target_folder, character), ".png")] for idx, character in enumerate(self._characters) ] self._flat_character_images: List[Tuple[str, int]] = sum(self._character_images, []) def __len__(self) -> int: return len(self._flat_character_images) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target character class. """ image_name, character_class = self._flat_character_images[index] image_path = join(self.target_folder, self._characters[character_class], image_name) image = Image.open(image_path, mode="r").convert("L") if self.transform: image = self.transform(image) if self.target_transform: character_class = self.target_transform(character_class) return image, character_class def _check_integrity(self) -> bool: zip_filename = self._get_target_folder() if not check_integrity(join(self.root, zip_filename + ".zip"), self.zips_md5[zip_filename]): return False return True def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return filename = self._get_target_folder() zip_filename = filename + ".zip" url = self.download_url_prefix + "/" + zip_filename download_and_extract_archive(url, self.root, filename=zip_filename, md5=self.zips_md5[filename]) def _get_target_folder(self) -> str: return "images_background" if self.background else "images_evaluation" msadapter/ms_adapter/torchvision/datasets/oxford_iiit_pet.py000066400000000000000000000117171436713554100251150ustar00rootroot00000000000000import os import os.path import pathlib from typing import Any, Callable, Optional, Union, Tuple from typing import Sequence from PIL import Image from .utils import download_and_extract_archive, verify_str_arg from .vision import VisionDataset class OxfordIIITPet(VisionDataset): """`Oxford-IIIT Pet Dataset `_. Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"trainval"`` (default) or ``"test"``. target_types (string, sequence of strings, optional): Types of target to use. Can be ``category`` (default) or ``segmentation``. Can also be a list to output a tuple with all specified target types. The types represent: - ``category`` (int): Label for one of the 37 pet categories. - ``segmentation`` (PIL image): Segmentation trimap of the image. If empty, ``None`` will be returned as target. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it into ``root/oxford-iiit-pet``. If dataset is already downloaded, it is not downloaded again. """ _RESOURCES = ( ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz", "5c4f3ee8e5d25df40f4fd59a7f44e54c"), ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz", "95a8c909bbe2e81eed6a22bccdf3f68f"), ) _VALID_TARGET_TYPES = ("category", "segmentation") def __init__( self, root: str, split: str = "trainval", target_types: Union[Sequence[str], str] = "category", transforms: Optional[Callable] = None, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ): self._split = verify_str_arg(split, "split", ("trainval", "test")) if isinstance(target_types, str): target_types = [target_types] self._target_types = [ verify_str_arg(target_type, "target_types", self._VALID_TARGET_TYPES) for target_type in target_types ] super().__init__(root, transforms=transforms, transform=transform, target_transform=target_transform) self._base_folder = pathlib.Path(self.root) / "oxford-iiit-pet" self._images_folder = self._base_folder / "images" self._anns_folder = self._base_folder / "annotations" self._segs_folder = self._anns_folder / "trimaps" if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") image_ids = [] self._labels = [] with open(self._anns_folder / f"{self._split}.txt") as file: for line in file: image_id, label, *_ = line.strip().split() image_ids.append(image_id) self._labels.append(int(label) - 1) self.classes = [ " ".join(part.title() for part in raw_cls.split("_")) for raw_cls, _ in sorted( {(image_id.rsplit("_", 1)[0], label) for image_id, label in zip(image_ids, self._labels)}, key=lambda image_id_and_label: image_id_and_label[1], ) ] self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) self._images = [self._images_folder / f"{image_id}.jpg" for image_id in image_ids] self._segs = [self._segs_folder / f"{image_id}.png" for image_id in image_ids] def __len__(self) -> int: return len(self._images) def __getitem__(self, idx: int) -> Tuple[Any, Any]: image = Image.open(self._images[idx]).convert("RGB") target: Any = [] for target_type in self._target_types: if target_type == "category": target.append(self._labels[idx]) else: # target_type == "segmentation" target.append(Image.open(self._segs[idx])) if not target: target = None elif len(target) == 1: target = target[0] else: target = tuple(target) if self.transforms: image, target = self.transforms(image, target) return image, target def _check_exists(self) -> bool: for folder in (self._images_folder, self._anns_folder): if not (os.path.exists(folder) and os.path.isdir(folder)): return False else: return True def _download(self) -> None: if self._check_exists(): return for url, md5 in self._RESOURCES: download_and_extract_archive(url, download_root=str(self._base_folder), md5=md5) msadapter/ms_adapter/torchvision/datasets/pcam.py000066400000000000000000000124111436713554100226360ustar00rootroot00000000000000import pathlib import warnings from typing import Any, Callable, Optional, Tuple from PIL import Image from .utils import download_file_from_google_drive, _decompress, verify_str_arg from .vision import VisionDataset class PCAM(VisionDataset): """`PCAM Dataset `_. The PatchCamelyon dataset is a binary classification dataset with 327,680 color images (96px x 96px), extracted from histopathologic scans of lymph node sections. Each image is annotated with a binary label indicating presence of metastatic tissue. This dataset requires the ``h5py`` package which you can install with ``pip install h5py``. Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"train"`` (default), ``"test"`` or ``"val"``. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it into ``root/pcam``. If dataset is already downloaded, it is not downloaded again. """ _FILES = { "train": { "images": ( "camelyonpatch_level_2_split_train_x.h5", # Data file name "1Ka0XfEMiwgCYPdTI-vv6eUElOBnKFKQ2", # Google Drive ID "1571f514728f59376b705fc836ff4b63", # md5 hash ), "targets": ( "camelyonpatch_level_2_split_train_y.h5", "1269yhu3pZDP8UYFQs-NYs3FPwuK-nGSG", "35c2d7259d906cfc8143347bb8e05be7", ), }, "test": { "images": ( "camelyonpatch_level_2_split_test_x.h5", "1qV65ZqZvWzuIVthK8eVDhIwrbnsJdbg_", "d5b63470df7cfa627aeec8b9dc0c066e", ), "targets": ( "camelyonpatch_level_2_split_test_y.h5", "17BHrSrwWKjYsOgTMmoqrIjDy6Fa2o_gP", "2b85f58b927af9964a4c15b8f7e8f179", ), }, "val": { "images": ( "camelyonpatch_level_2_split_valid_x.h5", "1hgshYGWK8V-eGRy8LToWJJgDU_rXWVJ3", "d8c2d60d490dbd479f8199bdfa0cf6ec", ), "targets": ( "camelyonpatch_level_2_split_valid_y.h5", "1bH8ZRbhSVAhScTS0p9-ZzGnX91cHT3uO", "60a7035772fbdb7f34eb86d4420cf66a", ), }, } def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ): try: import h5py self.h5py = h5py except ImportError: raise RuntimeError( "h5py is not found. This dataset needs to have h5py installed: please run pip install h5py" ) self._split = verify_str_arg(split, "split", ("train", "test", "val")) super().__init__(root, transform=transform, target_transform=target_transform) self._base_folder = pathlib.Path(self.root) / "pcam" if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") def __len__(self) -> int: images_file = self._FILES[self._split]["images"][0] with self.h5py.File(self._base_folder / images_file) as images_data: return images_data["x"].shape[0] def __getitem__(self, idx: int) -> Tuple[Any, Any]: images_file = self._FILES[self._split]["images"][0] with self.h5py.File(self._base_folder / images_file) as images_data: image = Image.fromarray(images_data["x"][idx]).convert("RGB") targets_file = self._FILES[self._split]["targets"][0] with self.h5py.File(self._base_folder / targets_file) as targets_data: target = int(targets_data["y"][idx, 0, 0, 0]) # shape is [num_images, 1, 1, 1] if self.transform: image = self.transform(image) if self.target_transform: target = self.target_transform(target) return image, target def _check_exists(self) -> bool: images_file = self._FILES[self._split]["images"][0] targets_file = self._FILES[self._split]["targets"][0] return all(self._base_folder.joinpath(h5_file).exists() for h5_file in (images_file, targets_file)) def _download(self) -> None: if self._check_exists(): return warnings.warn("The dataset is stored on google drive, if you can't download it from google drive, " "please download it from the official website." "PCAM Dataset ") for file_name, file_id, md5 in self._FILES[self._split].values(): archive_name = file_name + ".gz" download_file_from_google_drive(file_id, str(self._base_folder), filename=archive_name, md5=md5) _decompress(str(self._base_folder / archive_name)) msadapter/ms_adapter/torchvision/datasets/phototour.py000066400000000000000000000201321436713554100237600ustar00rootroot00000000000000import os from typing import Any, Callable, List, Optional, Tuple, Union import numpy as np import ms_adapter.pytorch as torch from PIL import Image from .utils import download_url from .vision import VisionDataset class PhotoTour(VisionDataset): """`Multi-view Stereo Correspondence `_ Dataset. .. note:: We only provide the newer version of the dataset, since the authors state that it is more suitable for training descriptors based on difference of Gaussian, or Harris corners, as the patches are centred on real interest point detections, rather than being projections of 3D points as is the case in the old dataset. The original dataset is available under http://phototour.cs.washington.edu/patches/default.htm. Args: root (string): Root directory where images are. name (string): Name of the dataset to load. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ urls = { "notredame_harris": [ "http://matthewalunbrown.com/patchdata/notredame_harris.zip", "notredame_harris.zip", "69f8c90f78e171349abdf0307afefe4d", ], "yosemite_harris": [ "http://matthewalunbrown.com/patchdata/yosemite_harris.zip", "yosemite_harris.zip", "a73253d1c6fbd3ba2613c45065c00d46", ], "liberty_harris": [ "http://matthewalunbrown.com/patchdata/liberty_harris.zip", "liberty_harris.zip", "c731fcfb3abb4091110d0ae8c7ba182c", ], "notredame": [ "http://icvl.ee.ic.ac.uk/vbalnt/notredame.zip", "notredame.zip", "509eda8535847b8c0a90bbb210c83484", ], "yosemite": ["http://icvl.ee.ic.ac.uk/vbalnt/yosemite.zip", "yosemite.zip", "533b2e8eb7ede31be40abc317b2fd4f0"], "liberty": ["http://icvl.ee.ic.ac.uk/vbalnt/liberty.zip", "liberty.zip", "fdd9152f138ea5ef2091746689176414"], } means = { "notredame": 0.4854, "yosemite": 0.4844, "liberty": 0.4437, "notredame_harris": 0.4854, "yosemite_harris": 0.4844, "liberty_harris": 0.4437, } stds = { "notredame": 0.1864, "yosemite": 0.1818, "liberty": 0.2019, "notredame_harris": 0.1864, "yosemite_harris": 0.1818, "liberty_harris": 0.2019, } lens = { "notredame": 468159, "yosemite": 633587, "liberty": 450092, "liberty_harris": 379587, "yosemite_harris": 450912, "notredame_harris": 325295, } image_ext = "bmp" info_file = "info.txt" matches_files = "m50_100000_100000_0.txt" def __init__( self, root: str, name: str, train: bool = True, transform: Optional[Callable] = None, download: bool = False ) -> None: super().__init__(root, transform=transform) self.name = name self.data_dir = os.path.join(self.root, name) self.data_down = os.path.join(self.root, f"{name}.zip") # self.data_file = os.path.join(self.root, f"{name}.pt") self.train = train self.mean = self.means[name] self.std = self.stds[name] if download: self.download() # if not self._check_datafile_exists(): # self.cache() self.data, self.labels, self.matches = self.cache() # load the serialized data # self.data, self.labels, self.matches = torch.load(self.data_file) def __getitem__(self, index: int) -> Union[torch.Tensor, Tuple[Any, Any, torch.Tensor]]: """ Args: index (int): Index Returns: tuple: (data1, data2, matches) """ if self.train: data = self.data[index] if self.transform is not None: data = self.transform(data) return data m = self.matches[index] data1, data2 = self.data[m[0]], self.data[m[1]] if self.transform is not None: data1 = self.transform(data1) data2 = self.transform(data2) return data1, data2, m[2] def __len__(self) -> int: return len(self.data if self.train else self.matches) # def _check_datafile_exists(self) -> bool: # return os.path.exists(self.data_file) def _check_downloaded(self) -> bool: return os.path.exists(self.data_dir) def download(self) -> None: # if self._check_datafile_exists(): # print(f"# Found cached data {self.data_file}") # return if not self._check_downloaded(): # download files url = self.urls[self.name][0] filename = self.urls[self.name][1] md5 = self.urls[self.name][2] fpath = os.path.join(self.root, filename) download_url(url, self.root, filename, md5) print(f"# Extracting data {self.data_down}\n") import zipfile with zipfile.ZipFile(fpath, "r") as z: z.extractall(self.data_dir) os.unlink(fpath) def cache(self): # process and save as torch files # print(f"# Caching data {self.data_file}") # dataset = ( # read_image_file(self.data_dir, self.image_ext, self.lens[self.name]), # read_info_file(self.data_dir, self.info_file), # read_matches_files(self.data_dir, self.matches_files), # ) data = read_image_file(self.data_dir, self.image_ext, self.lens[self.name]) info = read_info_file(self.data_dir, self.info_file) matches = read_matches_files(self.data_dir, self.matches_files) return data, info, matches # with open(self.data_file, "wb") as f: # torch.save(dataset, f) def extra_repr(self) -> str: split = "Train" if self.train is True else "Test" return f"Split: {split}" def read_image_file(data_dir: str, image_ext: str, n: int) -> torch.Tensor: """Return a Tensor containing the patches""" def PIL2array(_img: Image.Image) -> np.ndarray: """Convert PIL image type to numpy 2D array""" return np.array(_img.getdata(), dtype=np.uint8).reshape(64, 64) def find_files(_data_dir: str, _image_ext: str) -> List[str]: """Return a list with the file names of the images containing the patches""" files = [] # find those files with the specified extension for file_dir in os.listdir(_data_dir): if file_dir.endswith(_image_ext): files.append(os.path.join(_data_dir, file_dir)) return sorted(files) # sort files in ascend order to keep relations patches = [] list_files = find_files(data_dir, image_ext) for fpath in list_files: img = Image.open(fpath) for y in range(0, img.height, 64): for x in range(0, img.width, 64): patch = img.crop((x, y, x + 64, y + 64)) patches.append(PIL2array(patch)) return torch.ByteTensor(np.array(patches[:n])) def read_info_file(data_dir: str, info_file: str) -> torch.Tensor: """Return a Tensor containing the list of labels Read the file and keep only the ID of the 3D point. """ with open(os.path.join(data_dir, info_file)) as f: labels = [int(line.split()[0]) for line in f] return torch.LongTensor(labels) def read_matches_files(data_dir: str, matches_file: str) -> torch.Tensor: """Return a Tensor containing the ground truth matches Read the file and keep only 3D point ID. Matches are represented with a 1, non matches with a 0. """ matches = [] with open(os.path.join(data_dir, matches_file)) as f: for line in f: line_split = line.split() matches.append([int(line_split[0]), int(line_split[3]), int(line_split[1] == line_split[4])]) return torch.LongTensor(matches) msadapter/ms_adapter/torchvision/datasets/places365.py000066400000000000000000000160411436713554100234260ustar00rootroot00000000000000import os from os import path from typing import Any, Callable, Dict, List, Optional, Tuple from urllib.parse import urljoin from .folder import default_loader from .utils import verify_str_arg, check_integrity, download_and_extract_archive from .vision import VisionDataset class Places365(VisionDataset): r"""`Places365 `_ classification dataset. Args: root (string): Root directory of the Places365 dataset. split (string, optional): The dataset split. Can be one of ``train-standard`` (default), ``train-challenge``, ``val``. small (bool, optional): If ``True``, uses the small images, i. e. resized to 256 x 256 pixels, instead of the high resolution ones. download (bool, optional): If ``True``, downloads the dataset components and places them in ``root``. Already downloaded archives are not downloaded again. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. loader (callable, optional): A function to load an image given its path. Attributes: classes (list): List of the class names. class_to_idx (dict): Dict with items (class_name, class_index). imgs (list): List of (image path, class_index) tuples targets (list): The class_index value for each image in the dataset Raises: RuntimeError: If ``download is False`` and the meta files, i. e. the devkit, are not present or corrupted. RuntimeError: If ``download is True`` and the image archive is already extracted. """ _SPLITS = ("train-standard", "train-challenge", "val") _BASE_URL = "http://data.csail.mit.edu/places/places365/" # {variant: (archive, md5)} _DEVKIT_META = { "standard": ("filelist_places365-standard.tar", "35a0585fee1fa656440f3ab298f8479c"), "challenge": ("filelist_places365-challenge.tar", "70a8307e459c3de41690a7c76c931734"), } # (file, md5) _CATEGORIES_META = ("categories_places365.txt", "06c963b85866bd0649f97cb43dd16673") # {split: (file, md5)} _FILE_LIST_META = { "train-standard": ("places365_train_standard.txt", "30f37515461640559006b8329efbed1a"), "train-challenge": ("places365_train_challenge.txt", "b2931dc997b8c33c27e7329c073a6b57"), "val": ("places365_val.txt", "e9f2fd57bfd9d07630173f4e8708e4b1"), } # {(split, small): (file, md5)} _IMAGES_META = { ("train-standard", False): ("train_large_places365standard.tar", "67e186b496a84c929568076ed01a8aa1"), ("train-challenge", False): ("train_large_places365challenge.tar", "605f18e68e510c82b958664ea134545f"), ("val", False): ("val_large.tar", "9b71c4993ad89d2d8bcbdc4aef38042f"), ("train-standard", True): ("train_256_places365standard.tar", "53ca1c756c3d1e7809517cc47c5561c5"), ("train-challenge", True): ("train_256_places365challenge.tar", "741915038a5e3471ec7332404dfb64ef"), ("val", True): ("val_256.tar", "e27b17d8d44f4af9a78502beb927f808"), } def __init__( self, root: str, split: str = "train-standard", small: bool = False, download: bool = False, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, loader: Callable[[str], Any] = default_loader, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.split = self._verify_split(split) self.small = small self.loader = loader self.classes, self.class_to_idx = self.load_categories(download) self.imgs, self.targets = self.load_file_list(download) if download: self.download_images() def __getitem__(self, index: int) -> Tuple[Any, Any]: file, target = self.imgs[index] image = self.loader(file) if self.transforms is not None: image, target = self.transforms(image, target) return image, target def __len__(self) -> int: return len(self.imgs) @property def variant(self) -> str: return "challenge" if "challenge" in self.split else "standard" @property def images_dir(self) -> str: size = "256" if self.small else "large" if self.split.startswith("train"): dir = f"data_{size}_{self.variant}" else: dir = f"{self.split}_{size}" return path.join(self.root, dir) def load_categories(self, download: bool = True) -> Tuple[List[str], Dict[str, int]]: def process(line: str) -> Tuple[str, int]: cls, idx = line.split() return cls, int(idx) file, md5 = self._CATEGORIES_META file = path.join(self.root, file) if not self._check_integrity(file, md5, download): self.download_devkit() with open(file) as fh: class_to_idx = dict(process(line) for line in fh) return sorted(class_to_idx.keys()), class_to_idx def load_file_list(self, download: bool = True) -> Tuple[List[Tuple[str, int]], List[int]]: def process(line: str, sep="/") -> Tuple[str, int]: image, idx = line.split() return path.join(self.images_dir, image.lstrip(sep).replace(sep, os.sep)), int(idx) file, md5 = self._FILE_LIST_META[self.split] file = path.join(self.root, file) if not self._check_integrity(file, md5, download): self.download_devkit() with open(file) as fh: images = [process(line) for line in fh] _, targets = zip(*images) return images, list(targets) def download_devkit(self) -> None: file, md5 = self._DEVKIT_META[self.variant] download_and_extract_archive(urljoin(self._BASE_URL, file), self.root, md5=md5) def download_images(self) -> None: if path.exists(self.images_dir): raise RuntimeError( f"The directory {self.images_dir} already exists. If you want to re-download or re-extract the images, " f"delete the directory." ) file, md5 = self._IMAGES_META[(self.split, self.small)] download_and_extract_archive(urljoin(self._BASE_URL, file), self.root, md5=md5) if self.split.startswith("train"): os.rename(self.images_dir.rsplit("_", 1)[0], self.images_dir) def extra_repr(self) -> str: return "\n".join(("Split: {split}", "Small: {small}")).format(**self.__dict__) def _verify_split(self, split: str) -> str: return verify_str_arg(split, "split", self._SPLITS) def _check_integrity(self, file: str, md5: str, download: bool) -> bool: integrity = check_integrity(file, md5=md5) if not integrity and not download: raise RuntimeError( f"The file {file} does not exist or is corrupted. You can set download=True to download it." ) return integrity msadapter/ms_adapter/torchvision/datasets/rendered_sst2.py000066400000000000000000000067451436713554100244760ustar00rootroot00000000000000from pathlib import Path from typing import Any, Tuple, Callable, Optional import PIL.Image from .folder import make_dataset from .utils import verify_str_arg, download_and_extract_archive from .vision import VisionDataset class RenderedSST2(VisionDataset): """`The Rendered SST2 Dataset `_. Rendered SST2 is an image classification dataset used to evaluate the models capability on optical character recognition. This dataset was generated by rendering sentences in the Standford Sentiment Treebank v2 dataset. This dataset contains two classes (positive and negative) and is divided in three splits: a train split containing 6920 images (3610 positive and 3310 negative), a validation split containing 872 images (444 positive and 428 negative), and a test split containing 1821 images (909 positive and 912 negative). Args: root (string): Root directory of the dataset. split (string, optional): The dataset split, supports ``"train"`` (default), `"val"` and ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. Default is False. """ _URL = "https://openaipublic.azureedge.net/clip/data/rendered-sst2.tgz" _MD5 = "2384d08e9dcfa4bd55b324e610496ee5" def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self._split = verify_str_arg(split, "split", ("train", "val", "test")) self._split_to_folder = {"train": "train", "val": "valid", "test": "test"} self._base_folder = Path(self.root) / "rendered-sst2" self.classes = ["negative", "positive"] self.class_to_idx = {"negative": 0, "positive": 1} if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") self._samples = make_dataset(str(self._base_folder / self._split_to_folder[self._split]), extensions=("png",)) def __len__(self) -> int: return len(self._samples) def __getitem__(self, idx) -> Tuple[Any, Any]: image_file, label = self._samples[idx] image = PIL.Image.open(image_file).convert("RGB") if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def extra_repr(self) -> str: return f"split={self._split}" def _check_exists(self) -> bool: for class_label in set(self.classes): if not (self._base_folder / self._split_to_folder[self._split] / class_label).is_dir(): return False return True def _download(self) -> None: if self._check_exists(): return download_and_extract_archive(self._URL, download_root=self.root, md5=self._MD5) msadapter/ms_adapter/torchvision/datasets/samplers/000077500000000000000000000000001436713554100231735ustar00rootroot00000000000000msadapter/ms_adapter/torchvision/datasets/samplers/__init__.py000066400000000000000000000002411436713554100253010ustar00rootroot00000000000000from .clip_sampler import DistributedSampler, UniformClipSampler, RandomClipSampler __all__ = ("DistributedSampler", "UniformClipSampler", "RandomClipSampler") msadapter/ms_adapter/torchvision/datasets/samplers/clip_sampler.py000066400000000000000000000146071436713554100262270ustar00rootroot00000000000000import math from typing import Optional, List, Iterator, Sized, Union, cast import numpy as np import ms_adapter.pytorch as torch # import torch.distributed as dist TODO:dist import mindspore.communication as dist from ms_adapter.pytorch.utils.data import Sampler from ms_adapter.torchvision.datasets.video_utils import VideoClips class DistributedSampler(Sampler): """ Example: dataset: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] num_replicas: 4 shuffle: False when group_size = 1 RANK | shard_dataset ========================= rank_0 | [0, 4, 8, 12] rank_1 | [1, 5, 9, 13] rank_2 | [2, 6, 10, 0] rank_3 | [3, 7, 11, 1] when group_size = 2 RANK | shard_dataset ========================= rank_0 | [0, 1, 8, 9] rank_1 | [2, 3, 10, 11] rank_2 | [4, 5, 12, 13] rank_3 | [6, 7, 0, 1] """ def __init__( self, dataset: Sized, num_replicas: Optional[int] = None, rank: Optional[int] = None, shuffle: bool = False, group_size: int = 1, ) -> None: if num_replicas is None: # if not dist.is_available(): TODO:dist # raise RuntimeError("Requires distributed package to be available") # num_replicas = dist.get_world_size() TODO:dist num_replicas = dist.get_group_size() if rank is None: # if not dist.is_available(): # raise RuntimeError("Requires distributed package to be available") rank = dist.get_rank() if len(dataset) % group_size != 0: raise ValueError( f"dataset length must be a multiplier of group size dataset length: {len(dataset)}, group size: {group_size}" ) self.dataset = dataset self.group_size = group_size self.num_replicas = num_replicas self.rank = rank self.epoch = 0 dataset_group_length = len(dataset) // group_size self.num_group_samples = int(math.ceil(dataset_group_length * 1.0 / self.num_replicas)) self.num_samples = self.num_group_samples * group_size self.total_size = self.num_samples * self.num_replicas self.shuffle = shuffle def __iter__(self) -> Iterator[int]: # deterministically shuffle based on epoch # g = torch.Generator() # g.manual_seed(self.epoch) # indices: Union[torch.Tensor, List[int]] rnd = np.random.RandomState(self.epoch) # indices = rnd.choice(np.arange(len(self.dataset)), len(self.dataset), replace=False).tolist() if self.shuffle: indices = rnd.choice(np.arange(len(self.dataset)), len(self.dataset), replace=False).tolist() else: indices = list(range(len(self.dataset))) # add extra samples to make it evenly divisible indices += indices[: (self.total_size - len(indices))] assert len(indices) == self.total_size total_group_size = self.total_size // self.group_size indices = torch.reshape(torch.LongTensor(indices), (total_group_size, self.group_size)) # subsample indices = indices[self.rank : total_group_size : self.num_replicas, :] indices = torch.reshape(indices, (-1,)).numpy().tolist() assert len(indices) == self.num_samples if isinstance(self.dataset, Sampler): orig_indices = list(iter(self.dataset)) indices = [orig_indices[i] for i in indices] return iter(indices) def __len__(self) -> int: return self.num_samples def set_epoch(self, epoch: int) -> None: self.epoch = epoch class UniformClipSampler(Sampler): """ Sample `num_video_clips_per_video` clips for each video, equally spaced. When number of unique clips in the video is fewer than num_video_clips_per_video, repeat the clips until `num_video_clips_per_video` clips are collected Args: video_clips (VideoClips): video clips to sample from num_clips_per_video (int): number of clips to be sampled per video """ def __init__(self, video_clips: VideoClips, num_clips_per_video: int) -> None: if not isinstance(video_clips, VideoClips): raise TypeError(f"Expected video_clips to be an instance of VideoClips, got {type(video_clips)}") self.video_clips = video_clips self.num_clips_per_video = num_clips_per_video def __iter__(self) -> Iterator[int]: idxs = [] s = 0 # select num_clips_per_video for each video, uniformly spaced for c in self.video_clips.clips: length = len(c) if length == 0: # corner case where video decoding fails continue sampled = torch.linspace(s, s + length - 1, steps=self.num_clips_per_video).floor().to(torch.int64) s += length idxs.append(sampled) return iter(cast(List[int], torch.cat(idxs).numpy().tolist())) def __len__(self) -> int: return sum(self.num_clips_per_video for c in self.video_clips.clips if len(c) > 0) class RandomClipSampler(Sampler): """ Samples at most `max_video_clips_per_video` clips for each video randomly Args: video_clips (VideoClips): video clips to sample from max_clips_per_video (int): maximum number of clips to be sampled per video """ def __init__(self, video_clips: VideoClips, max_clips_per_video: int) -> None: if not isinstance(video_clips, VideoClips): raise TypeError(f"Expected video_clips to be an instance of VideoClips, got {type(video_clips)}") self.video_clips = video_clips self.max_clips_per_video = max_clips_per_video def __iter__(self) -> Iterator[int]: idxs = [] s = 0 # select at most max_clips_per_video for each video, randomly for c in self.video_clips.clips: length = len(c) size = min(length, self.max_clips_per_video) sampled = torch.randperm(length)[:size] + s s += length idxs.append(sampled) idxs_ = torch.cat(idxs) # shuffle all clips randomly perm = torch.randperm(len(idxs_)) return iter(idxs_[perm].tolist()) def __len__(self) -> int: return sum(min(len(c), self.max_clips_per_video) for c in self.video_clips.clips) msadapter/ms_adapter/torchvision/datasets/sbd.py000066400000000000000000000121221436713554100224650ustar00rootroot00000000000000import os import shutil from typing import Any, Callable, Optional, Tuple import numpy as np from PIL import Image from .utils import download_url, verify_str_arg, download_and_extract_archive from .vision import VisionDataset class SBDataset(VisionDataset): """`Semantic Boundaries Dataset `_ The SBD currently contains annotations from 11355 images taken from the PASCAL VOC 2011 dataset. .. note :: Please note that the train and val splits included with this dataset are different from the splits in the PASCAL VOC dataset. In particular some "train" images might be part of VOC2012 val. If you are interested in testing on VOC 2012 val, then use `image_set='train_noval'`, which excludes all val images. .. warning:: This class needs `scipy `_ to load target files from `.mat` format. Args: root (string): Root directory of the Semantic Boundaries Dataset image_set (string, optional): Select the image_set to use, ``train``, ``val`` or ``train_noval``. Image set ``train_noval`` excludes VOC 2012 val images. mode (string, optional): Select target type. Possible values 'boundaries' or 'segmentation'. In case of 'boundaries', the target is an array of shape `[num_classes, H, W]`, where `num_classes=20`. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. Input sample is PIL image and target is a numpy array if `mode='boundaries'` or PIL image if `mode='segmentation'`. """ url = "https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz" md5 = "82b4d87ceb2ed10f6038a1cba92111cb" filename = "benchmark.tgz" voc_train_url = "http://home.bharathh.info/pubs/codes/SBD/train_noval.txt" voc_split_filename = "train_noval.txt" voc_split_md5 = "79bff800c5f0b1ec6b21080a3c066722" def __init__( self, root: str, image_set: str = "train", mode: str = "boundaries", download: bool = False, transforms: Optional[Callable] = None, ) -> None: try: from scipy.io import loadmat self._loadmat = loadmat except ImportError: raise RuntimeError("Scipy is not found. This dataset needs to have scipy installed: pip install scipy") super().__init__(root, transforms) self.image_set = verify_str_arg(image_set, "image_set", ("train", "val", "train_noval")) self.mode = verify_str_arg(mode, "mode", ("segmentation", "boundaries")) self.num_classes = 20 sbd_root = self.root image_dir = os.path.join(sbd_root, "img") mask_dir = os.path.join(sbd_root, "cls") if download: download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.md5) extracted_ds_root = os.path.join(self.root, "benchmark_RELEASE", "dataset") for f in ["cls", "img", "inst", "train.txt", "val.txt"]: old_path = os.path.join(extracted_ds_root, f) shutil.move(old_path, sbd_root) download_url(self.voc_train_url, sbd_root, self.voc_split_filename, self.voc_split_md5) if not os.path.isdir(sbd_root): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") split_f = os.path.join(sbd_root, image_set.rstrip("\n") + ".txt") with open(os.path.join(split_f)) as fh: file_names = [x.strip() for x in fh.readlines()] self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names] self.masks = [os.path.join(mask_dir, x + ".mat") for x in file_names] self._get_target = self._get_segmentation_target if self.mode == "segmentation" else self._get_boundaries_target def _get_segmentation_target(self, filepath: str) -> Image.Image: mat = self._loadmat(filepath) return Image.fromarray(mat["GTcls"][0]["Segmentation"][0]) def _get_boundaries_target(self, filepath: str) -> np.ndarray: mat = self._loadmat(filepath) return np.concatenate( [np.expand_dims(mat["GTcls"][0]["Boundaries"][0][i][0].toarray(), axis=0) for i in range(self.num_classes)], axis=0, ) def __getitem__(self, index: int) -> Tuple[Any, Any]: img = Image.open(self.images[index]).convert("RGB") target = self._get_target(self.masks[index]) if self.transforms is not None: img, target = self.transforms(img, target) return img, target def __len__(self) -> int: return len(self.images) def extra_repr(self) -> str: lines = ["Image set: {image_set}", "Mode: {mode}"] return "\n".join(lines).format(**self.__dict__) msadapter/ms_adapter/torchvision/datasets/sbu.py000066400000000000000000000101531436713554100225100ustar00rootroot00000000000000import os from typing import Any, Callable, Optional, Tuple from PIL import Image from .utils import download_url, check_integrity from .vision import VisionDataset class SBU(VisionDataset): """`SBU Captioned Photo `_ Dataset. Args: root (string): Root directory of dataset where tarball ``SBUCaptionedPhotoDataset.tar.gz`` exists. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ url = "http://www.cs.virginia.edu/~vicente/sbucaptions/SBUCaptionedPhotoDataset.tar.gz" filename = "SBUCaptionedPhotoDataset.tar.gz" md5_checksum = "9aec147b3488753cf758b4d493422285" def __init__( self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = True, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") # Read the caption for each photo self.photos = [] self.captions = [] file1 = os.path.join(self.root, "dataset", "SBU_captioned_photo_dataset_urls.txt") file2 = os.path.join(self.root, "dataset", "SBU_captioned_photo_dataset_captions.txt") for line1, line2 in zip(open(file1), open(file2)): url = line1.rstrip() photo = os.path.basename(url) filename = os.path.join(self.root, "dataset", photo) if os.path.exists(filename): caption = line2.rstrip() self.photos.append(photo) self.captions.append(caption) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is a caption for the photo. """ filename = os.path.join(self.root, "dataset", self.photos[index]) img = Image.open(filename).convert("RGB") if self.transform is not None: img = self.transform(img) target = self.captions[index] if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: """The number of photos in the dataset.""" return len(self.photos) def _check_integrity(self) -> bool: """Check the md5 checksum of the downloaded tarball.""" root = self.root fpath = os.path.join(root, self.filename) if not check_integrity(fpath, self.md5_checksum): return False return True def download(self) -> None: """Download and extract the tarball, and download each individual photo.""" import tarfile if self._check_integrity(): print("Files already downloaded and verified") return download_url(self.url, self.root, self.filename, self.md5_checksum) # Extract file with tarfile.open(os.path.join(self.root, self.filename), "r:gz") as tar: tar.extractall(path=self.root) # Download individual photos with open(os.path.join(self.root, "dataset", "SBU_captioned_photo_dataset_urls.txt")) as fh: for line in fh: url = line.rstrip() try: download_url(url, os.path.join(self.root, "dataset")) except OSError: # The images point to public images on Flickr. # Note: Images might be removed by users at anytime. pass msadapter/ms_adapter/torchvision/datasets/semeion.py000066400000000000000000000060201436713554100233540ustar00rootroot00000000000000import os.path from typing import Any, Callable, Optional, Tuple import numpy as np from PIL import Image from .utils import download_url, check_integrity from .vision import VisionDataset class SEMEION(VisionDataset): r"""`SEMEION `_ Dataset. Args: root (string): Root directory of dataset where directory ``semeion.py`` exists. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ url = "http://archive.ics.uci.edu/ml/machine-learning-databases/semeion/semeion.data" filename = "semeion.data" md5_checksum = "cb545d371d2ce14ec121470795a77432" def __init__( self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = True, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") fp = os.path.join(self.root, self.filename) data = np.loadtxt(fp) # convert value to 8 bit unsigned integer # color (white #255) the pixels self.data = (data[:, :256] * 255).astype("uint8") self.data = np.reshape(self.data, (-1, 16, 16)) self.labels = np.nonzero(data[:, 256:])[1] def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target class. """ img, target = self.data[index], int(self.labels[index]) # doing this so that it is consistent with all other datasets # to return a PIL Image img = Image.fromarray(img, mode="L") if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.data) def _check_integrity(self) -> bool: root = self.root fpath = os.path.join(root, self.filename) if not check_integrity(fpath, self.md5_checksum): return False return True def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return root = self.root download_url(self.url, root, self.filename, self.md5_checksum) msadapter/ms_adapter/torchvision/datasets/stanford_cars.py000066400000000000000000000113531436713554100245520ustar00rootroot00000000000000import pathlib from typing import Callable, Optional, Any, Tuple from PIL import Image from .utils import download_and_extract_archive, download_url, verify_str_arg from .vision import VisionDataset class StanfordCars(VisionDataset): """`Stanford Cars `_ Dataset The Cars dataset contains 16,185 images of 196 classes of cars. The data is split into 8,144 training images and 8,041 testing images, where each class has been split roughly in a 50-50 split .. note:: This class needs `scipy `_ to load target files from `.mat` format. Args: root (string): Root directory of dataset split (string, optional): The dataset split, supports ``"train"`` (default) or ``"test"``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again.""" def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: try: import scipy.io as sio except ImportError: raise RuntimeError("Scipy is not found. This dataset needs to have scipy installed: pip install scipy") super().__init__(root, transform=transform, target_transform=target_transform) self._split = verify_str_arg(split, "split", ("train", "test")) self._base_folder = pathlib.Path(root) / "stanford_cars" devkit = self._base_folder / "devkit" if self._split == "train": self._annotations_mat_path = devkit / "cars_train_annos.mat" self._images_base_path = self._base_folder / "cars_train" else: self._annotations_mat_path = self._base_folder / "cars_test_annos_withlabels.mat" self._images_base_path = self._base_folder / "cars_test" if download: self.download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") self._samples = [ ( str(self._images_base_path / annotation["fname"]), annotation["class"] - 1, # Original target mapping starts from 1, hence -1 ) for annotation in sio.loadmat(self._annotations_mat_path, squeeze_me=True)["annotations"] ] self.classes = sio.loadmat(str(devkit / "cars_meta.mat"), squeeze_me=True)["class_names"].tolist() self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)} def __len__(self) -> int: return len(self._samples) def __getitem__(self, idx: int) -> Tuple[Any, Any]: """Returns pil_image and class_id for given index""" image_path, target = self._samples[idx] pil_image = Image.open(image_path).convert("RGB") if self.transform is not None: pil_image = self.transform(pil_image) if self.target_transform is not None: target = self.target_transform(target) return pil_image, target def download(self) -> None: if self._check_exists(): return download_and_extract_archive( url="https://ai.stanford.edu/~jkrause/cars/car_devkit.tgz", download_root=str(self._base_folder), md5="c3b158d763b6e2245038c8ad08e45376", ) if self._split == "train": download_and_extract_archive( url="https://ai.stanford.edu/~jkrause/car196/cars_train.tgz", download_root=str(self._base_folder), md5="065e5b463ae28d29e77c1b4b166cfe61", ) else: download_and_extract_archive( url="https://ai.stanford.edu/~jkrause/car196/cars_test.tgz", download_root=str(self._base_folder), md5="4ce7ebf6a94d07f1952d94dd34c4d501", ) download_url( url="https://ai.stanford.edu/~jkrause/car196/cars_test_annos_withlabels.mat", root=str(self._base_folder), md5="b0a2b23655a3edd16d84508592a98d10", ) def _check_exists(self) -> bool: if not (self._base_folder / "devkit").is_dir(): return False return self._annotations_mat_path.exists() and self._images_base_path.is_dir() msadapter/ms_adapter/torchvision/datasets/stl10.py000066400000000000000000000161761436713554100226750ustar00rootroot00000000000000import os.path from typing import Any, Callable, Optional, Tuple, cast import numpy as np from PIL import Image from .utils import check_integrity, download_and_extract_archive, verify_str_arg from .vision import VisionDataset class STL10(VisionDataset): """`STL10 `_ Dataset. Args: root (string): Root directory of dataset where directory ``stl10_binary`` exists. split (string): One of {'train', 'test', 'unlabeled', 'train+unlabeled'}. Accordingly dataset is selected. folds (int, optional): One of {0-9} or None. For training, loads one of the 10 pre-defined folds of 1k samples for the standard evaluation procedure. If no value is passed, loads the 5k samples. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ base_folder = "stl10_binary" url = "http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz" filename = "stl10_binary.tar.gz" tgz_md5 = "91f7769df0f17e558f3565bffb0c7dfb" class_names_file = "class_names.txt" folds_list_file = "fold_indices.txt" train_list = [ ["train_X.bin", "918c2871b30a85fa023e0c44e0bee87f"], ["train_y.bin", "5a34089d4802c674881badbb80307741"], ["unlabeled_X.bin", "5242ba1fed5e4be9e1e742405eb56ca4"], ] test_list = [["test_X.bin", "7f263ba9f9e0b06b93213547f721ac82"], ["test_y.bin", "36f9794fa4beb8a2c72628de14fa638e"]] splits = ("train", "train+unlabeled", "unlabeled", "test") def __init__( self, root: str, split: str = "train", folds: Optional[int] = None, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.split = verify_str_arg(split, "split", self.splits) self.folds = self._verify_folds(folds) if download: self.download() elif not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") # now load the picked numpy arrays self.labels: Optional[np.ndarray] if self.split == "train": self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0]) self.labels = cast(np.ndarray, self.labels) self.__load_folds(folds) elif self.split == "train+unlabeled": self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0]) self.labels = cast(np.ndarray, self.labels) self.__load_folds(folds) unlabeled_data, _ = self.__loadfile(self.train_list[2][0]) self.data = np.concatenate((self.data, unlabeled_data)) self.labels = np.concatenate((self.labels, np.asarray([-1] * unlabeled_data.shape[0]))) elif self.split == "unlabeled": self.data, _ = self.__loadfile(self.train_list[2][0]) self.labels = np.asarray([-1] * self.data.shape[0]) else: # self.split == 'test': self.data, self.labels = self.__loadfile(self.test_list[0][0], self.test_list[1][0]) class_file = os.path.join(self.root, self.base_folder, self.class_names_file) if os.path.isfile(class_file): with open(class_file) as f: self.classes = f.read().splitlines() def _verify_folds(self, folds: Optional[int]) -> Optional[int]: if folds is None: return folds elif isinstance(folds, int): if folds in range(10): return folds msg = "Value for argument folds should be in the range [0, 10), but got {}." raise ValueError(msg.format(folds)) else: msg = "Expected type None or int for argument folds, but got type {}." raise ValueError(msg.format(type(folds))) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target class. """ target: Optional[int] if self.labels is not None: img, target = self.data[index], int(self.labels[index]) else: img, target = self.data[index], None # doing this so that it is consistent with all other datasets # to return a PIL Image img = Image.fromarray(np.transpose(img, (1, 2, 0))) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return self.data.shape[0] def __loadfile(self, data_file: str, labels_file: Optional[str] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]: labels = None if labels_file: path_to_labels = os.path.join(self.root, self.base_folder, labels_file) with open(path_to_labels, "rb") as f: labels = np.fromfile(f, dtype=np.uint8) - 1 # 0-based path_to_data = os.path.join(self.root, self.base_folder, data_file) with open(path_to_data, "rb") as f: # read whole file in uint8 chunks everything = np.fromfile(f, dtype=np.uint8) images = np.reshape(everything, (-1, 3, 96, 96)) images = np.transpose(images, (0, 1, 3, 2)) return images, labels def _check_integrity(self) -> bool: root = self.root for fentry in self.train_list + self.test_list: filename, md5 = fentry[0], fentry[1] fpath = os.path.join(root, self.base_folder, filename) if not check_integrity(fpath, md5): return False return True def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5) self._check_integrity() def extra_repr(self) -> str: return "Split: {split}".format(**self.__dict__) def __load_folds(self, folds: Optional[int]) -> None: # loads one of the folds if specified if folds is None: return path_to_folds = os.path.join(self.root, self.base_folder, self.folds_list_file) with open(path_to_folds) as f: str_idx = f.read().splitlines()[folds] list_idx = np.fromstring(str_idx, dtype=np.int64, sep=" ") self.data = self.data[list_idx, :, :, :] if self.labels is not None: self.labels = self.labels[list_idx] msadapter/ms_adapter/torchvision/datasets/sun397.py000066400000000000000000000052671436713554100230010ustar00rootroot00000000000000from pathlib import Path from typing import Any, Tuple, Callable, Optional import PIL.Image from .utils import download_and_extract_archive from .vision import VisionDataset class SUN397(VisionDataset): """`The SUN397 Data Set `_. The SUN397 or Scene UNderstanding (SUN) is a dataset for scene recognition consisting of 397 categories with 108'754 images. Args: root (string): Root directory of the dataset. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop``. target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ _DATASET_URL = "http://vision.princeton.edu/projects/2010/SUN/SUN397.tar.gz" _DATASET_MD5 = "8ca2778205c41d23104230ba66911c7a" def __init__( self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self._data_dir = Path(self.root) / "SUN397" if download: self._download() if not self._check_exists(): raise RuntimeError("Dataset not found. You can use download=True to download it") with open(self._data_dir / "ClassName.txt") as f: self.classes = [c[3:].strip() for c in f] self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) self._image_files = list(self._data_dir.rglob("sun_*.jpg")) self._labels = [ self.class_to_idx["/".join(path.relative_to(self._data_dir).parts[1:-1])] for path in self._image_files ] def __len__(self) -> int: return len(self._image_files) def __getitem__(self, idx) -> Tuple[Any, Any]: image_file, label = self._image_files[idx], self._labels[idx] image = PIL.Image.open(image_file).convert("RGB") if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def _check_exists(self) -> bool: return self._data_dir.is_dir() def _download(self) -> None: if self._check_exists(): return download_and_extract_archive(self._DATASET_URL, download_root=self.root, md5=self._DATASET_MD5) msadapter/ms_adapter/torchvision/datasets/svhn.py000066400000000000000000000112361436713554100227000ustar00rootroot00000000000000import os.path from typing import Any, Callable, Optional, Tuple import numpy as np from PIL import Image from .utils import download_url, check_integrity, verify_str_arg from .vision import VisionDataset class SVHN(VisionDataset): """`SVHN `_ Dataset. Note: The SVHN dataset assigns the label `10` to the digit `0`. However, in this Dataset, we assign the label `0` to the digit `0` to be compatible with PyTorch loss functions which expect the class labels to be in the range `[0, C-1]` .. warning:: This class needs `scipy `_ to load data from `.mat` format. Args: root (string): Root directory of the dataset where the data is stored. split (string): One of {'train', 'test', 'extra'}. Accordingly dataset is selected. 'extra' is Extra training set. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ split_list = { "train": [ "http://ufldl.stanford.edu/housenumbers/train_32x32.mat", "train_32x32.mat", "e26dedcc434d2e4c54c9b2d4a06d8373", ], "test": [ "http://ufldl.stanford.edu/housenumbers/test_32x32.mat", "test_32x32.mat", "eb5a983be6a315427106f1b164d9cef3", ], "extra": [ "http://ufldl.stanford.edu/housenumbers/extra_32x32.mat", "extra_32x32.mat", "a93ce644f1a588dc4d68dda5feec44a7", ], } def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) self.split = verify_str_arg(split, "split", tuple(self.split_list.keys())) self.url = self.split_list[split][0] self.filename = self.split_list[split][1] self.file_md5 = self.split_list[split][2] if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") # import here rather than at top of file because this is # an optional dependency for torchvision import scipy.io as sio # reading(loading) mat file as array loaded_mat = sio.loadmat(os.path.join(self.root, self.filename)) self.data = loaded_mat["X"] # loading from the .mat file gives an np array of type np.uint8 # converting to np.int64, so that we have a LongTensor after # the conversion from the numpy array # the squeeze is needed to obtain a 1D tensor self.labels = loaded_mat["y"].astype(np.int64).squeeze() # the svhn dataset assigns the class label "10" to the digit 0 # this makes it inconsistent with several loss functions # which expect the class labels to be in the range [0, C-1] np.place(self.labels, self.labels == 10, 0) self.data = np.transpose(self.data, (3, 2, 0, 1)) def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target class. """ img, target = self.data[index], int(self.labels[index]) # doing this so that it is consistent with all other datasets # to return a PIL Image img = Image.fromarray(np.transpose(img, (1, 2, 0))) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.data) def _check_integrity(self) -> bool: root = self.root md5 = self.split_list[self.split][2] fpath = os.path.join(root, self.filename) return check_integrity(fpath, md5) def download(self) -> None: md5 = self.split_list[self.split][2] download_url(self.url, self.root, self.filename, md5) def extra_repr(self) -> str: return "Split: {split}".format(**self.__dict__) msadapter/ms_adapter/torchvision/datasets/ucf101.py000066400000000000000000000124511436713554100227210ustar00rootroot00000000000000import os from typing import Any, Dict, List, Tuple, Optional, Callable from .folder import find_classes, make_dataset from .video_utils import VideoClips from .vision import VisionDataset class UCF101(VisionDataset): """ `UCF101 `_ dataset. UCF101 is an action recognition video dataset. This dataset consider every video as a collection of video clips of fixed size, specified by ``frames_per_clip``, where the step in frames between each clip is given by ``step_between_clips``. The dataset itself can be downloaded from the dataset website; annotations that ``annotation_path`` should be pointing to can be downloaded from `here `. To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5`` and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two elements will come from video 1, and the next three elements from video 2. Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all frames in a video might be present. Internally, it uses a VideoClips object to handle clip creation. Args: root (string): Root directory of the UCF101 Dataset. annotation_path (str): path to the folder containing the split files; see docstring above for download instructions of these files frames_per_clip (int): number of frames in a clip. step_between_clips (int, optional): number of frames between each clip. fold (int, optional): which fold to use. Should be between 1 and 3. train (bool, optional): if ``True``, creates a dataset from the train split, otherwise from the ``test`` split. transform (callable, optional): A function/transform that takes in a TxHxWxC video and returns a transformed version. output_format (str, optional): The format of the output video tensors (before transforms). Can be either "THWC" (default) or "TCHW". Returns: tuple: A 3-tuple with the following entries: - video (Tensor[T, H, W, C] or Tensor[T, C, H, W]): The `T` video frames - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels and `L` is the number of points - label (int): class of the video clip """ def __init__( self, root: str, annotation_path: str, frames_per_clip: int, step_between_clips: int = 1, frame_rate: Optional[int] = None, fold: int = 1, train: bool = True, transform: Optional[Callable] = None, _precomputed_metadata: Optional[Dict[str, Any]] = None, num_workers: int = 1, _video_width: int = 0, _video_height: int = 0, _video_min_dimension: int = 0, _audio_samples: int = 0, output_format: str = "THWC", ) -> None: super().__init__(root) if not 1 <= fold <= 3: raise ValueError(f"fold should be between 1 and 3, got {fold}") extensions = ("avi",) self.fold = fold self.train = train self.classes, class_to_idx = find_classes(self.root) self.samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file=None) video_list = [x[0] for x in self.samples] video_clips = VideoClips( video_list, frames_per_clip, step_between_clips, frame_rate, _precomputed_metadata, num_workers=num_workers, _video_width=_video_width, _video_height=_video_height, _video_min_dimension=_video_min_dimension, _audio_samples=_audio_samples, output_format=output_format, ) # we bookkeep the full version of video clips because we want to be able # to return the meta data of full version rather than the subset version of # video clips self.full_video_clips = video_clips self.indices = self._select_fold(video_list, annotation_path, fold, train) self.video_clips = video_clips.subset(self.indices) self.transform = transform @property def metadata(self) -> Dict[str, Any]: return self.full_video_clips.metadata def _select_fold(self, video_list: List[str], annotation_path: str, fold: int, train: bool) -> List[int]: name = "train" if train else "test" name = f"{name}list{fold:02d}.txt" f = os.path.join(annotation_path, name) selected_files = set() with open(f) as fid: data = fid.readlines() data = [x.strip().split(" ")[0] for x in data] data = [os.path.join(self.root, *x.split("/")) for x in data] selected_files.update(data) indices = [i for i in range(len(video_list)) if video_list[i] in selected_files] return indices def __len__(self) -> int: return self.video_clips.num_clips() def __getitem__(self, idx: int): video, audio, info, video_idx = self.video_clips.get_clip(idx) label = self.samples[self.indices[video_idx]][1] if self.transform is not None: video = self.transform(video) return video, audio, label msadapter/ms_adapter/torchvision/datasets/usps.py000066400000000000000000000065601436713554100227200ustar00rootroot00000000000000import os from typing import Any, Callable, Optional, Tuple import numpy as np from PIL import Image from .utils import download_url from .vision import VisionDataset class USPS(VisionDataset): """`USPS `_ Dataset. The data-format is : [label [index:value ]*256 \\n] * num_lines, where ``label`` lies in ``[1, 10]``. The value for each pixel lies in ``[-1, 1]``. Here we transform the ``label`` into ``[0, 9]`` and make pixel values in ``[0, 255]``. Args: root (string): Root directory of dataset to store``USPS`` data files. train (bool, optional): If True, creates dataset from ``usps.bz2``, otherwise from ``usps.t.bz2``. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ split_list = { "train": [ "https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/usps.bz2", "usps.bz2", "ec16c51db3855ca6c91edd34d0e9b197", ], "test": [ "https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/usps.t.bz2", "usps.t.bz2", "8ea070ee2aca1ac39742fdd1ef5ed118", ], } def __init__( self, root: str, train: bool = True, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__(root, transform=transform, target_transform=target_transform) split = "train" if train else "test" url, filename, checksum = self.split_list[split] full_path = os.path.join(self.root, filename) if download and not os.path.exists(full_path): download_url(url, self.root, filename, md5=checksum) import bz2 with bz2.open(full_path) as fp: raw_data = [line.decode().split() for line in fp.readlines()] tmp_list = [[x.split(":")[-1] for x in data[1:]] for data in raw_data] imgs = np.asarray(tmp_list, dtype=np.float32).reshape((-1, 16, 16)) imgs = ((imgs + 1) / 2 * 255).astype(dtype=np.uint8) targets = [int(d[0]) - 1 for d in raw_data] self.data = imgs self.targets = targets def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is index of the target class. """ img, target = self.data[index], int(self.targets[index]) # doing this so that it is consistent with all other datasets # to return a PIL Image img = Image.fromarray(img, mode="L") if self.transform is not None: img = self.transform(img) if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.data) msadapter/ms_adapter/torchvision/datasets/utils.py000066400000000000000000000415541436713554100230700ustar00rootroot00000000000000import bz2 import contextlib import gzip import hashlib import itertools import lzma import os import os.path import pathlib import re import sys import tarfile import urllib import urllib.error import urllib.request import warnings import zipfile from typing import Any, Callable, List, Iterable, Optional, TypeVar, Dict, IO, Tuple, Iterator from urllib.parse import urlparse import requests from ms_adapter.pytorch._six import string_classes # from torch.utils.model_zoo import tqdm from tqdm import tqdm from .._internally_replaced_utils import ( _download_file_from_remote_location, _is_remote_location_available, ) USER_AGENT = "pytorch/vision" def _save_response_content( content: Iterator[bytes], destination: str, length: Optional[int] = None, ) -> None: with open(destination, "wb") as fh, tqdm(total=length) as pbar: for chunk in content: # filter out keep-alive new chunks if not chunk: continue fh.write(chunk) pbar.update(len(chunk)) def _urlretrieve(url: str, filename: str, chunk_size: int = 1024 * 32) -> None: with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": USER_AGENT})) as response: _save_response_content(iter(lambda: response.read(chunk_size), b""), filename, length=response.length) def gen_bar_updater() -> Callable[[int, int, int], None]: warnings.warn("The function `gen_bar_update` is deprecated since 0.13 and will be removed in 0.15.") pbar = tqdm(total=None) def bar_update(count, block_size, total_size): if pbar.total is None and total_size: pbar.total = total_size progress_bytes = count * block_size pbar.update(progress_bytes - pbar.n) return bar_update def calculate_md5(fpath: str, chunk_size: int = 1024 * 1024) -> str: # Setting the `usedforsecurity` flag does not change anything about the functionality, but indicates that we are # not using the MD5 checksum for cryptography. This enables its usage in restricted environments like FIPS. Without # it torchvision.datasets is unusable in these environments since we perform a MD5 check everywhere. md5 = hashlib.md5(**dict(usedforsecurity=False) if sys.version_info >= (3, 9) else dict()) with open(fpath, "rb") as f: for chunk in iter(lambda: f.read(chunk_size), b""): md5.update(chunk) return md5.hexdigest() def check_md5(fpath: str, md5: str, **kwargs: Any) -> bool: return md5 == calculate_md5(fpath, **kwargs) def check_integrity(fpath: str, md5: Optional[str] = None) -> bool: if not os.path.isfile(fpath): return False if md5 is None: return True return check_md5(fpath, md5) def _get_redirect_url(url: str, max_hops: int = 3) -> str: initial_url = url headers = {"Method": "HEAD", "User-Agent": USER_AGENT} for _ in range(max_hops + 1): with urllib.request.urlopen(urllib.request.Request(url, headers=headers)) as response: if response.url == url or response.url is None: return url url = response.url else: raise RecursionError( f"Request to {initial_url} exceeded {max_hops} redirects. The last redirect points to {url}." ) def _get_google_drive_file_id(url: str) -> Optional[str]: parts = urlparse(url) if re.match(r"(drive|docs)[.]google[.]com", parts.netloc) is None: return None match = re.match(r"/file/d/(?P[^/]*)", parts.path) if match is None: return None return match.group("id") def download_url( url: str, root: str, filename: Optional[str] = None, md5: Optional[str] = None, max_redirect_hops: int = 3 ) -> None: """Download a file from a url and place it in root. Args: url (str): URL to download file from root (str): Directory to place downloaded file in filename (str, optional): Name to save the file under. If None, use the basename of the URL md5 (str, optional): MD5 checksum of the download. If None, do not check max_redirect_hops (int, optional): Maximum number of redirect hops allowed """ root = os.path.expanduser(root) if not filename: filename = os.path.basename(url) fpath = os.path.join(root, filename) os.makedirs(root, exist_ok=True) # check if file is already present locally if check_integrity(fpath, md5): print("Using downloaded and verified file: " + fpath) return if _is_remote_location_available(): _download_file_from_remote_location(fpath, url) else: # expand redirect chain if needed url = _get_redirect_url(url, max_hops=max_redirect_hops) # check if file is located on Google Drive file_id = _get_google_drive_file_id(url) if file_id is not None: return download_file_from_google_drive(file_id, root, filename, md5) # download the file try: print("Downloading " + url + " to " + fpath) _urlretrieve(url, fpath) except (urllib.error.URLError, OSError) as e: # type: ignore[attr-defined] if url[:5] == "https": url = url.replace("https:", "http:") print("Failed download. Trying https -> http instead. Downloading " + url + " to " + fpath) _urlretrieve(url, fpath) else: raise e # check integrity of downloaded file if not check_integrity(fpath, md5): raise RuntimeError("File not found or corrupted.") def list_dir(root: str, prefix: bool = False) -> List[str]: """List all directories at a given root Args: root (str): Path to directory whose folders need to be listed prefix (bool, optional): If true, prepends the path to each result, otherwise only returns the name of the directories found """ root = os.path.expanduser(root) directories = [p for p in os.listdir(root) if os.path.isdir(os.path.join(root, p))] if prefix is True: directories = [os.path.join(root, d) for d in directories] return directories def list_files(root: str, suffix: str, prefix: bool = False) -> List[str]: """List all files ending with a suffix at a given root Args: root (str): Path to directory whose folders need to be listed suffix (str or tuple): Suffix of the files to match, e.g. '.png' or ('.jpg', '.png'). It uses the Python "str.endswith" method and is passed directly prefix (bool, optional): If true, prepends the path to each result, otherwise only returns the name of the files found """ root = os.path.expanduser(root) files = [p for p in os.listdir(root) if os.path.isfile(os.path.join(root, p)) and p.endswith(suffix)] if prefix is True: files = [os.path.join(root, d) for d in files] return files def _extract_gdrive_api_response(response, chunk_size: int = 32 * 1024) -> Tuple[bytes, Iterator[bytes]]: content = response.iter_content(chunk_size) first_chunk = None # filter out keep-alive new chunks while not first_chunk: first_chunk = next(content) content = itertools.chain([first_chunk], content) try: match = re.search("Google Drive - (?P<api_response>.+?)", first_chunk.decode()) api_response = match["api_response"] if match is not None else None except UnicodeDecodeError: api_response = None return api_response, content def download_file_from_google_drive(file_id: str, root: str, filename: Optional[str] = None, md5: Optional[str] = None): """Download a Google Drive file from and place it in root. Args: file_id (str): id of file to be downloaded root (str): Directory to place downloaded file in filename (str, optional): Name to save the file under. If None, use the id of the file. md5 (str, optional): MD5 checksum of the download. If None, do not check """ root = os.path.expanduser(root) if not filename: filename = file_id fpath = os.path.join(root, filename) os.makedirs(root, exist_ok=True) if check_integrity(fpath, md5): print(f"Using downloaded {'and verified ' if md5 else ''}file: {fpath}") return url = "https://drive.google.com/uc" params = dict(id=file_id, export="download") with requests.Session() as session: response = session.get(url, params=params, stream=True) for key, value in response.cookies.items(): if key.startswith("download_warning"): token = value break else: api_response, content = _extract_gdrive_api_response(response) token = "t" if api_response == "Virus scan warning" else None if token is not None: response = session.get(url, params=dict(params, confirm=token), stream=True) api_response, content = _extract_gdrive_api_response(response) if api_response == "Quota exceeded": raise RuntimeError( f"The daily quota of the file {filename} is exceeded and it " f"can't be downloaded. This is a limitation of Google Drive " f"and can only be overcome by trying again later." ) _save_response_content(content, fpath) # In case we deal with an unhandled GDrive API response, the file should be smaller than 10kB and contain only text if os.stat(fpath).st_size < 10 * 1024: with contextlib.suppress(UnicodeDecodeError), open(fpath) as fh: text = fh.read() # Regular expression to detect HTML. Copied from https://stackoverflow.com/a/70585604 if re.search(r"]*\s*>|(&(?:[\w\d]+|#\d+|#x[a-f\d]+);)", text): warnings.warn( f"We detected some HTML elements in the downloaded file. " f"This most likely means that the download triggered an unhandled API response by GDrive. " f"Please report this to torchvision at https://github.com/pytorch/vision/issues including " f"the response:\n\n{text}" ) if md5 and not check_md5(fpath, md5): raise RuntimeError( f"The MD5 checksum of the download file {fpath} does not match the one on record." f"Please delete the file and try again. " f"If the issue persists, please report this to torchvision at https://github.com/pytorch/vision/issues." ) def _extract_tar(from_path: str, to_path: str, compression: Optional[str]) -> None: with tarfile.open(from_path, f"r:{compression[1:]}" if compression else "r") as tar: tar.extractall(to_path) _ZIP_COMPRESSION_MAP: Dict[str, int] = { ".bz2": zipfile.ZIP_BZIP2, ".xz": zipfile.ZIP_LZMA, } def _extract_zip(from_path: str, to_path: str, compression: Optional[str]) -> None: with zipfile.ZipFile( from_path, "r", compression=_ZIP_COMPRESSION_MAP[compression] if compression else zipfile.ZIP_STORED ) as zip: zip.extractall(to_path) _ARCHIVE_EXTRACTORS: Dict[str, Callable[[str, str, Optional[str]], None]] = { ".tar": _extract_tar, ".zip": _extract_zip, } _COMPRESSED_FILE_OPENERS: Dict[str, Callable[..., IO]] = { ".bz2": bz2.open, ".gz": gzip.open, ".xz": lzma.open, } _FILE_TYPE_ALIASES: Dict[str, Tuple[Optional[str], Optional[str]]] = { ".tbz": (".tar", ".bz2"), ".tbz2": (".tar", ".bz2"), ".tgz": (".tar", ".gz"), } def _detect_file_type(file: str) -> Tuple[str, Optional[str], Optional[str]]: """Detect the archive type and/or compression of a file. Args: file (str): the filename Returns: (tuple): tuple of suffix, archive type, and compression Raises: RuntimeError: if file has no suffix or suffix is not supported """ suffixes = pathlib.Path(file).suffixes if not suffixes: raise RuntimeError( f"File '{file}' has no suffixes that could be used to detect the archive type and compression." ) suffix = suffixes[-1] # check if the suffix is a known alias if suffix in _FILE_TYPE_ALIASES: return (suffix, *_FILE_TYPE_ALIASES[suffix]) # check if the suffix is an archive type if suffix in _ARCHIVE_EXTRACTORS: return suffix, suffix, None # check if the suffix is a compression if suffix in _COMPRESSED_FILE_OPENERS: # check for suffix hierarchy if len(suffixes) > 1: suffix2 = suffixes[-2] # check if the suffix2 is an archive type if suffix2 in _ARCHIVE_EXTRACTORS: return suffix2 + suffix, suffix2, suffix return suffix, None, suffix valid_suffixes = sorted(set(_FILE_TYPE_ALIASES) | set(_ARCHIVE_EXTRACTORS) | set(_COMPRESSED_FILE_OPENERS)) raise RuntimeError(f"Unknown compression or archive type: '{suffix}'.\nKnown suffixes are: '{valid_suffixes}'.") def _decompress(from_path: str, to_path: Optional[str] = None, remove_finished: bool = False) -> str: r"""Decompress a file. The compression is automatically detected from the file name. Args: from_path (str): Path to the file to be decompressed. to_path (str): Path to the decompressed file. If omitted, ``from_path`` without compression extension is used. remove_finished (bool): If ``True``, remove the file after the extraction. Returns: (str): Path to the decompressed file. """ suffix, archive_type, compression = _detect_file_type(from_path) if not compression: raise RuntimeError(f"Couldn't detect a compression from suffix {suffix}.") if to_path is None: to_path = from_path.replace(suffix, archive_type if archive_type is not None else "") # We don't need to check for a missing key here, since this was already done in _detect_file_type() compressed_file_opener = _COMPRESSED_FILE_OPENERS[compression] with compressed_file_opener(from_path, "rb") as rfh, open(to_path, "wb") as wfh: wfh.write(rfh.read()) if remove_finished: os.remove(from_path) return to_path def extract_archive(from_path: str, to_path: Optional[str] = None, remove_finished: bool = False) -> str: """Extract an archive. The archive type and a possible compression is automatically detected from the file name. If the file is compressed but not an archive the call is dispatched to :func:`decompress`. Args: from_path (str): Path to the file to be extracted. to_path (str): Path to the directory the file will be extracted to. If omitted, the directory of the file is used. remove_finished (bool): If ``True``, remove the file after the extraction. Returns: (str): Path to the directory the file was extracted to. """ if to_path is None: to_path = os.path.dirname(from_path) suffix, archive_type, compression = _detect_file_type(from_path) if not archive_type: return _decompress( from_path, os.path.join(to_path, os.path.basename(from_path).replace(suffix, "")), remove_finished=remove_finished, ) # We don't need to check for a missing key here, since this was already done in _detect_file_type() extractor = _ARCHIVE_EXTRACTORS[archive_type] extractor(from_path, to_path, compression) if remove_finished: os.remove(from_path) return to_path def download_and_extract_archive( url: str, download_root: str, extract_root: Optional[str] = None, filename: Optional[str] = None, md5: Optional[str] = None, remove_finished: bool = False, ) -> None: download_root = os.path.expanduser(download_root) if extract_root is None: extract_root = download_root if not filename: filename = os.path.basename(url) download_url(url, download_root, filename, md5) archive = os.path.join(download_root, filename) print(f"Extracting {archive} to {extract_root}") extract_archive(archive, extract_root, remove_finished) def iterable_to_str(iterable: Iterable) -> str: return "'" + "', '".join([str(item) for item in iterable]) + "'" T = TypeVar("T", str, bytes) def verify_str_arg( value: T, arg: Optional[str] = None, valid_values: Iterable[T] = None, custom_msg: Optional[str] = None, ) -> T: if not isinstance(value, string_classes): if arg is None: msg = "Expected type str, but got type {type}." else: msg = "Expected type str for argument {arg}, but got type {type}." msg = msg.format(type=type(value), arg=arg) raise ValueError(msg) if valid_values is None: return value if value not in valid_values: if custom_msg is not None: msg = custom_msg else: msg = "Unknown value '{value}' for argument {arg}. Valid values are {{{valid_values}}}." msg = msg.format(value=value, arg=arg, valid_values=iterable_to_str(valid_values)) raise ValueError(msg) return value msadapter/ms_adapter/torchvision/datasets/video_utils.py000066400000000000000000000414671436713554100242610ustar00rootroot00000000000000import bisect import math import warnings from fractions import Fraction from typing import Any, Dict, List, Optional, Callable, Union, Tuple, TypeVar, cast import ms_adapter.pytorch as torch from ms_adapter.torchvision.io import ( # _probe_video_from_file, # _read_video_from_file, read_video, read_video_timestamps, ) from .utils import tqdm T = TypeVar("T") def pts_convert(pts: int, timebase_from: Fraction, timebase_to: Fraction, round_func: Callable = math.floor) -> int: """convert pts between different time bases Args: pts: presentation timestamp, float timebase_from: original timebase. Fraction timebase_to: new timebase. Fraction round_func: rounding function. """ new_pts = Fraction(pts, 1) * timebase_from / timebase_to return round_func(new_pts) def unfold(tensor: torch.Tensor, size: int, step: int, dilation: int = 1) -> torch.Tensor: """ similar to tensor.unfold, but with the dilation and specialized for 1d tensors Returns all consecutive windows of `size` elements, with `step` between windows. The distance between each element in a window is given by `dilation`. """ if tensor.dim() != 1: raise ValueError(f"tensor should have 1 dimension instead of {tensor.dim()}") o_stride = tensor.stride(0) numel = tensor.numel() new_stride = (step * o_stride, dilation * o_stride) new_size = ((numel - (dilation * (size - 1) + 1)) // step + 1, size) if new_size[0] < 1: new_size = (0, size) return torch.as_strided(tensor, new_size, new_stride) class _VideoTimestampsDataset: """ Dataset used to parallelize the reading of the timestamps of a list of videos, given their paths in the filesystem. Used in VideoClips and defined at top level so it can be pickled when forking. """ def __init__(self, video_paths: List[str]) -> None: self.video_paths = video_paths def __len__(self) -> int: return len(self.video_paths) def __getitem__(self, idx: int) -> Tuple[List[int], Optional[float]]: return read_video_timestamps(self.video_paths[idx]) def _collate_fn(x: T) -> T: """ Dummy collate function to be used with _VideoTimestampsDataset """ return x class VideoClips: """ Given a list of video files, computes all consecutive subvideos of size `clip_length_in_frames`, where the distance between each subvideo in the same video is defined by `frames_between_clips`. If `frame_rate` is specified, it will also resample all the videos to have the same frame rate, and the clips will refer to this frame rate. Creating this instance the first time is time-consuming, as it needs to decode all the videos in `video_paths`. It is recommended that you cache the results after instantiation of the class. Recreating the clips for different clip lengths is fast, and can be done with the `compute_clips` method. Args: video_paths (List[str]): paths to the video files clip_length_in_frames (int): size of a clip in number of frames frames_between_clips (int): step (in frames) between each clip frame_rate (int, optional): if specified, it will resample the video so that it has `frame_rate`, and then the clips will be defined on the resampled video num_workers (int): how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: 0) output_format (str): The format of the output video tensors. Can be either "THWC" (default) or "TCHW". """ def __init__( self, video_paths: List[str], clip_length_in_frames: int = 16, frames_between_clips: int = 1, frame_rate: Optional[int] = None, _precomputed_metadata: Optional[Dict[str, Any]] = None, num_workers: int = 0, _video_width: int = 0, _video_height: int = 0, _video_min_dimension: int = 0, _video_max_dimension: int = 0, _audio_samples: int = 0, _audio_channels: int = 0, output_format: str = "THWC", ) -> None: self.video_paths = video_paths self.num_workers = num_workers # these options are not valid for pyav backend self._video_width = _video_width self._video_height = _video_height self._video_min_dimension = _video_min_dimension self._video_max_dimension = _video_max_dimension self._audio_samples = _audio_samples self._audio_channels = _audio_channels self.output_format = output_format.upper() if self.output_format not in ("THWC", "TCHW"): raise ValueError(f"output_format should be either 'THWC' or 'TCHW', got {output_format}.") if _precomputed_metadata is None: self._compute_frame_pts() else: self._init_from_metadata(_precomputed_metadata) self.compute_clips(clip_length_in_frames, frames_between_clips, frame_rate) def _compute_frame_pts(self) -> None: self.video_pts = [] self.video_fps = [] # strategy: use a DataLoader to parallelize read_video_timestamps # so need to create a dummy dataset first import ms_adapter.pytorch.utils.data dl: ms_adapter.pytorch.utils.data.DataLoader = ms_adapter.pytorch.utils.data.DataLoader( _VideoTimestampsDataset(self.video_paths), # type: ignore[arg-type] batch_size=16, num_workers=self.num_workers, collate_fn=_collate_fn, ) with tqdm(total=len(dl)) as pbar: for batch in dl: pbar.update(1) clips, fps = list(zip(*batch)) # we need to specify dtype=torch.long because for empty list, # torch.as_tensor will use torch.float as default dtype. This # happens when decoding fails and no pts is returned in the list. clips = [torch.as_tensor(c, dtype=torch.long) for c in clips] self.video_pts.extend(clips) self.video_fps.extend(fps) def _init_from_metadata(self, metadata: Dict[str, Any]) -> None: self.video_paths = metadata["video_paths"] assert len(self.video_paths) == len(metadata["video_pts"]) self.video_pts = metadata["video_pts"] assert len(self.video_paths) == len(metadata["video_fps"]) self.video_fps = metadata["video_fps"] @property def metadata(self) -> Dict[str, Any]: _metadata = { "video_paths": self.video_paths, "video_pts": self.video_pts, "video_fps": self.video_fps, } return _metadata def subset(self, indices: List[int]) -> "VideoClips": video_paths = [self.video_paths[i] for i in indices] video_pts = [self.video_pts[i] for i in indices] video_fps = [self.video_fps[i] for i in indices] metadata = { "video_paths": video_paths, "video_pts": video_pts, "video_fps": video_fps, } return type(self)( video_paths, self.num_frames, self.step, self.frame_rate, _precomputed_metadata=metadata, num_workers=self.num_workers, _video_width=self._video_width, _video_height=self._video_height, _video_min_dimension=self._video_min_dimension, _video_max_dimension=self._video_max_dimension, _audio_samples=self._audio_samples, _audio_channels=self._audio_channels, ) @staticmethod def compute_clips_for_video( video_pts: torch.Tensor, num_frames: int, step: int, fps: int, frame_rate: Optional[int] = None ) -> Tuple[torch.Tensor, Union[List[slice], torch.Tensor]]: if fps is None: # if for some reason the video doesn't have fps (because doesn't have a video stream) # set the fps to 1. The value doesn't matter, because video_pts is empty anyway fps = 1 if frame_rate is None: frame_rate = fps total_frames = len(video_pts) * (float(frame_rate) / fps) _idxs = VideoClips._resample_video_idx(int(math.floor(total_frames)), fps, frame_rate) video_pts = video_pts[_idxs] clips = unfold(video_pts, num_frames, step) if not clips.numel(): warnings.warn( "There aren't enough frames in the current video to get a clip for the given clip length and " "frames between clips. The video (and potentially others) will be skipped." ) idxs: Union[List[slice], torch.Tensor] if isinstance(_idxs, slice): idxs = [_idxs] * len(clips) else: idxs = unfold(_idxs, num_frames, step) return clips, idxs def compute_clips(self, num_frames: int, step: int, frame_rate: Optional[int] = None) -> None: """ Compute all consecutive sequences of clips from video_pts. Always returns clips of size `num_frames`, meaning that the last few frames in a video can potentially be dropped. Args: num_frames (int): number of frames for the clip step (int): distance between two clips frame_rate (int, optional): The frame rate """ self.num_frames = num_frames self.step = step self.frame_rate = frame_rate self.clips = [] self.resampling_idxs = [] for video_pts, fps in zip(self.video_pts, self.video_fps): clips, idxs = self.compute_clips_for_video(video_pts, num_frames, step, fps, frame_rate) self.clips.append(clips) self.resampling_idxs.append(idxs) clip_lengths = torch.as_tensor([len(v) for v in self.clips]) self.cumulative_sizes = clip_lengths.cumsum(0).asnumpy().tolist() def __len__(self) -> int: return self.num_clips() def num_videos(self) -> int: return len(self.video_paths) def num_clips(self) -> int: """ Number of subclips that are available in the video list. """ return self.cumulative_sizes[-1] def get_clip_location(self, idx: int) -> Tuple[int, int]: """ Converts a flattened representation of the indices into a video_idx, clip_idx representation. """ video_idx = bisect.bisect_right(self.cumulative_sizes, idx) if video_idx == 0: clip_idx = idx else: clip_idx = idx - self.cumulative_sizes[video_idx - 1] return video_idx, clip_idx @staticmethod def _resample_video_idx(num_frames: int, original_fps: int, new_fps: int) -> Union[slice, torch.Tensor]: step = float(original_fps) / new_fps if step.is_integer(): # optimization: if step is integer, don't need to perform # advanced indexing step = int(step) return slice(None, None, step) idxs = torch.arange(start=0, end = num_frames, dtype=torch.float32) * step idxs = idxs.floor().to(torch.int64) return idxs def get_clip(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any], int]: """ Gets a subclip from a list of videos. Args: idx (int): index of the subclip. Must be between 0 and num_clips(). Returns: video (Tensor) audio (Tensor) info (Dict) video_idx (int): index of the video in `video_paths` """ if idx >= self.num_clips(): raise IndexError(f"Index {idx} out of range ({self.num_clips()} number of clips)") video_idx, clip_idx = self.get_clip_location(idx) video_path = self.video_paths[video_idx] clip_pts = self.clips[video_idx][clip_idx] from ms_adapter.torchvision import get_video_backend backend = get_video_backend() if backend == "pyav": # check for invalid options if self._video_width != 0: raise ValueError("pyav backend doesn't support _video_width != 0") if self._video_height != 0: raise ValueError("pyav backend doesn't support _video_height != 0") if self._video_min_dimension != 0: raise ValueError("pyav backend doesn't support _video_min_dimension != 0") if self._video_max_dimension != 0: raise ValueError("pyav backend doesn't support _video_max_dimension != 0") if self._audio_samples != 0: raise ValueError("pyav backend doesn't support _audio_samples != 0") if backend == "pyav": start_pts = clip_pts[0].item() end_pts = clip_pts[-1].item() video, audio, info = read_video(video_path, start_pts, end_pts) # else: #TODO just support pyav # _info = _probe_video_from_file(video_path) # video_fps = _info.video_fps # audio_fps = None # # video_start_pts = cast(int, clip_pts[0].item()) # video_end_pts = cast(int, clip_pts[-1].item()) # # audio_start_pts, audio_end_pts = 0, -1 # audio_timebase = Fraction(0, 1) # video_timebase = Fraction(_info.video_timebase.numerator, _info.video_timebase.denominator) # if _info.has_audio: # audio_timebase = Fraction(_info.audio_timebase.numerator, _info.audio_timebase.denominator) # audio_start_pts = pts_convert(video_start_pts, video_timebase, audio_timebase, math.floor) # audio_end_pts = pts_convert(video_end_pts, video_timebase, audio_timebase, math.ceil) # audio_fps = _info.audio_sample_rate # video, audio, _ = _read_video_from_file( # video_path, # video_width=self._video_width, # video_height=self._video_height, # video_min_dimension=self._video_min_dimension, # video_max_dimension=self._video_max_dimension, # video_pts_range=(video_start_pts, video_end_pts), # video_timebase=video_timebase, # audio_samples=self._audio_samples, # audio_channels=self._audio_channels, # audio_pts_range=(audio_start_pts, audio_end_pts), # audio_timebase=audio_timebase, # ) # # info = {"video_fps": video_fps} # if audio_fps is not None: # info["audio_fps"] = audio_fps if self.frame_rate is not None: resampling_idx = self.resampling_idxs[video_idx][clip_idx] if isinstance(resampling_idx, torch.Tensor): resampling_idx = resampling_idx - resampling_idx[0] video = video[resampling_idx] info["video_fps"] = self.frame_rate assert len(video) == self.num_frames, f"{video.shape} x {self.num_frames}" if self.output_format == "TCHW": # [T,H,W,C] --> [T,C,H,W] video = video.permute(0, 3, 1, 2) return video, audio, info, video_idx def __getstate__(self) -> Dict[str, Any]: video_pts_sizes = [len(v) for v in self.video_pts] # To be back-compatible, we convert data to dtype torch.long as needed # because for empty list, in legacy implementation, torch.as_tensor will # use torch.float as default dtype. This happens when decoding fails and # no pts is returned in the list. video_pts = [x.to(torch.int64) for x in self.video_pts] # video_pts can be an empty list if no frames have been decoded if video_pts: video_pts = torch.cat(video_pts) # type: ignore[assignment] # avoid bug in https://github.com/pytorch/pytorch/issues/32351 # TODO: Revert it once the bug is fixed. video_pts = video_pts.numpy() # type: ignore[attr-defined] # make a copy of the fields of self d = self.__dict__.copy() d["video_pts_sizes"] = video_pts_sizes d["video_pts"] = video_pts # delete the following attributes to reduce the size of dictionary. They # will be re-computed in "__setstate__()" del d["clips"] del d["resampling_idxs"] del d["cumulative_sizes"] # for backwards-compatibility d["_version"] = 2 return d def __setstate__(self, d: Dict[str, Any]) -> None: # for backwards-compatibility if "_version" not in d: self.__dict__ = d return video_pts = torch.as_tensor(d["video_pts"], dtype=torch.int64) video_pts = torch.split(video_pts, d["video_pts_sizes"], dim=0) # don't need this info anymore del d["video_pts_sizes"] d["video_pts"] = video_pts self.__dict__ = d # recompute attributes "clips", "resampling_idxs" and other derivative ones self.compute_clips(self.num_frames, self.step, self.frame_rate) msadapter/ms_adapter/torchvision/datasets/vision.py000066400000000000000000000101761436713554100232330ustar00rootroot00000000000000import os from typing import Any, Callable, List, Optional, Tuple import ms_adapter.pytorch.utils.data as data from ms_adapter.pytorch._six import string_classes # from ..utils import _log_api_usage_once class VisionDataset(data.Dataset): """ Base Class For making datasets which are compatible with torchvision. It is necessary to override the ``__getitem__`` and ``__len__`` method. Args: root (string): Root directory of dataset. transforms (callable, optional): A function/transforms that takes in an image and a label and returns the transformed versions of both. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. .. note:: :attr:`transforms` and the combination of :attr:`transform` and :attr:`target_transform` are mutually exclusive. """ def __init__( self, root: str, transforms: Optional[Callable] = None, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, ) -> None: # _log_api_usage_once(self) if isinstance(root, string_classes): root = os.path.expanduser(root) self.root = root self._repr_indent = 4 has_transforms = transforms is not None has_separate_transform = transform is not None or target_transform is not None if has_transforms and has_separate_transform: raise ValueError("Only transforms or transform/target_transform can be passed as argument") # for backwards-compatibility self.transform = transform self.target_transform = target_transform if has_separate_transform: transforms = StandardTransform(transform, target_transform) self.transforms = transforms def __getitem__(self, index: int) -> Any: """ Args: index (int): Index Returns: (Any): Sample and meta data, optionally transformed by the respective transforms. """ raise NotImplementedError def __len__(self) -> int: raise NotImplementedError def __repr__(self) -> str: head = "Dataset " + self.__class__.__name__ body = [f"Number of datapoints: {self.__len__()}"] if self.root is not None: body.append(f"Root location: {self.root}") body += self.extra_repr().splitlines() if hasattr(self, "transforms") and self.transforms is not None: body += [repr(self.transforms)] lines = [head] + [" " * self._repr_indent + line for line in body] return "\n".join(lines) def _format_transform_repr(self, transform: Callable, head: str) -> List[str]: lines = transform.__repr__().splitlines() return [f"{head}{lines[0]}"] + ["{}{}".format(" " * len(head), line) for line in lines[1:]] def extra_repr(self) -> str: return "" class StandardTransform: def __init__(self, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None) -> None: self.transform = transform self.target_transform = target_transform def __call__(self, input: Any, target: Any) -> Tuple[Any, Any]: if self.transform is not None: input = self.transform(input) if self.target_transform is not None: target = self.target_transform(target) return input, target def _format_transform_repr(self, transform: Callable, head: str) -> List[str]: lines = transform.__repr__().splitlines() return [f"{head}{lines[0]}"] + ["{}{}".format(" " * len(head), line) for line in lines[1:]] def __repr__(self) -> str: body = [self.__class__.__name__] if self.transform is not None: body += self._format_transform_repr(self.transform, "Transform: ") if self.target_transform is not None: body += self._format_transform_repr(self.target_transform, "Target transform: ") return "\n".join(body) msadapter/ms_adapter/torchvision/datasets/voc.py000066400000000000000000000221771436713554100225170ustar00rootroot00000000000000import collections import os from xml.etree.ElementTree import Element as ET_Element from .vision import VisionDataset try: from defusedxml.ElementTree import parse as ET_parse except ImportError: from xml.etree.ElementTree import parse as ET_parse import warnings from typing import Any, Callable, Dict, Optional, Tuple, List from PIL import Image from .utils import download_and_extract_archive, verify_str_arg DATASET_YEAR_DICT = { "2012": { "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar", "filename": "VOCtrainval_11-May-2012.tar", "md5": "6cd6e144f989b92b3379bac3b3de84fd", "base_dir": os.path.join("VOCdevkit", "VOC2012"), }, "2011": { "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2011/VOCtrainval_25-May-2011.tar", "filename": "VOCtrainval_25-May-2011.tar", "md5": "6c3384ef61512963050cb5d687e5bf1e", "base_dir": os.path.join("TrainVal", "VOCdevkit", "VOC2011"), }, "2010": { "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2010/VOCtrainval_03-May-2010.tar", "filename": "VOCtrainval_03-May-2010.tar", "md5": "da459979d0c395079b5c75ee67908abb", "base_dir": os.path.join("VOCdevkit", "VOC2010"), }, "2009": { "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2009/VOCtrainval_11-May-2009.tar", "filename": "VOCtrainval_11-May-2009.tar", "md5": "a3e00b113cfcfebf17e343f59da3caa1", "base_dir": os.path.join("VOCdevkit", "VOC2009"), }, "2008": { "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2008/VOCtrainval_14-Jul-2008.tar", "filename": "VOCtrainval_11-May-2012.tar", "md5": "2629fa636546599198acfcfbfcf1904a", "base_dir": os.path.join("VOCdevkit", "VOC2008"), }, "2007": { "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar", "filename": "VOCtrainval_06-Nov-2007.tar", "md5": "c52e279531787c972589f7e41ab4ae64", "base_dir": os.path.join("VOCdevkit", "VOC2007"), }, "2007-test": { "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar", "filename": "VOCtest_06-Nov-2007.tar", "md5": "b6e924de25625d8de591ea690078ad9f", "base_dir": os.path.join("VOCdevkit", "VOC2007"), }, } class _VOCBase(VisionDataset): _SPLITS_DIR: str _TARGET_DIR: str _TARGET_FILE_EXT: str def __init__( self, root: str, year: str = "2012", image_set: str = "train", download: bool = False, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, transforms: Optional[Callable] = None, ): super().__init__(root, transforms, transform, target_transform) if year == "2007-test": if image_set == "test": warnings.warn( "Accessing the test image set of the year 2007 with year='2007-test' is deprecated " "since 0.12 and will be removed in 0.14. " "Please use the combination year='2007' and image_set='test' instead." ) year = "2007" else: raise ValueError( "In the test image set of the year 2007 only image_set='test' is allowed. " "For all other image sets use year='2007' instead." ) self.year = year valid_image_sets = ["train", "trainval", "val"] if year == "2007": valid_image_sets.append("test") self.image_set = verify_str_arg(image_set, "image_set", valid_image_sets) key = "2007-test" if year == "2007" and image_set == "test" else year dataset_year_dict = DATASET_YEAR_DICT[key] self.url = dataset_year_dict["url"] self.filename = dataset_year_dict["filename"] self.md5 = dataset_year_dict["md5"] base_dir = dataset_year_dict["base_dir"] voc_root = os.path.join(self.root, base_dir) if download: download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.md5) if not os.path.isdir(voc_root): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it") splits_dir = os.path.join(voc_root, "ImageSets", self._SPLITS_DIR) split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt") with open(os.path.join(split_f)) as f: file_names = [x.strip() for x in f.readlines()] image_dir = os.path.join(voc_root, "JPEGImages") self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names] target_dir = os.path.join(voc_root, self._TARGET_DIR) self.targets = [os.path.join(target_dir, x + self._TARGET_FILE_EXT) for x in file_names] assert len(self.images) == len(self.targets) def __len__(self) -> int: return len(self.images) class VOCSegmentation(_VOCBase): """`Pascal VOC `_ Segmentation Dataset. Args: root (string): Root directory of the VOC Dataset. year (string, optional): The dataset year, supports years ``"2007"`` to ``"2012"``. image_set (string, optional): Select the image_set to use, ``"train"``, ``"trainval"`` or ``"val"``. If ``year=="2007"``, can also be ``"test"``. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. """ _SPLITS_DIR = "Segmentation" _TARGET_DIR = "SegmentationClass" _TARGET_FILE_EXT = ".png" @property def masks(self) -> List[str]: return self.targets def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is the image segmentation. """ img = Image.open(self.images[index]).convert("RGB") target = Image.open(self.masks[index]) if self.transforms is not None: img, target = self.transforms(img, target) return img, target class VOCDetection(_VOCBase): """`Pascal VOC `_ Detection Dataset. Args: root (string): Root directory of the VOC Dataset. year (string, optional): The dataset year, supports years ``"2007"`` to ``"2012"``. image_set (string, optional): Select the image_set to use, ``"train"``, ``"trainval"`` or ``"val"``. If ``year=="2007"``, can also be ``"test"``. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. (default: alphabetic indexing of VOC's 20 classes). transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, required): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. """ _SPLITS_DIR = "Main" _TARGET_DIR = "Annotations" _TARGET_FILE_EXT = ".xml" @property def annotations(self) -> List[str]: return self.targets def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is a dictionary of the XML tree. """ img = Image.open(self.images[index]).convert("RGB") target = self.parse_voc_xml(ET_parse(self.annotations[index]).getroot()) if self.transforms is not None: img, target = self.transforms(img, target) return img, target @staticmethod def parse_voc_xml(node: ET_Element) -> Dict[str, Any]: voc_dict: Dict[str, Any] = {} children = list(node) if children: def_dic: Dict[str, Any] = collections.defaultdict(list) for dc in map(VOCDetection.parse_voc_xml, children): for ind, v in dc.items(): def_dic[ind].append(v) if node.tag == "annotation": def_dic["object"] = [def_dic["object"]] voc_dict = {node.tag: {ind: v[0] if len(v) == 1 else v for ind, v in def_dic.items()}} if node.text: text = node.text.strip() if not children: voc_dict[node.tag] = text return voc_dict msadapter/ms_adapter/torchvision/datasets/widerface.py000066400000000000000000000201661436713554100236550ustar00rootroot00000000000000import os import warnings from os.path import abspath, expanduser from typing import Any, Callable, List, Dict, Optional, Tuple, Union import ms_adapter.pytorch as torch from PIL import Image from .utils import ( download_file_from_google_drive, download_and_extract_archive, extract_archive, verify_str_arg, ) from .vision import VisionDataset class WIDERFace(VisionDataset): """`WIDERFace `_ Dataset. Args: root (string): Root directory where images and annotations are downloaded to. Expects the following folder structure if download=False: .. code:: └── widerface ├── wider_face_split ('wider_face_split.zip' if compressed) ├── WIDER_train ('WIDER_train.zip' if compressed) ├── WIDER_val ('WIDER_val.zip' if compressed) └── WIDER_test ('WIDER_test.zip' if compressed) split (string): The dataset split to use. One of {``train``, ``val``, ``test``}. Defaults to ``train``. transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version. E.g, ``transforms.RandomCrop`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. download (bool, optional): If true, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again. """ BASE_FOLDER = "widerface" FILE_LIST = [ # File ID MD5 Hash Filename ("15hGDLhsx8bLgLcIRD5DhYt5iBxnjNF1M", "3fedf70df600953d25982bcd13d91ba2", "WIDER_train.zip"), ("1GUCogbp16PMGa39thoMMeWxp7Rp5oM8Q", "dfa7d7e790efa35df3788964cf0bbaea", "WIDER_val.zip"), ("1HIfDbVEWKmsYKJZm4lchTBDLW5N7dY5T", "e5d8f4248ed24c334bbd12f49c29dd40", "WIDER_test.zip"), ] ANNOTATIONS_FILE = ( "http://shuoyang1213.me/WIDERFACE/support/bbx_annotation/wider_face_split.zip", "0e3767bcf0e326556d407bf5bff5d27c", "wider_face_split.zip", ) def __init__( self, root: str, split: str = "train", transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False, ) -> None: super().__init__( root=os.path.join(root, self.BASE_FOLDER), transform=transform, target_transform=target_transform ) # check arguments self.split = verify_str_arg(split, "split", ("train", "val", "test")) if download: self.download() if not self._check_integrity(): raise RuntimeError("Dataset not found or corrupted. You can use download=True to download and prepare it") self.img_info = [] if self.split in ("train", "val"): self.parse_train_val_annotations_file() else: self.parse_test_annotations_file() def __getitem__(self, index: int) -> Tuple[Any, Any]: """ Args: index (int): Index Returns: tuple: (image, target) where target is a dict of annotations for all faces in the image. target=None for the test split. """ # stay consistent with other datasets and return a PIL Image img = Image.open(self.img_info[index]["img_path"]) if self.transform is not None: img = self.transform(img) target = None if self.split == "test" else self.img_info[index]["annotations"] if self.target_transform is not None: target = self.target_transform(target) return img, target def __len__(self) -> int: return len(self.img_info) def extra_repr(self) -> str: lines = ["Split: {split}"] return "\n".join(lines).format(**self.__dict__) def parse_train_val_annotations_file(self) -> None: filename = "wider_face_train_bbx_gt.txt" if self.split == "train" else "wider_face_val_bbx_gt.txt" filepath = os.path.join(self.root, "wider_face_split", filename) with open(filepath) as f: lines = f.readlines() file_name_line, num_boxes_line, box_annotation_line = True, False, False num_boxes, box_counter = 0, 0 labels = [] for line in lines: line = line.rstrip() if file_name_line: img_path = os.path.join(self.root, "WIDER_" + self.split, "images", line) img_path = abspath(expanduser(img_path)) file_name_line = False num_boxes_line = True elif num_boxes_line: num_boxes = int(line) num_boxes_line = False box_annotation_line = True elif box_annotation_line: box_counter += 1 line_split = line.split(" ") line_values = [int(x) for x in line_split] labels.append(line_values) if box_counter >= num_boxes: box_annotation_line = False file_name_line = True labels_tensor = torch.tensor(labels) self.img_info.append( { "img_path": img_path, "annotations": { "bbox": labels_tensor[:, 0:4], # x, y, width, height "blur": labels_tensor[:, 4], "expression": labels_tensor[:, 5], "illumination": labels_tensor[:, 6], "occlusion": labels_tensor[:, 7], "pose": labels_tensor[:, 8], "invalid": labels_tensor[:, 9], }, } ) box_counter = 0 labels.clear() else: raise RuntimeError(f"Error parsing annotation file {filepath}") def parse_test_annotations_file(self) -> None: filepath = os.path.join(self.root, "wider_face_split", "wider_face_test_filelist.txt") filepath = abspath(expanduser(filepath)) with open(filepath) as f: lines = f.readlines() for line in lines: line = line.rstrip() img_path = os.path.join(self.root, "WIDER_test", "images", line) img_path = abspath(expanduser(img_path)) self.img_info.append({"img_path": img_path}) def _check_integrity(self) -> bool: # Allow original archive to be deleted (zip). Only need the extracted images all_files = self.FILE_LIST.copy() all_files.append(self.ANNOTATIONS_FILE) for (_, md5, filename) in all_files: file, ext = os.path.splitext(filename) extracted_dir = os.path.join(self.root, file) if not os.path.exists(extracted_dir): return False return True def download(self) -> None: if self._check_integrity(): print("Files already downloaded and verified") return warnings.warn("The dataset is stored on google drive. if you can't download it from google drive, " "please download it from official website." "WIDERFace ") # download and extract image data for (file_id, md5, filename) in self.FILE_LIST: download_file_from_google_drive(file_id, self.root, filename, md5) filepath = os.path.join(self.root, filename) extract_archive(filepath) # download and extract annotation files download_and_extract_archive( url=self.ANNOTATIONS_FILE[0], download_root=self.root, md5=self.ANNOTATIONS_FILE[1] ) msadapter/ms_adapter/torchvision/extension.py000066400000000000000000000061761436713554100221350ustar00rootroot00000000000000import ctypes import os import sys from warnings import warn from ._internally_replaced_utils import _get_extension_path # # _HAS_OPS = False def _has_ops(): return False # # # try: # lib_path = _get_extension_path("_C") # torch.ops.load_library(lib_path) # _HAS_OPS = True # # def _has_ops(): # noqa: F811 # return True # # # except (ImportError, OSError): # pass def _assert_has_ops(): if not _has_ops(): raise RuntimeError( "Couldn't load custom C++ ops. This can happen if your PyTorch and " "torchvision versions are incompatible, or if you had errors while compiling " "torchvision from source. For further information on the compatible versions, check " "https://github.com/pytorch/vision#installation for the compatibility matrix. " "Please check your PyTorch version with torch.__version__ and your torchvision " "version with torchvision.__version__ and verify if they are compatible, and if not " "please reinstall torchvision so that it matches your PyTorch install." ) # def _check_cuda_version(): # """ # Make sure that CUDA versions match between the pytorch install and torchvision install # """ # if not _HAS_OPS: # return -1 # import torch # # _version = torch.ops.torchvision._cuda_version() # if _version != -1 and torch.version.cuda is not None: # tv_version = str(_version) # if int(tv_version) < 10000: # tv_major = int(tv_version[0]) # tv_minor = int(tv_version[2]) # else: # tv_major = int(tv_version[0:2]) # tv_minor = int(tv_version[3]) # t_version = torch.version.cuda # t_version = t_version.split(".") # t_major = int(t_version[0]) # t_minor = int(t_version[1]) # if t_major != tv_major or t_minor != tv_minor: # raise RuntimeError( # "Detected that PyTorch and torchvision were compiled with different CUDA versions. " # f"PyTorch has CUDA Version={t_major}.{t_minor} and torchvision has " # f"CUDA Version={tv_major}.{tv_minor}. " # "Please reinstall the torchvision that matches your PyTorch install." # ) # return _version # # # def _load_library(lib_name): # lib_path = _get_extension_path(lib_name) # # On Windows Python-3.8+ has `os.add_dll_directory` call, # # which is called from _get_extension_path to configure dll search path # # Condition below adds a workaround for older versions by # # explicitly calling `LoadLibraryExW` with the following flags: # # - LOAD_LIBRARY_SEARCH_DEFAULT_DIRS (0x1000) # # - LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR (0x100) # if os.name == "nt" and sys.version_info < (3, 8): # _kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) # if hasattr(_kernel32, "LoadLibraryExW"): # _kernel32.LoadLibraryExW(lib_path, None, 0x00001100) # else: # warn("LoadLibraryExW is missing in kernel32.dll") # # torch.ops.load_library(lib_path) # # # _check_cuda_version() msadapter/ms_adapter/torchvision/io/000077500000000000000000000000001436713554100201445ustar00rootroot00000000000000msadapter/ms_adapter/torchvision/io/__init__.py000066400000000000000000000027701436713554100222630ustar00rootroot00000000000000from typing import Any, Dict, Iterator try: from ._load_gpu_decoder import _HAS_GPU_VIDEO_DECODER except ModuleNotFoundError: _HAS_GPU_VIDEO_DECODER = False from ._video_opt import ( Timebase, VideoMetaData, _HAS_VIDEO_OPT, _probe_video_from_file, _probe_video_from_memory, _read_video_from_file, _read_video_from_memory, _read_video_timestamps_from_file, _read_video_timestamps_from_memory, ) from .image import ( ImageReadMode, decode_image, decode_jpeg, decode_png, encode_jpeg, encode_png, read_file, read_image, write_file, write_jpeg, write_png, ) from .video import ( read_video, read_video_timestamps, write_video, ) # from .video_reader import VideoReader __all__ = [ "write_video", "read_video", "read_video_timestamps", "_read_video_from_file", "_read_video_timestamps_from_file", "_probe_video_from_file", "_read_video_from_memory", "_read_video_timestamps_from_memory", "_probe_video_from_memory", "_HAS_VIDEO_OPT", "_HAS_GPU_VIDEO_DECODER", # "_read_video_clip_from_memory", # TODO depend on torch.ops # "_read_video_meta_data",# TODO depend on torch.ops "VideoMetaData", "Timebase", "ImageReadMode", "decode_image", "decode_jpeg", "decode_png", "encode_jpeg", "encode_png", "read_file", "read_image", "write_file", "write_jpeg", "write_png", # "Video", # "VideoReader", # TODO depend on torch.ops ] msadapter/ms_adapter/torchvision/io/_load_gpu_decoder.py000066400000000000000000000003301436713554100241300ustar00rootroot00000000000000# from ..extension import _load_library # try: # _load_library("Decoder") # _HAS_GPU_VIDEO_DECODER = True # except (ImportError, OSError): # _HAS_GPU_VIDEO_DECODER = False _HAS_GPU_VIDEO_DECODER = Falsemsadapter/ms_adapter/torchvision/io/_video_opt.py000066400000000000000000000465111436713554100226540ustar00rootroot00000000000000import math import warnings from fractions import Fraction from typing import List, Tuple, Dict, Optional, Union import ms_adapter.pytorch as torch # from ..extension import _load_library # try: # _load_library("video_reader") # _HAS_VIDEO_OPT = True # except (ImportError, OSError): # _HAS_VIDEO_OPT = False _HAS_VIDEO_OPT = False default_timebase = Fraction(0, 1) # simple class for torch scripting # the complex Fraction class from fractions module is not scriptable class Timebase: __annotations__ = {"numerator": int, "denominator": int} __slots__ = ["numerator", "denominator"] def __init__( self, numerator: int, denominator: int, ) -> None: self.numerator = numerator self.denominator = denominator class VideoMetaData: __annotations__ = { "has_video": bool, "video_timebase": Timebase, "video_duration": float, "video_fps": float, "has_audio": bool, "audio_timebase": Timebase, "audio_duration": float, "audio_sample_rate": float, } __slots__ = [ "has_video", "video_timebase", "video_duration", "video_fps", "has_audio", "audio_timebase", "audio_duration", "audio_sample_rate", ] def __init__(self) -> None: self.has_video = False self.video_timebase = Timebase(0, 1) self.video_duration = 0.0 self.video_fps = 0.0 self.has_audio = False self.audio_timebase = Timebase(0, 1) self.audio_duration = 0.0 self.audio_sample_rate = 0.0 def _validate_pts(pts_range: Tuple[int, int]) -> None: if pts_range[0] > pts_range[1] > 0: raise ValueError( f"Start pts should not be smaller than end pts, got start pts: {pts_range[0]} and end pts: {pts_range[1]}" ) def _fill_info( vtimebase: torch.Tensor, vfps: torch.Tensor, vduration: torch.Tensor, atimebase: torch.Tensor, asample_rate: torch.Tensor, aduration: torch.Tensor, ) -> VideoMetaData: """ Build update VideoMetaData struct with info about the video """ meta = VideoMetaData() if vtimebase.numel() > 0: meta.video_timebase = Timebase(int(vtimebase[0].item()), int(vtimebase[1].item())) timebase = vtimebase[0].item() / float(vtimebase[1].item()) if vduration.numel() > 0: meta.has_video = True meta.video_duration = float(vduration.item()) * timebase if vfps.numel() > 0: meta.video_fps = float(vfps.item()) if atimebase.numel() > 0: meta.audio_timebase = Timebase(int(atimebase[0].item()), int(atimebase[1].item())) timebase = atimebase[0].item() / float(atimebase[1].item()) if aduration.numel() > 0: meta.has_audio = True meta.audio_duration = float(aduration.item()) * timebase if asample_rate.numel() > 0: meta.audio_sample_rate = float(asample_rate.item()) return meta def _align_audio_frames( aframes: torch.Tensor, aframe_pts: torch.Tensor, audio_pts_range: Tuple[int, int] ) -> torch.Tensor: start, end = aframe_pts[0], aframe_pts[-1] num_samples = aframes.size(0) step_per_aframe = float(end - start + 1) / float(num_samples) s_idx = 0 e_idx = num_samples if start < audio_pts_range[0]: s_idx = int((audio_pts_range[0] - start) / step_per_aframe) if audio_pts_range[1] != -1 and end > audio_pts_range[1]: e_idx = int((audio_pts_range[1] - end) / step_per_aframe) return aframes[s_idx:e_idx, :] def _read_video_from_file( filename: str, seek_frame_margin: float = 0.25, read_video_stream: bool = True, video_width: int = 0, video_height: int = 0, video_min_dimension: int = 0, video_max_dimension: int = 0, video_pts_range: Tuple[int, int] = (0, -1), video_timebase: Fraction = default_timebase, read_audio_stream: bool = True, audio_samples: int = 0, audio_channels: int = 0, audio_pts_range: Tuple[int, int] = (0, -1), audio_timebase: Fraction = default_timebase, ) -> Tuple[torch.Tensor, torch.Tensor, VideoMetaData]: """ Reads a video from a file, returning both the video frames as well as the audio frames Args: filename (str): path to the video file seek_frame_margin (double, optional): seeking frame in the stream is imprecise. Thus, when video_start_pts is specified, we seek the pts earlier by seek_frame_margin seconds read_video_stream (int, optional): whether read video stream. If yes, set to 1. Otherwise, 0 video_width/video_height/video_min_dimension/video_max_dimension (int): together decide the size of decoded frames: - When video_width = 0, video_height = 0, video_min_dimension = 0, and video_max_dimension = 0, keep the original frame resolution - When video_width = 0, video_height = 0, video_min_dimension != 0, and video_max_dimension = 0, keep the aspect ratio and resize the frame so that shorter edge size is video_min_dimension - When video_width = 0, video_height = 0, video_min_dimension = 0, and video_max_dimension != 0, keep the aspect ratio and resize the frame so that longer edge size is video_max_dimension - When video_width = 0, video_height = 0, video_min_dimension != 0, and video_max_dimension != 0, resize the frame so that shorter edge size is video_min_dimension, and longer edge size is video_max_dimension. The aspect ratio may not be preserved - When video_width = 0, video_height != 0, video_min_dimension = 0, and video_max_dimension = 0, keep the aspect ratio and resize the frame so that frame video_height is $video_height - When video_width != 0, video_height == 0, video_min_dimension = 0, and video_max_dimension = 0, keep the aspect ratio and resize the frame so that frame video_width is $video_width - When video_width != 0, video_height != 0, video_min_dimension = 0, and video_max_dimension = 0, resize the frame so that frame video_width and video_height are set to $video_width and $video_height, respectively video_pts_range (list(int), optional): the start and end presentation timestamp of video stream video_timebase (Fraction, optional): a Fraction rational number which denotes timebase in video stream read_audio_stream (int, optional): whether read audio stream. If yes, set to 1. Otherwise, 0 audio_samples (int, optional): audio sampling rate audio_channels (int optional): audio channels audio_pts_range (list(int), optional): the start and end presentation timestamp of audio stream audio_timebase (Fraction, optional): a Fraction rational number which denotes time base in audio stream Returns vframes (Tensor[T, H, W, C]): the `T` video frames aframes (Tensor[L, K]): the audio frames, where `L` is the number of points and `K` is the number of audio_channels info (Dict): metadata for the video and audio. Can contain the fields video_fps (float) and audio_fps (int) """ _validate_pts(video_pts_range) _validate_pts(audio_pts_range) result = torch.ops.video_reader.read_video_from_file( filename, seek_frame_margin, 0, # getPtsOnly read_video_stream, video_width, video_height, video_min_dimension, video_max_dimension, video_pts_range[0], video_pts_range[1], video_timebase.numerator, video_timebase.denominator, read_audio_stream, audio_samples, audio_channels, audio_pts_range[0], audio_pts_range[1], audio_timebase.numerator, audio_timebase.denominator, ) vframes, _vframe_pts, vtimebase, vfps, vduration, aframes, aframe_pts, atimebase, asample_rate, aduration = result info = _fill_info(vtimebase, vfps, vduration, atimebase, asample_rate, aduration) if aframes.numel() > 0: # when audio stream is found aframes = _align_audio_frames(aframes, aframe_pts, audio_pts_range) return vframes, aframes, info def _read_video_timestamps_from_file(filename: str) -> Tuple[List[int], List[int], VideoMetaData]: """ Decode all video- and audio frames in the video. Only pts (presentation timestamp) is returned. The actual frame pixel data is not copied. Thus, it is much faster than read_video(...) """ result = torch.ops.video_reader.read_video_from_file( filename, 0, # seek_frame_margin 1, # getPtsOnly 1, # read_video_stream 0, # video_width 0, # video_height 0, # video_min_dimension 0, # video_max_dimension 0, # video_start_pts -1, # video_end_pts 0, # video_timebase_num 1, # video_timebase_den 1, # read_audio_stream 0, # audio_samples 0, # audio_channels 0, # audio_start_pts -1, # audio_end_pts 0, # audio_timebase_num 1, # audio_timebase_den ) _vframes, vframe_pts, vtimebase, vfps, vduration, _aframes, aframe_pts, atimebase, asample_rate, aduration = result info = _fill_info(vtimebase, vfps, vduration, atimebase, asample_rate, aduration) vframe_pts = vframe_pts.numpy().tolist() aframe_pts = aframe_pts.numpy().tolist() return vframe_pts, aframe_pts, info def _probe_video_from_file(filename: str) -> VideoMetaData: """ Probe a video file and return VideoMetaData with info about the video """ result = torch.ops.video_reader.probe_video_from_file(filename) vtimebase, vfps, vduration, atimebase, asample_rate, aduration = result info = _fill_info(vtimebase, vfps, vduration, atimebase, asample_rate, aduration) return info def _read_video_from_memory( video_data: torch.Tensor, seek_frame_margin: float = 0.25, read_video_stream: int = 1, video_width: int = 0, video_height: int = 0, video_min_dimension: int = 0, video_max_dimension: int = 0, video_pts_range: Tuple[int, int] = (0, -1), video_timebase_numerator: int = 0, video_timebase_denominator: int = 1, read_audio_stream: int = 1, audio_samples: int = 0, audio_channels: int = 0, audio_pts_range: Tuple[int, int] = (0, -1), audio_timebase_numerator: int = 0, audio_timebase_denominator: int = 1, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Reads a video from memory, returning both the video frames as well as the audio frames This function is torchscriptable. Args: video_data (data type could be 1) torch.Tensor, dtype=torch.int8 or 2) python bytes): compressed video content stored in either 1) torch.Tensor 2) python bytes seek_frame_margin (double, optional): seeking frame in the stream is imprecise. Thus, when video_start_pts is specified, we seek the pts earlier by seek_frame_margin seconds read_video_stream (int, optional): whether read video stream. If yes, set to 1. Otherwise, 0 video_width/video_height/video_min_dimension/video_max_dimension (int): together decide the size of decoded frames: - When video_width = 0, video_height = 0, video_min_dimension = 0, and video_max_dimension = 0, keep the original frame resolution - When video_width = 0, video_height = 0, video_min_dimension != 0, and video_max_dimension = 0, keep the aspect ratio and resize the frame so that shorter edge size is video_min_dimension - When video_width = 0, video_height = 0, video_min_dimension = 0, and video_max_dimension != 0, keep the aspect ratio and resize the frame so that longer edge size is video_max_dimension - When video_width = 0, video_height = 0, video_min_dimension != 0, and video_max_dimension != 0, resize the frame so that shorter edge size is video_min_dimension, and longer edge size is video_max_dimension. The aspect ratio may not be preserved - When video_width = 0, video_height != 0, video_min_dimension = 0, and video_max_dimension = 0, keep the aspect ratio and resize the frame so that frame video_height is $video_height - When video_width != 0, video_height == 0, video_min_dimension = 0, and video_max_dimension = 0, keep the aspect ratio and resize the frame so that frame video_width is $video_width - When video_width != 0, video_height != 0, video_min_dimension = 0, and video_max_dimension = 0, resize the frame so that frame video_width and video_height are set to $video_width and $video_height, respectively video_pts_range (list(int), optional): the start and end presentation timestamp of video stream video_timebase_numerator / video_timebase_denominator (float, optional): a rational number which denotes timebase in video stream read_audio_stream (int, optional): whether read audio stream. If yes, set to 1. Otherwise, 0 audio_samples (int, optional): audio sampling rate audio_channels (int optional): audio audio_channels audio_pts_range (list(int), optional): the start and end presentation timestamp of audio stream audio_timebase_numerator / audio_timebase_denominator (float, optional): a rational number which denotes time base in audio stream Returns: vframes (Tensor[T, H, W, C]): the `T` video frames aframes (Tensor[L, K]): the audio frames, where `L` is the number of points and `K` is the number of channels """ _validate_pts(video_pts_range) _validate_pts(audio_pts_range) if not isinstance(video_data, torch.Tensor): video_data = torch.frombuffer(video_data, dtype=torch.uint8) result = torch.ops.video_reader.read_video_from_memory( video_data, seek_frame_margin, 0, # getPtsOnly read_video_stream, video_width, video_height, video_min_dimension, video_max_dimension, video_pts_range[0], video_pts_range[1], video_timebase_numerator, video_timebase_denominator, read_audio_stream, audio_samples, audio_channels, audio_pts_range[0], audio_pts_range[1], audio_timebase_numerator, audio_timebase_denominator, ) vframes, _vframe_pts, vtimebase, vfps, vduration, aframes, aframe_pts, atimebase, asample_rate, aduration = result if aframes.numel() > 0: # when audio stream is found aframes = _align_audio_frames(aframes, aframe_pts, audio_pts_range) return vframes, aframes def _read_video_timestamps_from_memory( video_data: torch.Tensor, ) -> Tuple[List[int], List[int], VideoMetaData]: """ Decode all frames in the video. Only pts (presentation timestamp) is returned. The actual frame pixel data is not copied. Thus, read_video_timestamps(...) is much faster than read_video(...) """ if not isinstance(video_data, torch.Tensor): video_data = torch.frombuffer(video_data, dtype=torch.uint8) result = torch.ops.video_reader.read_video_from_memory( video_data, 0, # seek_frame_margin 1, # getPtsOnly 1, # read_video_stream 0, # video_width 0, # video_height 0, # video_min_dimension 0, # video_max_dimension 0, # video_start_pts -1, # video_end_pts 0, # video_timebase_num 1, # video_timebase_den 1, # read_audio_stream 0, # audio_samples 0, # audio_channels 0, # audio_start_pts -1, # audio_end_pts 0, # audio_timebase_num 1, # audio_timebase_den ) _vframes, vframe_pts, vtimebase, vfps, vduration, _aframes, aframe_pts, atimebase, asample_rate, aduration = result info = _fill_info(vtimebase, vfps, vduration, atimebase, asample_rate, aduration) vframe_pts = vframe_pts.numpy().tolist() aframe_pts = aframe_pts.numpy().tolist() return vframe_pts, aframe_pts, info def _probe_video_from_memory( video_data: torch.Tensor, ) -> VideoMetaData: """ Probe a video in memory and return VideoMetaData with info about the video This function is torchscriptable """ if not isinstance(video_data, torch.Tensor): video_data = torch.frombuffer(video_data, dtype=torch.uint8) result = torch.ops.video_reader.probe_video_from_memory(video_data) vtimebase, vfps, vduration, atimebase, asample_rate, aduration = result info = _fill_info(vtimebase, vfps, vduration, atimebase, asample_rate, aduration) return info def _read_video( filename: str, start_pts: Union[float, Fraction] = 0, end_pts: Optional[Union[float, Fraction]] = None, pts_unit: str = "pts", ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, float]]: if end_pts is None: end_pts = float("inf") if pts_unit == "pts": warnings.warn( "The pts_unit 'pts' gives wrong results and will be removed in a " + "follow-up version. Please use pts_unit 'sec'." ) info = _probe_video_from_file(filename) has_video = info.has_video has_audio = info.has_audio def get_pts(time_base): start_offset = start_pts end_offset = end_pts if pts_unit == "sec": start_offset = int(math.floor(start_pts * (1 / time_base))) if end_offset != float("inf"): end_offset = int(math.ceil(end_pts * (1 / time_base))) if end_offset == float("inf"): end_offset = -1 return start_offset, end_offset video_pts_range = (0, -1) video_timebase = default_timebase if has_video: video_timebase = Fraction(info.video_timebase.numerator, info.video_timebase.denominator) video_pts_range = get_pts(video_timebase) audio_pts_range = (0, -1) audio_timebase = default_timebase if has_audio: audio_timebase = Fraction(info.audio_timebase.numerator, info.audio_timebase.denominator) audio_pts_range = get_pts(audio_timebase) vframes, aframes, info = _read_video_from_file( filename, read_video_stream=True, video_pts_range=video_pts_range, video_timebase=video_timebase, read_audio_stream=True, audio_pts_range=audio_pts_range, audio_timebase=audio_timebase, ) _info = {} if has_video: _info["video_fps"] = info.video_fps if has_audio: _info["audio_fps"] = info.audio_sample_rate return vframes, aframes, _info def _read_video_timestamps( filename: str, pts_unit: str = "pts" ) -> Tuple[Union[List[int], List[Fraction]], Optional[float]]: if pts_unit == "pts": warnings.warn( "The pts_unit 'pts' gives wrong results and will be removed in a " + "follow-up version. Please use pts_unit 'sec'." ) pts: Union[List[int], List[Fraction]] pts, _, info = _read_video_timestamps_from_file(filename) if pts_unit == "sec": video_time_base = Fraction(info.video_timebase.numerator, info.video_timebase.denominator) pts = [x * video_time_base for x in pts] video_fps = info.video_fps if info.has_video else None return pts, video_fps msadapter/ms_adapter/torchvision/io/image.py000066400000000000000000000233331436713554100216040ustar00rootroot00000000000000from enum import Enum from warnings import warn import ms_adapter.pytorch as torch from mindspore.dataset import vision # try: # _load_library("image") # except (ImportError, OSError) as e: # warn(f"Failed to load image Python extension: {e}") class ImageReadMode(Enum): """ Support for various modes while reading images. Use ``ImageReadMode.UNCHANGED`` for loading the image as-is, ``ImageReadMode.GRAY`` for converting to grayscale, ``ImageReadMode.GRAY_ALPHA`` for grayscale with transparency, ``ImageReadMode.RGB`` for RGB and ``ImageReadMode.RGB_ALPHA`` for RGB with transparency. """ UNCHANGED = 0 GRAY = 1 GRAY_ALPHA = 2 RGB = 3 RGB_ALPHA = 4 def read_file(path: str) -> torch.Tensor: """ Reads and outputs the bytes contents of a file as a uint8 Tensor with one dimension. Args: path (str): the path to the file to be read Returns: data (Tensor) """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(read_file) data = vision.read_file(path) return data def write_file(filename: str, data: torch.Tensor) -> None: """ Writes the contents of a uint8 tensor with one dimension to a file. Args: filename (str): the path to the file to be written data (Tensor): the contents to be written to the output file """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(write_file) vision.write_file(filename, data) def decode_png(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor: """ Decodes a PNG image into a 3 dimensional RGB or grayscale Tensor. Optionally converts the image to the desired format. The values of the output tensor are uint8 in [0, 255]. Args: input (Tensor[1]): a one dimensional uint8 tensor containing the raw bytes of the PNG image. mode (ImageReadMode): the read mode used for optionally converting the image. Default: ``ImageReadMode.UNCHANGED``. See `ImageReadMode` class for more information on various available modes. Returns: output (Tensor[image_channels, image_height, image_width]) """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(decode_png) # output = vision.decode_png(input, mode.value, False) output = vision.Decode()(input) output = torch.tensor(output, dtype=torch.uint8) output = output.permute(2, 0, 1) return output def encode_png(input: torch.Tensor, compression_level: int = 6) -> torch.Tensor: """ Takes an input tensor in CHW layout and returns a buffer with the contents of its corresponding PNG file. Args: input (Tensor[channels, image_height, image_width]): int8 image tensor of ``c`` channels, where ``c`` must 3 or 1. compression_level (int): Compression factor for the resulting file, it must be a number between 0 and 9. Default: 6 Returns: Tensor[1]: A one dimensional int8 tensor that contains the raw bytes of the PNG file. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(encode_png) output = vision.encode_png(input, compression_level) return output def write_png(input: torch.Tensor, filename: str, compression_level: int = 6): """ Takes an input tensor in CHW layout (or HW in the case of grayscale images) and saves it in a PNG file. Args: input (Tensor[channels, image_height, image_width]): int8 image tensor of ``c`` channels, where ``c`` must be 1 or 3. filename (str): Path to save the image. compression_level (int): Compression factor for the resulting file, it must be a number between 0 and 9. Default: 6 """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(write_png) output = encode_png(input, compression_level) write_file(filename, output) def decode_jpeg( input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED, device: str = "cpu" ) -> torch.Tensor: """ Decodes a JPEG image into a 3 dimensional RGB or grayscale Tensor. Optionally converts the image to the desired format. The values of the output tensor are uint8 between 0 and 255. Args: input (Tensor[1]): a one dimensional uint8 tensor containing the raw bytes of the JPEG image. This tensor must be on CPU, regardless of the ``device`` parameter. mode (ImageReadMode): the read mode used for optionally converting the image. Default: ``ImageReadMode.UNCHANGED``. See ``ImageReadMode`` class for more information on various available modes. device (str or torch.device): The device on which the decoded image will be stored. If a cuda device is specified, the image will be decoded with `nvjpeg `_. This is only supported for CUDA version >= 10.1 .. betastatus:: device parameter .. warning:: There is a memory leak in the nvjpeg library for CUDA versions < 11.6. Make sure to rely on CUDA 11.6 or above before using ``device="cuda"``. Returns: output (Tensor[image_channels, image_height, image_width]) """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(decode_jpeg) # device = torch.device(device) # if device.type == "cuda": # output = torch.ops.image.decode_jpeg_cuda(input, mode.value, device) # else: # output = torch.ops.image.decode_jpeg(input, mode.value) output = vision.Decode()(input) output = torch.tensor(output, dtype=torch.uint8) output = output.permute(2, 0, 1) return output def encode_jpeg(input: torch.Tensor, quality: int = 75) -> torch.Tensor: """ Takes an input tensor in CHW layout and returns a buffer with the contents of its corresponding JPEG file. Args: input (Tensor[channels, image_height, image_width])): int8 image tensor of ``c`` channels, where ``c`` must be 1 or 3. quality (int): Quality of the resulting JPEG file, it must be a number between 1 and 100. Default: 75 Returns: output (Tensor[1]): A one dimensional int8 tensor that contains the raw bytes of the JPEG file. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(encode_jpeg) if quality < 1 or quality > 100: raise ValueError("Image quality should be a positive number between 1 and 100") output = vision.encode_jpeg(input, quality) return output def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75): """ Takes an input tensor in CHW layout and saves it in a JPEG file. Args: input (Tensor[channels, image_height, image_width]): int8 image tensor of ``c`` channels, where ``c`` must be 1 or 3. filename (str): Path to save the image. quality (int): Quality of the resulting JPEG file, it must be a number between 1 and 100. Default: 75 """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(write_jpeg) output = encode_jpeg(input, quality) write_file(filename, output) def decode_image(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor: """ Detects whether an image is a JPEG or PNG and performs the appropriate operation to decode the image into a 3 dimensional RGB or grayscale Tensor. Optionally converts the image to the desired format. The values of the output tensor are uint8 in [0, 255]. Args: input (Tensor): a one dimensional uint8 tensor containing the raw bytes of the PNG or JPEG image. mode (ImageReadMode): the read mode used for optionally converting the image. Default: ``ImageReadMode.UNCHANGED``. See ``ImageReadMode`` class for more information on various available modes. Returns: output (Tensor[image_channels, image_height, image_width]) """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(decode_image) # output = vision.decode_image(input, mode.value) output = vision.Decode()(input) output = torch.tensor(output, dtype=torch.uint8) output = output.permute(2, 0, 1) return output def read_image(path: str, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor: """ Reads a JPEG or PNG image into a 3 dimensional RGB or grayscale Tensor. Optionally converts the image to the desired format. The values of the output tensor are uint8 in [0, 255]. Args: path (str): path of the JPEG or PNG image. mode (ImageReadMode): the read mode used for optionally converting the image. Default: ``ImageReadMode.UNCHANGED``. See ``ImageReadMode`` class for more information on various available modes. Returns: output (Tensor[image_channels, image_height, image_width]) """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(read_image) data = read_file(path) return decode_image(data, mode) def _read_png_16(path: str, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor: data = read_file(path) # return vision.decode_png(data, mode.value, True) output = vision.Decode()(data) output = torch.tensor(output, dtype=torch.uint8) output = output.permute(2, 0, 1) return outputmsadapter/ms_adapter/torchvision/io/video.py000066400000000000000000000361011436713554100216250ustar00rootroot00000000000000import gc import math import os import re import warnings from fractions import Fraction from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import ms_adapter.pytorch as torch from . import _video_opt try: import av av.logging.set_level(av.logging.ERROR) if not hasattr(av.video.frame.VideoFrame, "pict_type"): av = ImportError( """\ Your version of PyAV is too old for the necessary video operations in torchvision. If you are on Python 3.5, you will have to build from source (the conda-forge packages are not up-to-date). See https://github.com/mikeboers/PyAV#installation for instructions on how to install PyAV on your system. """ ) except ImportError: av = ImportError( """\ PyAV is not installed, and is necessary for the video operations in torchvision. See https://github.com/mikeboers/PyAV#installation for instructions on how to install PyAV on your system. """ ) def _check_av_available() -> None: if isinstance(av, Exception): raise av def _av_available() -> bool: return not isinstance(av, Exception) # PyAV has some reference cycles _CALLED_TIMES = 0 _GC_COLLECTION_INTERVAL = 10 def write_video( filename: str, video_array: torch.Tensor, fps: float, video_codec: str = "libx264", options: Optional[Dict[str, Any]] = None, audio_array: Optional[torch.Tensor] = None, audio_fps: Optional[float] = None, audio_codec: Optional[str] = None, audio_options: Optional[Dict[str, Any]] = None, ) -> None: """ Writes a 4d tensor in [T, H, W, C] format in a video file Args: filename (str): path where the video will be saved video_array (Tensor[T, H, W, C]): tensor containing the individual frames, as a uint8 tensor in [T, H, W, C] format fps (Number): video frames per second video_codec (str): the name of the video codec, i.e. "libx264", "h264", etc. options (Dict): dictionary containing options to be passed into the PyAV video stream audio_array (Tensor[C, N]): tensor containing the audio, where C is the number of channels and N is the number of samples audio_fps (Number): audio sample rate, typically 44100 or 48000 audio_codec (str): the name of the audio codec, i.e. "mp3", "aac", etc. audio_options (Dict): dictionary containing options to be passed into the PyAV audio stream """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(write_video) _check_av_available() video_array = torch.as_tensor(video_array, dtype=torch.uint8).numpy() # PyAV does not support floating point numbers with decimal point # and will throw OverflowException in case this is not the case if isinstance(fps, float): fps = np.round(fps) with av.open(filename, mode="w") as container: stream = container.add_stream(video_codec, rate=fps) stream.width = video_array.shape[2] stream.height = video_array.shape[1] stream.pix_fmt = "yuv420p" if video_codec != "libx264rgb" else "rgb24" stream.options = options or {} if audio_array is not None: audio_format_dtypes = { "dbl": " 1 else "mono" audio_sample_fmt = container.streams.audio[0].format.name format_dtype = np.dtype(audio_format_dtypes[audio_sample_fmt]) audio_array = torch.as_tensor(audio_array).numpy().astype(format_dtype) frame = av.AudioFrame.from_ndarray(audio_array, format=audio_sample_fmt, layout=audio_layout) frame.sample_rate = audio_fps for packet in a_stream.encode(frame): container.mux(packet) for packet in a_stream.encode(): container.mux(packet) for img in video_array: frame = av.VideoFrame.from_ndarray(img, format="rgb24") frame.pict_type = "NONE" for packet in stream.encode(frame): container.mux(packet) # Flush stream for packet in stream.encode(): container.mux(packet) def _read_from_stream( container: "av.container.Container", start_offset: float, end_offset: float, pts_unit: str, stream: "av.stream.Stream", stream_name: Dict[str, Optional[Union[int, Tuple[int, ...], List[int]]]], ) -> List["av.frame.Frame"]: global _CALLED_TIMES, _GC_COLLECTION_INTERVAL _CALLED_TIMES += 1 if _CALLED_TIMES % _GC_COLLECTION_INTERVAL == _GC_COLLECTION_INTERVAL - 1: gc.collect() if pts_unit == "sec": # TODO: we should change all of this from ground up to simply take # sec and convert to MS in C++ start_offset = int(math.floor(start_offset * (1 / stream.time_base))) if end_offset != float("inf"): end_offset = int(math.ceil(end_offset * (1 / stream.time_base))) else: warnings.warn("The pts_unit 'pts' gives wrong results. Please use pts_unit 'sec'.") frames = {} should_buffer = True max_buffer_size = 5 if stream.type == "video": # DivX-style packed B-frames can have out-of-order pts (2 frames in a single pkt) # so need to buffer some extra frames to sort everything # properly extradata = stream.codec_context.extradata # overly complicated way of finding if `divx_packed` is set, following # https://github.com/FFmpeg/FFmpeg/commit/d5a21172283572af587b3d939eba0091484d3263 if extradata and b"DivX" in extradata: # can't use regex directly because of some weird characters sometimes... pos = extradata.find(b"DivX") d = extradata[pos:] o = re.search(rb"DivX(\d+)Build(\d+)(\w)", d) if o is None: o = re.search(rb"DivX(\d+)b(\d+)(\w)", d) if o is not None: should_buffer = o.group(3) == b"p" seek_offset = start_offset # some files don't seek to the right location, so better be safe here seek_offset = max(seek_offset - 1, 0) if should_buffer: # FIXME this is kind of a hack, but we will jump to the previous keyframe # so this will be safe seek_offset = max(seek_offset - max_buffer_size, 0) try: # TODO check if stream needs to always be the video stream here or not container.seek(seek_offset, any_frame=False, backward=True, stream=stream) except av.AVError: # TODO add some warnings in this case # print("Corrupted file?", container.name) return [] buffer_count = 0 try: for _idx, frame in enumerate(container.decode(**stream_name)): frames[frame.pts] = frame if frame.pts >= end_offset: if should_buffer and buffer_count < max_buffer_size: buffer_count += 1 continue break except av.AVError: # TODO add a warning pass # ensure that the results are sorted wrt the pts result = [frames[i] for i in sorted(frames) if start_offset <= frames[i].pts <= end_offset] if len(frames) > 0 and start_offset > 0 and start_offset not in frames: # if there is no frame that exactly matches the pts of start_offset # add the last frame smaller than start_offset, to guarantee that # we will have all the necessary data. This is most useful for audio preceding_frames = [i for i in frames if i < start_offset] if len(preceding_frames) > 0: first_frame_pts = max(preceding_frames) result.insert(0, frames[first_frame_pts]) return result def _align_audio_frames( aframes: torch.Tensor, audio_frames: List["av.frame.Frame"], ref_start: int, ref_end: float ) -> torch.Tensor: start, end = audio_frames[0].pts, audio_frames[-1].pts total_aframes = aframes.shape[1] step_per_aframe = (end - start + 1) / total_aframes s_idx = 0 e_idx = total_aframes if start < ref_start: s_idx = int((ref_start - start) / step_per_aframe) if end > ref_end: e_idx = int((ref_end - end) / step_per_aframe) return aframes[:, s_idx:e_idx] def read_video( filename: str, start_pts: Union[float, Fraction] = 0, end_pts: Optional[Union[float, Fraction]] = None, pts_unit: str = "pts", output_format: str = "THWC", ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]: """ Reads a video from a file, returning both the video frames as well as the audio frames Args: filename (str): path to the video file start_pts (int if pts_unit = 'pts', float / Fraction if pts_unit = 'sec', optional): The start presentation time of the video end_pts (int if pts_unit = 'pts', float / Fraction if pts_unit = 'sec', optional): The end presentation time pts_unit (str, optional): unit in which start_pts and end_pts values will be interpreted, either 'pts' or 'sec'. Defaults to 'pts'. output_format (str, optional): The format of the output video tensors. Can be either "THWC" (default) or "TCHW". Returns: vframes (Tensor[T, H, W, C] or Tensor[T, C, H, W]): the `T` video frames aframes (Tensor[K, L]): the audio frames, where `K` is the number of channels and `L` is the number of points info (Dict): metadata for the video and audio. Can contain the fields video_fps (float) and audio_fps (int) """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(read_video) output_format = output_format.upper() if output_format not in ("THWC", "TCHW"): raise ValueError(f"output_format should be either 'THWC' or 'TCHW', got {output_format}.") from ms_adapter.torchvision import get_video_backend if not os.path.exists(filename): raise RuntimeError(f"File not found: {filename}") if get_video_backend() != "pyav": return _video_opt._read_video(filename, start_pts, end_pts, pts_unit) _check_av_available() if end_pts is None: end_pts = float("inf") if end_pts < start_pts: raise ValueError(f"end_pts should be larger than start_pts, got start_pts={start_pts} and end_pts={end_pts}") info = {} video_frames = [] audio_frames = [] audio_timebase = _video_opt.default_timebase try: with av.open(filename, metadata_errors="ignore") as container: if container.streams.audio: audio_timebase = container.streams.audio[0].time_base if container.streams.video: video_frames = _read_from_stream( container, start_pts, end_pts, pts_unit, container.streams.video[0], {"video": 0}, ) video_fps = container.streams.video[0].average_rate # guard against potentially corrupted files if video_fps is not None: info["video_fps"] = float(video_fps) if container.streams.audio: audio_frames = _read_from_stream( container, start_pts, end_pts, pts_unit, container.streams.audio[0], {"audio": 0}, ) info["audio_fps"] = container.streams.audio[0].rate except av.AVError: # TODO raise a warning? pass vframes_list = [frame.to_rgb().to_ndarray() for frame in video_frames] aframes_list = [frame.to_ndarray() for frame in audio_frames] if vframes_list: vframes = torch.as_tensor(np.stack(vframes_list)) else: vframes = torch.empty((0, 1, 1, 3), dtype=torch.uint8) if aframes_list: aframes = np.concatenate(aframes_list, 1) aframes = torch.as_tensor(aframes) if pts_unit == "sec": start_pts = int(math.floor(start_pts * (1 / audio_timebase))) if end_pts != float("inf"): end_pts = int(math.ceil(end_pts * (1 / audio_timebase))) aframes = _align_audio_frames(aframes, audio_frames, start_pts, end_pts) else: aframes = torch.empty((1, 0), dtype=torch.float32) if output_format == "TCHW": # [T,H,W,C] --> [T,C,H,W] vframes = vframes.permute(0, 3, 1, 2) return vframes, aframes, info def _can_read_timestamps_from_packets(container: "av.container.Container") -> bool: extradata = container.streams[0].codec_context.extradata if extradata is None: return False if b"Lavc" in extradata: return True return False def _decode_video_timestamps(container: "av.container.Container") -> List[int]: if _can_read_timestamps_from_packets(container): # fast path return [x.pts for x in container.demux(video=0) if x.pts is not None] else: return [x.pts for x in container.decode(video=0) if x.pts is not None] def read_video_timestamps(filename: str, pts_unit: str = "pts") -> Tuple[List[int], Optional[float]]: """ List the video frames timestamps. Note that the function decodes the whole video frame-by-frame. Args: filename (str): path to the video file pts_unit (str, optional): unit in which timestamp values will be returned either 'pts' or 'sec'. Defaults to 'pts'. Returns: pts (List[int] if pts_unit = 'pts', List[Fraction] if pts_unit = 'sec'): presentation timestamps for each one of the frames in the video. video_fps (float, optional): the frame rate for the video """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(read_video_timestamps) from ms_adapter.torchvision import get_video_backend if get_video_backend() != "pyav": return _video_opt._read_video_timestamps(filename, pts_unit) _check_av_available() video_fps = None pts = [] try: with av.open(filename, metadata_errors="ignore") as container: if container.streams.video: video_stream = container.streams.video[0] video_time_base = video_stream.time_base try: pts = _decode_video_timestamps(container) except av.AVError: warnings.warn(f"Failed decoding frames for file {filename}") video_fps = float(video_stream.average_rate) except av.AVError as e: msg = f"Failed to open container for {filename}; Caught error: {e}" warnings.warn(msg, RuntimeWarning) pts.sort() if pts_unit == "sec": pts = [x * video_time_base for x in pts] return pts, video_fps msadapter/ms_adapter/torchvision/io/video_reader.py000066400000000000000000000147261436713554100231600ustar00rootroot00000000000000from typing import Any, Dict, Iterator import torch try: from ._load_gpu_decoder import _HAS_GPU_VIDEO_DECODER except ModuleNotFoundError: _HAS_GPU_VIDEO_DECODER = False from ._video_opt import ( _HAS_VIDEO_OPT, ) if _HAS_VIDEO_OPT: def _has_video_opt() -> bool: return True else: def _has_video_opt() -> bool: return False class VideoReader: """ Fine-grained video-reading API. Supports frame-by-frame reading of various streams from a single video container. .. betastatus:: VideoReader class Example: The following examples creates a :mod:`VideoReader` object, seeks into 2s point, and returns a single frame:: import torchvision video_path = "path_to_a_test_video" reader = torchvision.io.VideoReader(video_path, "video") reader.seek(2.0) frame = next(reader) :mod:`VideoReader` implements the iterable API, which makes it suitable to using it in conjunction with :mod:`itertools` for more advanced reading. As such, we can use a :mod:`VideoReader` instance inside for loops:: reader.seek(2) for frame in reader: frames.append(frame['data']) # additionally, `seek` implements a fluent API, so we can do for frame in reader.seek(2): frames.append(frame['data']) With :mod:`itertools`, we can read all frames between 2 and 5 seconds with the following code:: for frame in itertools.takewhile(lambda x: x['pts'] <= 5, reader.seek(2)): frames.append(frame['data']) and similarly, reading 10 frames after the 2s timestamp can be achieved as follows:: for frame in itertools.islice(reader.seek(2), 10): frames.append(frame['data']) .. note:: Each stream descriptor consists of two parts: stream type (e.g. 'video') and a unique stream id (which are determined by the video encoding). In this way, if the video contaner contains multiple streams of the same type, users can acces the one they want. If only stream type is passed, the decoder auto-detects first stream of that type. Args: path (string): Path to the video file in supported format stream (string, optional): descriptor of the required stream, followed by the stream id, in the format ``{stream_type}:{stream_id}``. Defaults to ``"video:0"``. Currently available options include ``['video', 'audio']`` num_threads (int, optional): number of threads used by the codec to decode video. Default value (0) enables multithreading with codec-dependent heuristic. The performance will depend on the version of FFMPEG codecs supported. device (str, optional): Device to be used for decoding. Defaults to ``"cpu"``. To use GPU decoding, pass ``device="cuda"``. """ def __init__(self, path: str, stream: str = "video", num_threads: int = 0, device: str = "cpu") -> None: # _log_api_usage_once(self) self.is_cuda = False device = torch.device(device) if device.type == "cuda": if not _HAS_GPU_VIDEO_DECODER: raise RuntimeError("Not compiled with GPU decoder support.") self.is_cuda = True self._c = torch.classes.torchvision.GPUDecoder(path, device) return if not _has_video_opt(): raise RuntimeError( "Not compiled with video_reader support, " + "to enable video_reader support, please install " + "ffmpeg (version 4.2 is currently supported) and " + "build torchvision from source." ) self._c = torch.classes.torchvision.Video(path, stream, num_threads) def __next__(self) -> Dict[str, Any]: """Decodes and returns the next frame of the current stream. Frames are encoded as a dict with mandatory data and pts fields, where data is a tensor, and pts is a presentation timestamp of the frame expressed in seconds as a float. Returns: (dict): a dictionary and containing decoded frame (``data``) and corresponding timestamp (``pts``) in seconds """ if self.is_cuda: frame = self._c.next() if frame.numel() == 0: raise StopIteration return {"data": frame} frame, pts = self._c.next() if frame.numel() == 0: raise StopIteration return {"data": frame, "pts": pts} def __iter__(self) -> Iterator[Dict[str, Any]]: return self def seek(self, time_s: float, keyframes_only: bool = False) -> "VideoReader": """Seek within current stream. Args: time_s (float): seek time in seconds keyframes_only (bool): allow to seek only to keyframes .. note:: Current implementation is the so-called precise seek. This means following seek, call to :mod:`next()` will return the frame with the exact timestamp if it exists or the first frame with timestamp larger than ``time_s``. """ self._c.seek(time_s, keyframes_only) return self def get_metadata(self) -> Dict[str, Any]: """Returns video metadata Returns: (dict): dictionary containing duration and frame rate for every stream """ return self._c.get_metadata() def set_current_stream(self, stream: str) -> bool: """Set current stream. Explicitly define the stream we are operating on. Args: stream (string): descriptor of the required stream. Defaults to ``"video:0"`` Currently available stream types include ``['video', 'audio']``. Each descriptor consists of two parts: stream type (e.g. 'video') and a unique stream id (which are determined by video encoding). In this way, if the video contaner contains multiple streams of the same type, users can acces the one they want. If only stream type is passed, the decoder auto-detects first stream of that type and returns it. Returns: (bool): True on succes, False otherwise """ if self.is_cuda: print("GPU decoding only works with video stream.") return self._c.set_current_stream(stream) msadapter/ms_adapter/torchvision/ops/000077500000000000000000000000001436713554100203365ustar00rootroot00000000000000msadapter/ms_adapter/torchvision/ops/__init__.py000066400000000000000000000035501436713554100224520ustar00rootroot00000000000000from .boxes import ( nms, batched_nms, remove_small_boxes, clip_boxes_to_image, box_area, box_convert, box_iou, generalized_box_iou, distance_box_iou, complete_box_iou, masks_to_boxes, ) from .ciou_loss import complete_box_iou_loss from .deform_conv import deform_conv2d, DeformConv2d from .diou_loss import distance_box_iou_loss from .drop_block import drop_block2d, DropBlock2d, drop_block3d, DropBlock3d from .feature_pyramid_network import FeaturePyramidNetwork from .focal_loss import sigmoid_focal_loss from .giou_loss import generalized_box_iou_loss from .misc import FrozenBatchNorm2d, Conv2dNormActivation, Conv3dNormActivation, SqueezeExcitation, MLP, Permute from .poolers import MultiScaleRoIAlign # from .ps_roi_align import ps_roi_align, PSRoIAlign # from .ps_roi_pool import ps_roi_pool, PSRoIPool from .roi_align import roi_align, RoIAlign # from .roi_pool import roi_pool, RoIPool from .stochastic_depth import stochastic_depth, StochasticDepth # _register_custom_op() __all__ = [ "masks_to_boxes", "deform_conv2d", "DeformConv2d", "nms", "batched_nms", "remove_small_boxes", "clip_boxes_to_image", "box_convert", "box_area", "box_iou", "generalized_box_iou", "distance_box_iou", "complete_box_iou", "roi_align", "RoIAlign", # "roi_pool", # "RoIPool", # "ps_roi_align", # "PSRoIAlign", # "ps_roi_pool", # "PSRoIPool", "MultiScaleRoIAlign", "FeaturePyramidNetwork", "sigmoid_focal_loss", "stochastic_depth", "StochasticDepth", "FrozenBatchNorm2d", "Conv2dNormActivation", "Conv3dNormActivation", "SqueezeExcitation", "MLP", "Permute", "generalized_box_iou_loss", "distance_box_iou_loss", "complete_box_iou_loss", "drop_block2d", "DropBlock2d", "drop_block3d", "DropBlock3d", ] msadapter/ms_adapter/torchvision/ops/_box_convert.py000066400000000000000000000046121436713554100234020ustar00rootroot00000000000000import ms_adapter.pytorch as torch from ms_adapter.pytorch import Tensor def _box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor: """ Converts bounding boxes from (cx, cy, w, h) format to (x1, y1, x2, y2) format. (cx, cy) refers to center of bounding box (w, h) are width and height of bounding box Args: boxes (Tensor[N, 4]): boxes in (cx, cy, w, h) format which will be converted. Returns: boxes (Tensor(N, 4)): boxes in (x1, y1, x2, y2) format. """ # We need to change all 4 of them so some temporary variable is needed. cx, cy, w, h = boxes.unbind(-1) x1 = cx - 0.5 * w y1 = cy - 0.5 * h x2 = cx + 0.5 * w y2 = cy + 0.5 * h boxes = torch.stack((x1, y1, x2, y2), dim=-1) return boxes def _box_xyxy_to_cxcywh(boxes: Tensor) -> Tensor: """ Converts bounding boxes from (x1, y1, x2, y2) format to (cx, cy, w, h) format. (x1, y1) refer to top left of bounding box (x2, y2) refer to bottom right of bounding box Args: boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format which will be converted. Returns: boxes (Tensor(N, 4)): boxes in (cx, cy, w, h) format. """ x1, y1, x2, y2 = boxes.unbind(-1) cx = (x1 + x2) / 2 cy = (y1 + y2) / 2 w = x2 - x1 h = y2 - y1 boxes = torch.stack((cx, cy, w, h), dim=-1) return boxes def _box_xywh_to_xyxy(boxes: Tensor) -> Tensor: """ Converts bounding boxes from (x, y, w, h) format to (x1, y1, x2, y2) format. (x, y) refers to top left of bouding box. (w, h) refers to width and height of box. Args: boxes (Tensor[N, 4]): boxes in (x, y, w, h) which will be converted. Returns: boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format. """ x, y, w, h = boxes.unbind(-1) boxes = torch.stack([x, y, x + w, y + h], dim=-1) return boxes def _box_xyxy_to_xywh(boxes: Tensor) -> Tensor: """ Converts bounding boxes from (x1, y1, x2, y2) format to (x, y, w, h) format. (x1, y1) refer to top left of bounding box (x2, y2) refer to bottom right of bounding box Args: boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) which will be converted. Returns: boxes (Tensor[N, 4]): boxes in (x, y, w, h) format. """ x1, y1, x2, y2 = boxes.unbind(-1) w = x2 - x1 # x2 - x1 h = y2 - y1 # y2 - y1 boxes = torch.stack((x1, y1, w, h), dim=-1) return boxes msadapter/ms_adapter/torchvision/ops/_utils.py000066400000000000000000000074561436713554100222230ustar00rootroot00000000000000from typing import List, Optional, Tuple, Union import mindspore as ms import ms_adapter.pytorch as torch from ms_adapter.pytorch import nn, Tensor from ms_adapter.pytorch.tensor import cast_to_adapter_tensor def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor: """ Efficient version of torch.cat that avoids a copy if there is only a single element in a list """ # TODO add back the assert # assert isinstance(tensors, (list, tuple)) if len(tensors) == 1: return tensors[0] return torch.cat(tensors, dim) def convert_boxes_to_roi_format(boxes: List[Tensor]) -> Tensor: concat_boxes = _cat([b for b in boxes], dim=0) temp = [] for i, b in enumerate(boxes): temp.append(torch.full_like(b[:, :1], i)) ids = _cat(temp, dim=0) rois = torch.cat([ids, concat_boxes], dim=1) return rois def check_roi_boxes_shape(boxes: Union[Tensor, List[Tensor]]): if isinstance(boxes, (list, tuple)): for _tensor in boxes: torch._assert( _tensor.size(1) == 4, "The shape of the tensor in the boxes list is not correct as List[Tensor[L, 4]]" ) elif isinstance(boxes, torch.Tensor): torch._assert(boxes.size(1) == 5, "The boxes tensor shape is not correct as Tensor[K, 5]") else: torch._assert(False, "boxes is expected to be a Tensor[L, 5] or a List[Tensor[K, 4]]") return def split_normalization_params( model: nn.Module, norm_classes: Optional[List[type]] = None ) -> Tuple[List[Tensor], List[Tensor]]: # Adapted from https://github.com/facebookresearch/ClassyVision/blob/659d7f78/classy_vision/generic/util.py#L501 if not norm_classes: norm_classes = [ nn.modules.batchnorm._BatchNorm, nn.LayerNorm, nn.GroupNorm, # TODO: nn.modules.instancenorm._InstanceNorm, nn.LocalResponseNorm, ] for t in norm_classes: if not issubclass(t, nn.Module): raise ValueError(f"Class {t} is not a subclass of nn.Module.") classes = tuple(norm_classes) norm_params = [] other_params = [] for module in model.modules(): if next(module.children(), None): other_params.extend(p for p in module.parameters(recurse=False) if p.requires_grad) elif isinstance(module, classes): norm_params.extend(p for p in module.parameters() if p.requires_grad) else: other_params.extend(p for p in module.parameters() if p.requires_grad) return norm_params, other_params def _upcast(t: Tensor) -> Tensor: # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type if t.is_floating_point(): return t if t.dtype in (torch.float32, torch.float64) else t.float() else: return t if t.dtype in (torch.int32, torch.int64) else t.int() def _upcast_non_float(t: Tensor) -> Tensor: # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type if t.dtype not in (torch.float32, torch.float64): return t.float() return t def _loss_inter_union( boxes1: torch.Tensor, boxes2: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: x1, y1, x2, y2 = boxes1.unbind(dim=-1) x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) # Intersection keypoints xkis1 = torch.maximum(x1, x1g) ykis1 = torch.maximum(y1, y1g) xkis2 = torch.minimum(x2, x2g) ykis2 = torch.minimum(y2, y2g) intsctk = torch.zeros_like(x1) mask = (ykis2 > ykis1) & (xkis2 > xkis1) if mask.to(torch.int8).sum() != 0: out = (xkis2[mask] - xkis1[mask]) * (ykis2[mask] - ykis1[mask]) intsctk = cast_to_adapter_tensor(ms.numpy.where(mask, x=out, y=intsctk)) unionk = (x2 - x1) * (y2 - y1) + (x2g - x1g) * (y2g - y1g) - intsctk return intsctk, unionk msadapter/ms_adapter/torchvision/ops/boxes.py000066400000000000000000000343601436713554100220360ustar00rootroot00000000000000from typing import Tuple import mindspore as ms from mindspore.ops._primitive_cache import _get_cache_prim import ms_adapter.pytorch as torch from ms_adapter.pytorch import Tensor from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from ._box_convert import _box_cxcywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xywh_to_xyxy, _box_xyxy_to_xywh from ._utils import _upcast def nms(boxes, scores, iou_threshold): boxes = cast_to_ms_tensor(boxes) scores_keep = cast_to_ms_tensor(scores) ordered_scores, order = _get_cache_prim(ms.ops.Sort)(descending=True)(scores_keep) ordered_boxes = boxes[order, :] bboxes = ms.ops.concat((ordered_boxes, ordered_scores.reshape((-1, 1))), -1) _, output_idx, selected_mask = _get_cache_prim(ms.ops.NMSWithMask)(iou_threshold)(bboxes) output_idx = cast_to_adapter_tensor(output_idx) return output_idx[selected_mask] def batched_nms( boxes: Tensor, scores: Tensor, idxs: Tensor, iou_threshold: float, ) -> Tensor: """ Performs non-maximum suppression in a batched fashion. Each index value correspond to a category, and NMS will not be applied between elements of different categories. Args: boxes (Tensor[N, 4]): boxes where NMS will be performed. They are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. scores (Tensor[N]): scores for each one of the boxes idxs (Tensor[N]): indices of the categories for each one of the boxes. iou_threshold (float): discards all overlapping boxes with IoU > iou_threshold Returns: Tensor: int64 tensor with the indices of the elements that have been kept by NMS, sorted in decreasing order of scores """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(batched_nms) # Benchmarks that drove the following thresholds are at # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339 # if boxes.numel() > (4000 if boxes.device.type == "cpu" else 20000) and not torchvision._is_tracing(): if boxes.numel() > 4000: return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold) else: return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold) # @torch.jit._script_if_tracing def _batched_nms_coordinate_trick( boxes: Tensor, scores: Tensor, idxs: Tensor, iou_threshold: float, ) -> Tensor: # strategy: in order to perform NMS independently per class, # we add an offset to all the boxes. The offset is dependent # only on the class idx, and is large enough so that boxes # from different classes do not overlap if boxes.numel() == 0: return torch.empty((0,), dtype=torch.int64, device=boxes.device) max_coordinate = boxes.max() offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes)) boxes_for_nms = boxes + offsets[:, None] keep = nms(boxes_for_nms, scores, iou_threshold) return keep # @torch.jit._script_if_tracing def _batched_nms_vanilla( boxes: Tensor, scores: Tensor, idxs: Tensor, iou_threshold: float, ) -> Tensor: # Based on Detectron2 implementation, just manually call nms() on each class independently keep_mask = torch.zeros_like(scores, dtype=torch.bool) for class_id in torch.unique(idxs): curr_indices = torch.where(idxs == class_id)[0] curr_keep_indices = nms(boxes[curr_indices], scores[curr_indices], iou_threshold) keep_mask[curr_indices[curr_keep_indices]] = True keep_indices = torch.where(keep_mask)[0] return keep_indices[scores[keep_indices].sort(descending=True)[1]] def remove_small_boxes(boxes: Tensor, min_size: float) -> Tensor: """ Remove boxes which contains at least one side smaller than min_size. Args: boxes (Tensor[N, 4]): boxes in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. min_size (float): minimum size Returns: Tensor[K]: indices of the boxes that have both sides larger than min_size """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(remove_small_boxes) ws, hs = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1] keep = (ws >= min_size) & (hs >= min_size) keep = torch.where(keep)[0] return keep def clip_boxes_to_image(boxes: Tensor, size: Tuple[int, int]) -> Tensor: """ Clip boxes so that they lie inside an image of size `size`. Args: boxes (Tensor[N, 4]): boxes in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. size (Tuple[height, width]): size of the image Returns: Tensor[N, 4]: clipped boxes """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(clip_boxes_to_image) dim = boxes.dim() boxes_x = boxes[..., 0::2] boxes_y = boxes[..., 1::2] height, width = size boxes_x = boxes_x.clamp(min=0, max=width) boxes_y = boxes_y.clamp(min=0, max=height) clipped_boxes = torch.stack((boxes_x, boxes_y), dim=dim) return clipped_boxes.reshape(boxes.shape) def box_convert(boxes: Tensor, in_fmt: str, out_fmt: str) -> Tensor: """ Converts boxes from given in_fmt to out_fmt. Supported in_fmt and out_fmt are: 'xyxy': boxes are represented via corners, x1, y1 being top left and x2, y2 being bottom right. This is the format that torchvision utilities expect. 'xywh' : boxes are represented via corner, width and height, x1, y2 being top left, w, h being width and height. 'cxcywh' : boxes are represented via centre, width and height, cx, cy being center of box, w, h being width and height. Args: boxes (Tensor[N, 4]): boxes which will be converted. in_fmt (str): Input format of given boxes. Supported formats are ['xyxy', 'xywh', 'cxcywh']. out_fmt (str): Output format of given boxes. Supported formats are ['xyxy', 'xywh', 'cxcywh'] Returns: Tensor[N, 4]: Boxes into converted format. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(box_convert) allowed_fmts = ("xyxy", "xywh", "cxcywh") if in_fmt not in allowed_fmts or out_fmt not in allowed_fmts: raise ValueError("Unsupported Bounding Box Conversions for given in_fmt and out_fmt") if in_fmt == out_fmt: return boxes.clone() if in_fmt != "xyxy" and out_fmt != "xyxy": # convert to xyxy and change in_fmt xyxy if in_fmt == "xywh": boxes = _box_xywh_to_xyxy(boxes) elif in_fmt == "cxcywh": boxes = _box_cxcywh_to_xyxy(boxes) in_fmt = "xyxy" if in_fmt == "xyxy": if out_fmt == "xywh": boxes = _box_xyxy_to_xywh(boxes) elif out_fmt == "cxcywh": boxes = _box_xyxy_to_cxcywh(boxes) elif out_fmt == "xyxy": if in_fmt == "xywh": boxes = _box_xywh_to_xyxy(boxes) elif in_fmt == "cxcywh": boxes = _box_cxcywh_to_xyxy(boxes) return boxes def box_area(boxes: Tensor) -> Tensor: """ Computes the area of a set of bounding boxes, which are specified by their (x1, y1, x2, y2) coordinates. Args: boxes (Tensor[N, 4]): boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. Returns: Tensor[N]: the area for each box """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(box_area) boxes = _upcast(boxes) return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) # implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py # with slight modifications def _box_inter_union(boxes1: Tensor, boxes2: Tensor) -> Tuple[Tensor, Tensor]: area1 = box_area(boxes1) area2 = box_area(boxes2) lt = torch.maximum(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] rb = torch.minimum(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] wh = _upcast(rb - lt).clamp(min=0) # [N,M,2] inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] union = area1[:, None] + area2 - inter return inter, union def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor: """ Return intersection-over-union (Jaccard index) between two sets of boxes. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. Args: boxes1 (Tensor[N, 4]): first set of boxes boxes2 (Tensor[M, 4]): second set of boxes Returns: Tensor[N, M]: the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2 """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(box_iou) inter, union = _box_inter_union(boxes1, boxes2) iou = inter / union return iou # Implementation adapted from https://github.com/facebookresearch/detr/blob/master/util/box_ops.py def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor: """ Return generalized intersection-over-union (Jaccard index) between two sets of boxes. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. Args: boxes1 (Tensor[N, 4]): first set of boxes boxes2 (Tensor[M, 4]): second set of boxes Returns: Tensor[N, M]: the NxM matrix containing the pairwise generalized IoU values for every element in boxes1 and boxes2 """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(generalized_box_iou) inter, union = _box_inter_union(boxes1, boxes2) iou = inter / union lti = torch.minimum(boxes1[:, None, :2], boxes2[:, :2]) rbi = torch.maximum(boxes1[:, None, 2:], boxes2[:, 2:]) whi = _upcast(rbi - lti).clamp(min=0) # [N,M,2] areai = whi[:, :, 0] * whi[:, :, 1] return iou - (areai - union) / areai def complete_box_iou(boxes1: Tensor, boxes2: Tensor, eps: float = 1e-7) -> Tensor: """ Return complete intersection-over-union (Jaccard index) between two sets of boxes. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. Args: boxes1 (Tensor[N, 4]): first set of boxes boxes2 (Tensor[M, 4]): second set of boxes eps (float, optional): small number to prevent division by zero. Default: 1e-7 Returns: Tensor[N, M]: the NxM matrix containing the pairwise complete IoU values for every element in boxes1 and boxes2 """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(complete_box_iou) boxes1 = _upcast(boxes1) boxes2 = _upcast(boxes2) diou, iou = _box_diou_iou(boxes1, boxes2, eps) w_pred = boxes1[:, None, 2] - boxes1[:, None, 0] h_pred = boxes1[:, None, 3] - boxes1[:, None, 1] w_gt = boxes2[:, 2] - boxes2[:, 0] h_gt = boxes2[:, 3] - boxes2[:, 1] v = (4 / (torch.pi ** 2)) * torch.pow(torch.atan(w_pred / h_pred) - torch.atan(w_gt / h_gt), 2) # TODO: # with torch.no_grad(): # alpha = v / (1 - iou + v + eps) alpha = v / (1 - iou + v + eps) alpha = cast_to_adapter_tensor(ms.ops.stop_gradient(alpha)) return diou - alpha * v def distance_box_iou(boxes1: Tensor, boxes2: Tensor, eps: float = 1e-7) -> Tensor: """ Return distance intersection-over-union (Jaccard index) between two sets of boxes. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. Args: boxes1 (Tensor[N, 4]): first set of boxes boxes2 (Tensor[M, 4]): second set of boxes eps (float, optional): small number to prevent division by zero. Default: 1e-7 Returns: Tensor[N, M]: the NxM matrix containing the pairwise distance IoU values for every element in boxes1 and boxes2 """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(distance_box_iou) boxes1 = _upcast(boxes1) boxes2 = _upcast(boxes2) diou, _ = _box_diou_iou(boxes1, boxes2, eps=eps) return diou def _box_diou_iou(boxes1: Tensor, boxes2: Tensor, eps: float = 1e-7) -> Tuple[Tensor, Tensor]: iou = box_iou(boxes1, boxes2) lti = torch.minimum(boxes1[:, None, :2], boxes2[:, :2]) rbi = torch.maximum(boxes1[:, None, 2:], boxes2[:, 2:]) whi = _upcast(rbi - lti).clamp(min=0) # [N,M,2] diagonal_distance_squared = (whi[:, :, 0] ** 2) + (whi[:, :, 1] ** 2) + eps # centers of boxes x_p = (boxes1[:, 0] + boxes1[:, 2]) / 2 y_p = (boxes1[:, 1] + boxes1[:, 3]) / 2 x_g = (boxes2[:, 0] + boxes2[:, 2]) / 2 y_g = (boxes2[:, 1] + boxes2[:, 3]) / 2 # The distance between boxes' centers squared. centers_distance_squared = (_upcast((x_p[:, None] - x_g[None, :])) ** 2) + ( _upcast((y_p[:, None] - y_g[None, :])) ** 2 ) # The distance IoU is the IoU penalized by a normalized # distance between boxes' centers squared. return iou - (centers_distance_squared / diagonal_distance_squared), iou def masks_to_boxes(masks: torch.Tensor) -> torch.Tensor: """ Compute the bounding boxes around the provided masks. Returns a [N, 4] tensor containing bounding boxes. The boxes are in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``. Args: masks (Tensor[N, H, W]): masks to transform where N is the number of masks and (H, W) are the spatial dimensions. Returns: Tensor[N, 4]: bounding boxes """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(masks_to_boxes) if masks.numel() == 0: return torch.zeros((0, 4), device=masks.device, dtype=torch.float) n = masks.shape[0] bounding_boxes = torch.zeros((n, 4), device=masks.device, dtype=torch.float) for index, mask in enumerate(masks): y, x = torch.where(mask != 0) bounding_boxes[index, 0] = torch.min(x) bounding_boxes[index, 1] = torch.min(y) bounding_boxes[index, 2] = torch.max(x) bounding_boxes[index, 3] = torch.max(y) return bounding_boxes msadapter/ms_adapter/torchvision/ops/ciou_loss.py000066400000000000000000000051741436713554100227160ustar00rootroot00000000000000import ms_adapter.pytorch as torch import mindspore as ms from ms_adapter.pytorch.tensor import cast_to_adapter_tensor from ._utils import _upcast_non_float from .diou_loss import _diou_iou_loss def complete_box_iou_loss( boxes1: torch.Tensor, boxes2: torch.Tensor, reduction: str = "none", eps: float = 1e-7, ) -> torch.Tensor: """ Gradient-friendly IoU loss with an additional penalty that is non-zero when the boxes do not overlap. This loss function considers important geometrical factors such as overlap area, normalized central point distance and aspect ratio. This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the same dimensions. Args: boxes1 : (Tensor[N, 4] or Tensor[4]) first set of boxes boxes2 : (Tensor[N, 4] or Tensor[4]) second set of boxes reduction : (string, optional) Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be applied to the output. ``'mean'``: The output will be averaged. ``'sum'``: The output will be summed. Default: ``'none'`` eps : (float): small number to prevent division by zero. Default: 1e-7 Returns: Tensor: Loss tensor with the reduction option applied. Reference: Zhaohui Zheng et al.: Complete Intersection over Union Loss: https://arxiv.org/abs/1911.08287 """ # Original Implementation from https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/losses.py # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(complete_box_iou_loss) boxes1 = _upcast_non_float(boxes1) boxes2 = _upcast_non_float(boxes2) diou_loss, iou = _diou_iou_loss(boxes1, boxes2) x1, y1, x2, y2 = boxes1.unbind(dim=-1) x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) # width and height of boxes w_pred = x2 - x1 h_pred = y2 - y1 w_gt = x2g - x1g h_gt = y2g - y1g v = (4 / (torch.pi ** 2)) * torch.pow((torch.atan(w_gt / h_gt) - torch.atan(w_pred / h_pred)), 2) # TODO: # with torch.no_grad(): # alpha = v / (1 - iou + v + eps) alpha = v / (1 - iou + v + eps) alpha = cast_to_adapter_tensor(ms.ops.stop_gradient(alpha)) loss = diou_loss + alpha * v if reduction == "mean": loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() elif reduction == "sum": loss = loss.sum() return loss msadapter/ms_adapter/torchvision/ops/deform_conv.py000066400000000000000000000162551436713554100232220ustar00rootroot00000000000000import math from typing import Optional, Tuple import mindspore as ms import ms_adapter.pytorch as torch from ms_adapter.pytorch import nn, Tensor from ms_adapter.pytorch.nn import init from ms_adapter.pytorch.nn.modules.utils import _pair from ms_adapter.pytorch.nn.parameter import Parameter from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor def deform_conv2d( input: Tensor, offset: Tensor, weight: Tensor, bias: Optional[Tensor] = None, stride: Tuple[int, int] = (1, 1), padding: Tuple[int, int] = (0, 0), dilation: Tuple[int, int] = (1, 1), mask: Optional[Tensor] = None, ) -> Tensor: r""" Performs Deformable Convolution v2, described in `Deformable ConvNets v2: More Deformable, Better Results `__ if :attr:`mask` is not ``None`` and Performs Deformable Convolution, described in `Deformable Convolutional Networks `__ if :attr:`mask` is ``None``. Args: input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor offset (Tensor[batch_size, 2 * offset_groups * kernel_height * kernel_width, out_height, out_width]): offsets to be applied for each position in the convolution kernel. weight (Tensor[out_channels, in_channels // groups, kernel_height, kernel_width]): convolution weights, split into groups of size (in_channels // groups) bias (Tensor[out_channels]): optional bias of shape (out_channels,). Default: None stride (int or Tuple[int, int]): distance between convolution centers. Default: 1 padding (int or Tuple[int, int]): height/width of padding of zeroes around each image. Default: 0 dilation (int or Tuple[int, int]): the spacing between kernel elements. Default: 1 mask (Tensor[batch_size, offset_groups * kernel_height * kernel_width, out_height, out_width]): masks to be applied for each position in the convolution kernel. Default: None Returns: Tensor[batch_sz, out_channels, out_h, out_w]: result of convolution Examples:: >>> input = torch.rand(4, 3, 10, 10) >>> kh, kw = 3, 3 >>> weight = torch.rand(5, 3, kh, kw) >>> # offset and mask should have the same spatial size as the output >>> # of the convolution. In this case, for an input of 10, stride of 1 >>> # and kernel size of 3, without padding, the output size is 8 >>> offset = torch.rand(4, 2 * kh * kw, 8, 8) >>> mask = torch.rand(4, kh * kw, 8, 8) >>> out = deform_conv2d(input, offset, weight, mask=mask) >>> print(out.shape) >>> # returns >>> torch.Size([4, 5, 8, 8]) """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(deform_conv2d) # _assert_has_ops() out_channels = weight.shape[0] # use_mask = mask is not None if mask is None: # mask = torch.zeros((input.shape[0], 0), device=input.device, dtype=input.dtype) raise NotImplementedError("Only Deformable ConvNets v2 supported.") # if bias is None: # bias = torch.zeros(out_channels, device=input.device, dtype=input.dtype) stride_h, stride_w = _pair(stride) pad_h, pad_w = _pair(padding) dil_h, dil_w = _pair(dilation) weights_h, weights_w = weight.shape[-2:] _, n_in_channels, _, _ = input.shape n_offset_grps = offset.shape[1] // (2 * weights_h * weights_w) n_weight_grps = n_in_channels // weight.shape[1] if n_offset_grps == 0: raise RuntimeError( "the shape of the offset tensor at dimension 1 is not valid. It should " "be a multiple of 2 * weight.size[2] * weight.size[3].\n" f"Got offset.shape[1]={offset.shape[1]}, while 2 * weight.size[2] * weight.size[3]={2 * weights_h * weights_w}" ) input = cast_to_ms_tensor(input) offset = cast_to_ms_tensor(offset) weight = cast_to_ms_tensor(weight) if bias is not None: bias = cast_to_ms_tensor(bias) output = ms.ops.deformable_conv2d(input, weight, offset, (weights_h, weights_w), (1, 1, stride_h, stride_w), (pad_h, pad_h, pad_w, pad_w), bias, (1, 1, dil_h, dil_w), n_weight_grps, n_offset_grps) return cast_to_adapter_tensor(output) class DeformConv2d(nn.Module): """ See :func:`deform_conv2d`. """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, ): super().__init__() # _log_api_usage_once(self) if in_channels % groups != 0: raise ValueError("in_channels must be divisible by groups") if out_channels % groups != 0: raise ValueError("out_channels must be divisible by groups") self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = _pair(kernel_size) self.stride = _pair(stride) self.padding = _pair(padding) self.dilation = _pair(dilation) self.groups = groups self.weight = Parameter( torch.empty(out_channels, in_channels // groups, self.kernel_size[0], self.kernel_size[1]) ) if bias: self.bias = Parameter(torch.empty(out_channels)) else: self.register_parameter("bias", None) self.reset_parameters() def reset_parameters(self) -> None: init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) init.uniform_(self.bias, -bound, bound) def forward(self, input: Tensor, offset: Tensor, mask: Optional[Tensor] = None) -> Tensor: """ Args: input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor offset (Tensor[batch_size, 2 * offset_groups * kernel_height * kernel_width, out_height, out_width]): offsets to be applied for each position in the convolution kernel. mask (Tensor[batch_size, offset_groups * kernel_height * kernel_width, out_height, out_width]): masks to be applied for each position in the convolution kernel. """ return deform_conv2d( input, offset, self.weight, self.bias, stride=self.stride, padding=self.padding, dilation=self.dilation, mask=mask, ) def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(" f"{self.in_channels}" f", {self.out_channels}" f", kernel_size={self.kernel_size}" f", stride={self.stride}" ) s += f", padding={self.padding}" if self.padding != (0, 0) else "" s += f", dilation={self.dilation}" if self.dilation != (1, 1) else "" s += f", groups={self.groups}" if self.groups != 1 else "" s += ", bias=False" if self.bias is None else "" s += ")" return s msadapter/ms_adapter/torchvision/ops/diou_loss.py000066400000000000000000000060371436713554100227160ustar00rootroot00000000000000from typing import Tuple import ms_adapter.pytorch as torch from ._utils import _loss_inter_union, _upcast_non_float def distance_box_iou_loss( boxes1: torch.Tensor, boxes2: torch.Tensor, reduction: str = "none", eps: float = 1e-7, ) -> torch.Tensor: """ Gradient-friendly IoU loss with an additional penalty that is non-zero when the distance between boxes' centers isn't zero. Indeed, for two exactly overlapping boxes, the distance IoU is the same as the IoU loss. This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the same dimensions. Args: boxes1 (Tensor[N, 4]): first set of boxes boxes2 (Tensor[N, 4]): second set of boxes reduction (string, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be applied to the output. ``'mean'``: The output will be averaged. ``'sum'``: The output will be summed. Default: ``'none'`` eps (float, optional): small number to prevent division by zero. Default: 1e-7 Returns: Tensor: Loss tensor with the reduction option applied. Reference: Zhaohui Zheng et. al: Distance Intersection over Union Loss: https://arxiv.org/abs/1911.08287 """ # Original Implementation from https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/losses.py # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(distance_box_iou_loss) boxes1 = _upcast_non_float(boxes1) boxes2 = _upcast_non_float(boxes2) loss, _ = _diou_iou_loss(boxes1, boxes2, eps) if reduction == "mean": loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() elif reduction == "sum": loss = loss.sum() return loss def _diou_iou_loss( boxes1: torch.Tensor, boxes2: torch.Tensor, eps: float = 1e-7, ) -> Tuple[torch.Tensor, torch.Tensor]: intsct, union = _loss_inter_union(boxes1, boxes2) iou = intsct / (union + eps) # smallest enclosing box x1, y1, x2, y2 = boxes1.unbind(dim=-1) x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) xc1 = torch.minimum(x1, x1g) yc1 = torch.minimum(y1, y1g) xc2 = torch.maximum(x2, x2g) yc2 = torch.maximum(y2, y2g) # The diagonal distance of the smallest enclosing box squared diagonal_distance_squared = ((xc2 - xc1) ** 2) + ((yc2 - yc1) ** 2) + eps # centers of boxes x_p = (x2 + x1) / 2 y_p = (y2 + y1) / 2 x_g = (x1g + x2g) / 2 y_g = (y1g + y2g) / 2 # The distance between boxes' centers squared. centers_distance_squared = ((x_p - x_g) ** 2) + ((y_p - y_g) ** 2) # The distance IoU is the IoU penalized by a normalized # distance between boxes' centers squared. loss = 1 - iou + (centers_distance_squared / diagonal_distance_squared) return loss, iou msadapter/ms_adapter/torchvision/ops/drop_block.py000066400000000000000000000133461436713554100230350ustar00rootroot00000000000000import ms_adapter.pytorch as torch import ms_adapter.pytorch.nn.functional as F from ms_adapter.pytorch import nn, Tensor def drop_block2d( input: Tensor, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06, training: bool = True ) -> Tensor: """ Implements DropBlock2d from `"DropBlock: A regularization method for convolutional networks" `. Args: input (Tensor[N, C, H, W]): The input tensor or 4-dimensions with the first one being its batch i.e. a batch with ``N`` rows. p (float): Probability of an element to be dropped. block_size (int): Size of the block to drop. inplace (bool): If set to ``True``, will do this operation in-place. Default: ``False``. eps (float): A value added to the denominator for numerical stability. Default: 1e-6. training (bool): apply dropblock if is ``True``. Default: ``True``. Returns: Tensor[N, C, H, W]: The randomly zeroed tensor after dropblock. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(drop_block2d) if p < 0.0 or p > 1.0: raise ValueError(f"drop probability has to be between 0 and 1, but got {p}.") if input.ndim != 4: raise ValueError(f"input should be 4 dimensional. Got {input.ndim} dimensions.") if not training or p == 0.0: return input N, C, H, W = input.size() block_size = min(block_size, W, H) # compute the gamma of Bernoulli distribution gamma = (p * H * W) / ((block_size ** 2) * ((H - block_size + 1) * (W - block_size + 1))) noise = torch.empty((N, C, H - block_size + 1, W - block_size + 1), dtype=input.dtype, device=input.device) noise.bernoulli_(gamma) noise = F.pad(noise, [block_size // 2] * 4, value=0) noise = F.max_pool2d(noise, stride=(1, 1), kernel_size=(block_size, block_size), padding=block_size // 2) noise = 1 - noise normalize_scale = noise.numel() / (eps + noise.sum()) if inplace: input.mul_(noise).mul_(normalize_scale) else: input = input * noise * normalize_scale return input def drop_block3d( input: Tensor, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06, training: bool = True ) -> Tensor: """ Implements DropBlock3d from `"DropBlock: A regularization method for convolutional networks" `. Args: input (Tensor[N, C, D, H, W]): The input tensor or 5-dimensions with the first one being its batch i.e. a batch with ``N`` rows. p (float): Probability of an element to be dropped. block_size (int): Size of the block to drop. inplace (bool): If set to ``True``, will do this operation in-place. Default: ``False``. eps (float): A value added to the denominator for numerical stability. Default: 1e-6. training (bool): apply dropblock if is ``True``. Default: ``True``. Returns: Tensor[N, C, D, H, W]: The randomly zeroed tensor after dropblock. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(drop_block3d) if p < 0.0 or p > 1.0: raise ValueError(f"drop probability has to be between 0 and 1, but got {p}.") if input.ndim != 5: raise ValueError(f"input should be 5 dimensional. Got {input.ndim} dimensions.") if not training or p == 0.0: return input N, C, D, H, W = input.size() block_size = min(block_size, D, H, W) # compute the gamma of Bernoulli distribution gamma = (p * D * H * W) / ((block_size ** 3) * ((D - block_size + 1) * (H - block_size + 1) * (W - block_size + 1))) noise = torch.empty( (N, C, D - block_size + 1, H - block_size + 1, W - block_size + 1), dtype=input.dtype, device=input.device ) noise.bernoulli_(gamma) noise = F.pad(noise, [block_size // 2] * 6, value=0) noise = F.max_pool3d( noise, stride=(1, 1, 1), kernel_size=(block_size, block_size, block_size), padding=block_size // 2 ) noise = 1 - noise normalize_scale = noise.numel() / (eps + noise.sum()) if inplace: input.mul_(noise).mul_(normalize_scale) else: input = input * noise * normalize_scale return input # torch.fx.wrap("drop_block2d") class DropBlock2d(nn.Module): """ See :func:`drop_block2d`. """ def __init__(self, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06) -> None: super().__init__() self.p = p self.block_size = block_size self.inplace = inplace self.eps = eps def forward(self, input: Tensor) -> Tensor: """ Args: input (Tensor): Input feature map on which some areas will be randomly dropped. Returns: Tensor: The tensor after DropBlock layer. """ return drop_block2d(input, self.p, self.block_size, self.inplace, self.eps, self.training) def __repr__(self) -> str: s = f"{self.__class__.__name__}(p={self.p}, block_size={self.block_size}, inplace={self.inplace})" return s # torch.fx.wrap("drop_block3d") class DropBlock3d(DropBlock2d): """ See :func:`drop_block3d`. """ def __init__(self, p: float, block_size: int, inplace: bool = False, eps: float = 1e-06) -> None: super().__init__(p, block_size, inplace, eps) def forward(self, input: Tensor) -> Tensor: """ Args: input (Tensor): Input feature map on which some areas will be randomly dropped. Returns: Tensor: The tensor after DropBlock layer. """ return drop_block3d(input, self.p, self.block_size, self.inplace, self.eps, self.training) msadapter/ms_adapter/torchvision/ops/feature_pyramid_network.py000066400000000000000000000166411436713554100256510ustar00rootroot00000000000000from collections import OrderedDict from typing import Tuple, List, Dict, Callable, Optional import ms_adapter.pytorch.nn.functional as F from ms_adapter.pytorch import nn, Tensor from ..ops.misc import Conv2dNormActivation class ExtraFPNBlock(nn.Module): """ Base class for the extra block in the FPN. Args: results (List[Tensor]): the result of the FPN x (List[Tensor]): the original feature maps names (List[str]): the names for each one of the original feature maps Returns: results (List[Tensor]): the extended set of results of the FPN names (List[str]): the extended set of names for the results """ def forward( self, results: List[Tensor], x: List[Tensor], names: List[str], ) -> Tuple[List[Tensor], List[str]]: pass class FeaturePyramidNetwork(nn.Module): """ Module that adds a FPN from on top of a set of feature maps. This is based on `"Feature Pyramid Network for Object Detection" `_. The feature maps are currently supposed to be in increasing depth order. The input to the model is expected to be an OrderedDict[Tensor], containing the feature maps on top of which the FPN will be added. Args: in_channels_list (list[int]): number of channels for each feature map that is passed to the module out_channels (int): number of channels of the FPN representation extra_blocks (ExtraFPNBlock or None): if provided, extra operations will be performed. It is expected to take the fpn features, the original features and the names of the original features as input, and returns a new list of feature maps and their corresponding names norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None Examples:: >>> m = torchvision.ops.FeaturePyramidNetwork([10, 20, 30], 5) >>> # get some dummy data >>> x = OrderedDict() >>> x['feat0'] = torch.rand(1, 10, 64, 64) >>> x['feat2'] = torch.rand(1, 20, 16, 16) >>> x['feat3'] = torch.rand(1, 30, 8, 8) >>> # compute the FPN on top of x >>> output = m(x) >>> print([(k, v.shape) for k, v in output.items()]) >>> # returns >>> [('feat0', torch.Size([1, 5, 64, 64])), >>> ('feat2', torch.Size([1, 5, 16, 16])), >>> ('feat3', torch.Size([1, 5, 8, 8]))] """ _version = 2 def __init__( self, in_channels_list: List[int], out_channels: int, extra_blocks: Optional[ExtraFPNBlock] = None, norm_layer: Optional[Callable[..., nn.Module]] = None, ): super().__init__() # _log_api_usage_once(self) self.inner_blocks = nn.ModuleList() self.layer_blocks = nn.ModuleList() for in_channels in in_channels_list: if in_channels == 0: raise ValueError("in_channels=0 is currently not supported") inner_block_module = Conv2dNormActivation( in_channels, out_channels, kernel_size=1, padding=0, norm_layer=norm_layer, activation_layer=None ) layer_block_module = Conv2dNormActivation( out_channels, out_channels, kernel_size=3, norm_layer=norm_layer, activation_layer=None ) self.inner_blocks.append(inner_block_module) self.layer_blocks.append(layer_block_module) # initialize parameters now to avoid modifying the initialization of top_blocks for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_uniform_(m.weight, a=1) if m.bias is not None: nn.init.constant_(m.bias, 0) if extra_blocks is not None: if not isinstance(extra_blocks, ExtraFPNBlock): raise TypeError(f"extra_blocks should be of type ExtraFPNBlock not {type(extra_blocks)}") self.extra_blocks = extra_blocks def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor: """ This is equivalent to self.inner_blocks[idx](x), but torchscript doesn't support this yet """ num_blocks = len(self.inner_blocks) if idx < 0: idx += num_blocks out = x for i, module in enumerate(self.inner_blocks): if i == idx: out = module(x) return out def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor: """ This is equivalent to self.layer_blocks[idx](x), but torchscript doesn't support this yet """ num_blocks = len(self.layer_blocks) if idx < 0: idx += num_blocks out = x for i, module in enumerate(self.layer_blocks): if i == idx: out = module(x) return out def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]: """ Computes the FPN for a set of feature maps. Args: x (OrderedDict[Tensor]): feature maps for each feature level. Returns: results (OrderedDict[Tensor]): feature maps after FPN layers. They are ordered from highest resolution first. """ # unpack OrderedDict into two lists for easier handling names = list(x.keys()) x = list(x.values()) last_inner = self.get_result_from_inner_blocks(x[-1], -1) results = [] results.append(self.get_result_from_layer_blocks(last_inner, -1)) for idx in range(len(x) - 2, -1, -1): inner_lateral = self.get_result_from_inner_blocks(x[idx], idx) feat_shape = inner_lateral.shape[-2:] inner_top_down = F.interpolate(last_inner, size=feat_shape, mode="nearest") last_inner = inner_lateral + inner_top_down results.insert(0, self.get_result_from_layer_blocks(last_inner, idx)) if self.extra_blocks is not None: results, names = self.extra_blocks(results, x, names) # make it back an OrderedDict out = OrderedDict([(k, v) for k, v in zip(names, results)]) return out class LastLevelMaxPool(ExtraFPNBlock): """ Applies a max_pool2d on top of the last feature map """ def forward( self, x: List[Tensor], y: List[Tensor], names: List[str], ) -> Tuple[List[Tensor], List[str]]: names.append("pool") x.append(F.max_pool2d(x[-1], 1, 2, 0)) return x, names class LastLevelP6P7(ExtraFPNBlock): """ This module is used in RetinaNet to generate extra layers, P6 and P7. """ def __init__(self, in_channels: int, out_channels: int): super().__init__() self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1) self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1) for module in [self.p6, self.p7]: nn.init.kaiming_uniform_(module.weight, a=1) nn.init.constant_(module.bias, 0) self.use_P5 = in_channels == out_channels def forward( self, p: List[Tensor], c: List[Tensor], names: List[str], ) -> Tuple[List[Tensor], List[str]]: p5, c5 = p[-1], c[-1] x = p5 if self.use_P5 else c5 p6 = self.p6(x) p7 = self.p7(F.relu(p6)) p.extend([p6, p7]) names.extend(["p6", "p7"]) return p, names msadapter/ms_adapter/torchvision/ops/focal_loss.py000066400000000000000000000037161436713554100230430ustar00rootroot00000000000000import ms_adapter.pytorch as torch import ms_adapter.pytorch.nn.functional as F def sigmoid_focal_loss( inputs: torch.Tensor, targets: torch.Tensor, alpha: float = 0.25, gamma: float = 2, reduction: str = "none", ) -> torch.Tensor: """ Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. Args: inputs (Tensor): A float tensor of arbitrary shape. The predictions for each example. targets (Tensor): A float tensor with the same shape as inputs. Stores the binary classification label for each element in inputs (0 for the negative class and 1 for the positive class). alpha (float): Weighting factor in range (0,1) to balance positive vs negative examples or -1 for ignore. Default: ``0.25``. gamma (float): Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples. Default: ``2``. reduction (string): ``'none'`` | ``'mean'`` | ``'sum'`` ``'none'``: No reduction will be applied to the output. ``'mean'``: The output will be averaged. ``'sum'``: The output will be summed. Default: ``'none'``. Returns: Loss tensor with the reduction option applied. """ # Original implementation from https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/focal_loss.py # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(sigmoid_focal_loss) p = torch.sigmoid(inputs) ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") p_t = p * targets + (1 - p) * (1 - targets) loss = ce_loss * ((1 - p_t) ** gamma) if alpha >= 0: alpha_t = alpha * targets + (1 - alpha) * (1 - targets) loss = alpha_t * loss if reduction == "mean": loss = loss.mean() elif reduction == "sum": loss = loss.sum() return loss msadapter/ms_adapter/torchvision/ops/giou_loss.py000066400000000000000000000046061436713554100227210ustar00rootroot00000000000000import ms_adapter.pytorch as torch from ._utils import _upcast_non_float, _loss_inter_union def generalized_box_iou_loss( boxes1: torch.Tensor, boxes2: torch.Tensor, reduction: str = "none", eps: float = 1e-7, ) -> torch.Tensor: """ Gradient-friendly IoU loss with an additional penalty that is non-zero when the boxes do not overlap and scales with the size of their smallest enclosing box. This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the same dimensions. Args: boxes1 (Tensor[N, 4] or Tensor[4]): first set of boxes boxes2 (Tensor[N, 4] or Tensor[4]): second set of boxes reduction (string, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be applied to the output. ``'mean'``: The output will be averaged. ``'sum'``: The output will be summed. Default: ``'none'`` eps (float): small number to prevent division by zero. Default: 1e-7 Returns: Tensor: Loss tensor with the reduction option applied. Reference: Hamid Rezatofighi et. al: Generalized Intersection over Union: A Metric and A Loss for Bounding Box Regression: https://arxiv.org/abs/1902.09630 """ # Original implementation from https://github.com/facebookresearch/fvcore/blob/bfff2ef/fvcore/nn/giou_loss.py # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(generalized_box_iou_loss) boxes1 = _upcast_non_float(boxes1) boxes2 = _upcast_non_float(boxes2) intsctk, unionk = _loss_inter_union(boxes1, boxes2) iouk = intsctk / (unionk + eps) x1, y1, x2, y2 = boxes1.unbind(dim=-1) x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) # smallest enclosing box xc1 = torch.minimum(x1, x1g) yc1 = torch.minimum(y1, y1g) xc2 = torch.maximum(x2, x2g) yc2 = torch.maximum(y2, y2g) area_c = (xc2 - xc1) * (yc2 - yc1) miouk = iouk - ((area_c - unionk) / (area_c + eps)) loss = 1 - miouk if reduction == "mean": loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() elif reduction == "sum": loss = loss.sum() return loss msadapter/ms_adapter/torchvision/ops/misc.py000066400000000000000000000275021436713554100216510ustar00rootroot00000000000000import warnings from typing import Callable, List, Optional import ms_adapter.pytorch as torch from ms_adapter.pytorch import Tensor interpolate = torch.nn.functional.interpolate class FrozenBatchNorm2d(torch.nn.Module): """ BatchNorm2d where the batch statistics and the affine parameters are fixed Args: num_features (int): Number of features ``C`` from an expected input of size ``(N, C, H, W)`` eps (float): a value added to the denominator for numerical stability. Default: 1e-5 """ def __init__( self, num_features: int, eps: float = 1e-5, ): super().__init__() # _log_api_usage_once(self) self.eps = eps self.register_buffer("weight", torch.ones(num_features)) self.register_buffer("bias", torch.zeros(num_features)) self.register_buffer("running_mean", torch.zeros(num_features)) self.register_buffer("running_var", torch.ones(num_features)) def forward(self, x: Tensor) -> Tensor: # move reshapes to the beginning # to make it fuser-friendly w = self.weight.reshape(1, -1, 1, 1) b = self.bias.reshape(1, -1, 1, 1) rv = self.running_var.reshape(1, -1, 1, 1) rm = self.running_mean.reshape(1, -1, 1, 1) scale = w * (rv + self.eps).rsqrt() bias = b - rm * scale return x * scale + bias def __repr__(self) -> str: return f"{self.__class__.__name__}({self.weight.shape[0]}, eps={self.eps})" class ConvNormActivation(torch.nn.Sequential): def __init__( self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: int = 1, padding: Optional[int] = None, groups: int = 1, norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d, activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, dilation: int = 1, inplace: Optional[bool] = True, bias: Optional[bool] = None, conv_layer: Callable[..., torch.nn.Module] = torch.nn.Conv2d, ) -> None: if padding is None: padding = (kernel_size - 1) // 2 * dilation if bias is None: bias = norm_layer is None layers = [ conv_layer( in_channels, out_channels, kernel_size, stride, padding, dilation=dilation, groups=groups, bias=bias, ) ] if norm_layer is not None: layers.append(norm_layer(out_channels)) if activation_layer is not None: params = {} if inplace is None else {"inplace": inplace} layers.append(activation_layer(**params)) super().__init__(*layers) # _log_api_usage_once(self) self.out_channels = out_channels if self.__class__ == ConvNormActivation: warnings.warn( "Don't use ConvNormActivation directly, please use Conv2dNormActivation and Conv3dNormActivation instead." ) class Conv2dNormActivation(ConvNormActivation): """ Configurable block used for Convolution2d-Normalization-Activation blocks. Args: in_channels (int): Number of channels in the input image out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block kernel_size: (int, optional): Size of the convolving kernel. Default: 3 stride (int, optional): Stride of the convolution. Default: 1 padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation`` groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1 norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm2d`` activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU`` dilation (int): Spacing between kernel elements. Default: 1 inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True`` bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``. """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: int = 1, padding: Optional[int] = None, groups: int = 1, norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d, activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, dilation: int = 1, inplace: Optional[bool] = True, bias: Optional[bool] = None, ) -> None: super().__init__( in_channels, out_channels, kernel_size, stride, padding, groups, norm_layer, activation_layer, dilation, inplace, bias, torch.nn.Conv2d, ) class Conv3dNormActivation(ConvNormActivation): """ Configurable block used for Convolution3d-Normalization-Activation blocks. Args: in_channels (int): Number of channels in the input video. out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block kernel_size: (int, optional): Size of the convolving kernel. Default: 3 stride (int, optional): Stride of the convolution. Default: 1 padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation`` groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1 norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm3d`` activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU`` dilation (int): Spacing between kernel elements. Default: 1 inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True`` bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``. """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: int = 1, padding: Optional[int] = None, groups: int = 1, norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm3d, activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, dilation: int = 1, inplace: Optional[bool] = True, bias: Optional[bool] = None, ) -> None: super().__init__( in_channels, out_channels, kernel_size, stride, padding, groups, norm_layer, activation_layer, dilation, inplace, bias, torch.nn.Conv3d, ) class SqueezeExcitation(torch.nn.Module): """ This block implements the Squeeze-and-Excitation block from https://arxiv.org/abs/1709.01507 (see Fig. 1). Parameters ``activation``, and ``scale_activation`` correspond to ``delta`` and ``sigma`` in eq. 3. Args: input_channels (int): Number of channels in the input image squeeze_channels (int): Number of squeeze channels activation (Callable[..., torch.nn.Module], optional): ``delta`` activation. Default: ``torch.nn.ReLU`` scale_activation (Callable[..., torch.nn.Module]): ``sigma`` activation. Default: ``torch.nn.Sigmoid`` """ def __init__( self, input_channels: int, squeeze_channels: int, activation: Callable[..., torch.nn.Module] = torch.nn.ReLU, scale_activation: Callable[..., torch.nn.Module] = torch.nn.Sigmoid, ) -> None: super().__init__() # _log_api_usage_once(self) self.avgpool = torch.nn.AdaptiveAvgPool2d(1) self.fc1 = torch.nn.Conv2d(input_channels, squeeze_channels, 1) self.fc2 = torch.nn.Conv2d(squeeze_channels, input_channels, 1) self.activation = activation() self.scale_activation = scale_activation() def _scale(self, input: Tensor) -> Tensor: scale = self.avgpool(input) scale = self.fc1(scale) scale = self.activation(scale) scale = self.fc2(scale) return self.scale_activation(scale) def forward(self, input: Tensor) -> Tensor: scale = self._scale(input) return scale * input class MLP(torch.nn.Sequential): """This block implements the multi-layer perceptron (MLP) module. Args: in_channels (int): Number of channels of the input hidden_channels (List[int]): List of the hidden channel dimensions norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer wont be used. Default: ``None`` activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU`` inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True`` bias (bool): Whether to use bias in the linear layer. Default ``True`` dropout (float): The probability for the dropout layer. Default: 0.0 """ def __init__( self, in_channels: int, hidden_channels: List[int], norm_layer: Optional[Callable[..., torch.nn.Module]] = None, activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU, inplace: Optional[bool] = True, bias: bool = True, dropout: float = 0.0, ): # The addition of `norm_layer` is inspired from the implementation of TorchMultimodal: # https://github.com/facebookresearch/multimodal/blob/5dec8a/torchmultimodal/modules/layers/mlp.py params = {} if inplace is None else {"inplace": inplace} layers = [] in_dim = in_channels for hidden_dim in hidden_channels[:-1]: layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias)) if norm_layer is not None: layers.append(norm_layer(hidden_dim)) layers.append(activation_layer(**params)) layers.append(torch.nn.Dropout(dropout, **params)) in_dim = hidden_dim layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias)) layers.append(torch.nn.Dropout(dropout, **params)) super().__init__(*layers) # _log_api_usage_once(self) class Permute(torch.nn.Module): """This module returns a view of the tensor input with its dimensions permuted. Args: dims (List[int]): The desired ordering of dimensions """ def __init__(self, dims: List[int]): super().__init__() self.dims = dims def forward(self, x: Tensor) -> Tensor: return torch.permute(x, self.dims) msadapter/ms_adapter/torchvision/ops/poolers.py000066400000000000000000000270061436713554100224000ustar00rootroot00000000000000import warnings from typing import Optional, List, Dict, Tuple, Union import ms_adapter.pytorch as torch import ms_adapter.torchvision as torchvision from ms_adapter.pytorch import nn, Tensor from ms_adapter.torchvision.ops.boxes import box_area from .roi_align import roi_align # TODO: (eellison) T54974082 https://github.com/pytorch/pytorch/issues/26744/pytorch/issues/26744 def initLevelMapper( k_min: int, k_max: int, canonical_scale: int = 224, canonical_level: int = 4, eps: float = 1e-6, ): return LevelMapper(k_min, k_max, canonical_scale, canonical_level, eps) class LevelMapper: """Determine which FPN level each RoI in a set of RoIs should map to based on the heuristic in the FPN paper. Args: k_min (int) k_max (int) canonical_scale (int) canonical_level (int) eps (float) """ def __init__( self, k_min: int, k_max: int, canonical_scale: int = 224, canonical_level: int = 4, eps: float = 1e-6, ): self.k_min = k_min self.k_max = k_max self.s0 = canonical_scale self.lvl0 = canonical_level self.eps = eps def __call__(self, boxlists: List[Tensor]) -> Tensor: """ Args: boxlists (list[BoxList]) """ # Compute level ids s = torch.sqrt(torch.cat([box_area(boxlist) for boxlist in boxlists])) # Eqn.(1) in FPN paper target_lvls = torch.floor(self.lvl0 + torch.log2(s / self.s0) + torch.tensor(self.eps, dtype=s.dtype)) target_lvls = torch.clamp(target_lvls, min=self.k_min, max=self.k_max) return (target_lvls.to(torch.int64) - self.k_min).to(torch.int64) def _convert_to_roi_format(boxes: List[Tensor]) -> Tensor: concat_boxes = torch.cat(boxes, dim=0) device, dtype = concat_boxes.device, concat_boxes.dtype ids = torch.cat( [torch.full_like(b[:, :1], i, dtype=dtype, layout=torch.strided, device=device) for i, b in enumerate(boxes)], dim=0, ) rois = torch.cat([ids, concat_boxes], dim=1) return rois def _infer_scale(feature: Tensor, original_size: List[int]) -> float: # assumption: the scale is of the form 2 ** (-k), with k integer size = feature.shape[-2:] possible_scales: List[float] = [] for s1, s2 in zip(size, original_size): approx_scale = float(s1) / float(s2) scale = 2 ** float(torch.tensor(approx_scale).log2().round()) possible_scales.append(scale) return possible_scales[0] def _setup_scales( features: List[Tensor], image_shapes: List[Tuple[int, int]], canonical_scale: int, canonical_level: int ) -> Tuple[List[float], LevelMapper]: if not image_shapes: raise ValueError("images list should not be empty") max_x = 0 max_y = 0 for shape in image_shapes: max_x = max(shape[0], max_x) max_y = max(shape[1], max_y) original_input_shape = (max_x, max_y) scales = [_infer_scale(feat, original_input_shape) for feat in features] # get the levels in the feature map by leveraging the fact that the network always # downsamples by a factor of 2 at each level. lvl_min = -torch.log2(torch.tensor(scales[0], dtype=torch.float32)).item() lvl_max = -torch.log2(torch.tensor(scales[-1], dtype=torch.float32)).item() map_levels = initLevelMapper( int(lvl_min), int(lvl_max), canonical_scale=canonical_scale, canonical_level=canonical_level, ) return scales, map_levels def _filter_input(x: Dict[str, Tensor], featmap_names: List[str]) -> List[Tensor]: x_filtered = [] for k, v in x.items(): if k in featmap_names: x_filtered.append(v) return x_filtered def _multiscale_roi_align( x_filtered: List[Tensor], boxes: List[Tensor], output_size: List[int], sampling_ratio: int, scales: Optional[List[float]], mapper: Optional[LevelMapper], ) -> Tensor: """ Args: x_filtered (List[Tensor]): List of input tensors. boxes (List[Tensor[N, 4]]): boxes to be used to perform the pooling operation, in (x1, y1, x2, y2) format and in the image reference size, not the feature map reference. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. output_size (Union[List[Tuple[int, int]], List[int]]): size of the output sampling_ratio (int): sampling ratio for ROIAlign scales (Optional[List[float]]): If None, scales will be automatically infered. Default value is None. mapper (Optional[LevelMapper]): If none, mapper will be automatically infered. Default value is None. Returns: result (Tensor) """ if scales is None or mapper is None: raise ValueError("scales and mapper should not be None") num_levels = len(x_filtered) rois = _convert_to_roi_format(boxes) if num_levels == 1: return roi_align( x_filtered[0], rois, output_size=output_size, spatial_scale=scales[0], sampling_ratio=sampling_ratio, ) levels = mapper(boxes) num_rois = len(rois) num_channels = x_filtered[0].shape[1] dtype, device = x_filtered[0].dtype, x_filtered[0].device result = torch.zeros( ( num_rois, num_channels, ) + output_size, dtype=dtype, device=device, ) tracing_results = [] for level, (per_level_feature, scale) in enumerate(zip(x_filtered, scales)): idx_in_level = torch.where(levels == level)[0] rois_per_level = rois[idx_in_level] result_idx_in_level = roi_align( per_level_feature, rois_per_level, output_size=output_size, spatial_scale=scale, sampling_ratio=sampling_ratio, ) # if torchvision._is_tracing(): # tracing_results.append(result_idx_in_level.to(dtype)) # else: # # result and result_idx_in_level's dtypes are based on dtypes of different # # elements in x_filtered. x_filtered contains tensors output by different # # layers. When autocast is active, it may choose different dtypes for # # different layers' outputs. Therefore, we defensively match result's dtype # # before copying elements from result_idx_in_level in the following op. # # We need to cast manually (can't rely on autocast to cast for us) because # # the op acts on result in-place, and autocast only affects out-of-place ops. # result[idx_in_level] = result_idx_in_level.to(result.dtype) result[idx_in_level] = result_idx_in_level.to(result.dtype) return result class MultiScaleRoIAlign(nn.Module): """ Multi-scale RoIAlign pooling, which is useful for detection with or without FPN. It infers the scale of the pooling via the heuristics specified in eq. 1 of the `Feature Pyramid Network paper `_. They keyword-only parameters ``canonical_scale`` and ``canonical_level`` correspond respectively to ``224`` and ``k0=4`` in eq. 1, and have the following meaning: ``canonical_level`` is the target level of the pyramid from which to pool a region of interest with ``w x h = canonical_scale x canonical_scale``. Args: featmap_names (List[str]): the names of the feature maps that will be used for the pooling. output_size (List[Tuple[int, int]] or List[int]): output size for the pooled region sampling_ratio (int): sampling ratio for ROIAlign canonical_scale (int, optional): canonical_scale for LevelMapper canonical_level (int, optional): canonical_level for LevelMapper Examples:: >>> m = torchvision.ops.MultiScaleRoIAlign(['feat1', 'feat3'], 3, 2) >>> i = OrderedDict() >>> i['feat1'] = torch.rand(1, 5, 64, 64) >>> i['feat2'] = torch.rand(1, 5, 32, 32) # this feature won't be used in the pooling >>> i['feat3'] = torch.rand(1, 5, 16, 16) >>> # create some random bounding boxes >>> boxes = torch.rand(6, 4) * 256; boxes[:, 2:] += boxes[:, :2] >>> # original image size, before computing the feature maps >>> image_sizes = [(512, 512)] >>> output = m(i, [boxes], image_sizes) >>> print(output.shape) >>> torch.Size([6, 5, 3, 3]) """ __annotations__ = {"scales": Optional[List[float]], "map_levels": Optional[LevelMapper]} def __init__( self, featmap_names: List[str], output_size: Union[int, Tuple[int], List[int]], sampling_ratio: int, *, canonical_scale: int = 224, canonical_level: int = 4, ): super().__init__() # _log_api_usage_once(self) if isinstance(output_size, int): output_size = (output_size, output_size) self.featmap_names = featmap_names self.sampling_ratio = sampling_ratio self.output_size = tuple(output_size) self.scales = None self.map_levels = None self.canonical_scale = canonical_scale self.canonical_level = canonical_level def convert_to_roi_format(self, boxes: List[Tensor]) -> Tensor: warnings.warn("The 'convert_to_roi_format' method is deprecated since 0.12 and will be removed in 0.14.") return _convert_to_roi_format(boxes) def infer_scale(self, feature: Tensor, original_size: List[int]) -> float: warnings.warn("The 'infer_scale' method is deprecated since 0.12 and will be removed in 0.14.") return _infer_scale(feature, original_size) def setup_setup_scales( self, features: List[Tensor], image_shapes: List[Tuple[int, int]], ) -> None: warnings.warn("The 'setup_setup_scales' method is deprecated since 0.12 and will be removed in 0.14.") self.scales, self.map_levels = _setup_scales(features, image_shapes, self.canonical_scale, self.canonical_level) def forward( self, x: Dict[str, Tensor], boxes: List[Tensor], image_shapes: List[Tuple[int, int]], ) -> Tensor: """ Args: x (OrderedDict[Tensor]): feature maps for each level. They are assumed to have all the same number of channels, but they can have different sizes. boxes (List[Tensor[N, 4]]): boxes to be used to perform the pooling operation, in (x1, y1, x2, y2) format and in the image reference size, not the feature map reference. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. image_shapes (List[Tuple[height, width]]): the sizes of each image before they have been fed to a CNN to obtain feature maps. This allows us to infer the scale factor for each one of the levels to be pooled. Returns: result (Tensor) """ x_filtered = _filter_input(x, self.featmap_names) if self.scales is None or self.map_levels is None: self.scales, self.map_levels = _setup_scales( x_filtered, image_shapes, self.canonical_scale, self.canonical_level ) return _multiscale_roi_align( x_filtered, boxes, self.output_size, self.sampling_ratio, self.scales, self.map_levels, ) def __repr__(self) -> str: return ( f"{self.__class__.__name__}(featmap_names={self.featmap_names}, " f"output_size={self.output_size}, sampling_ratio={self.sampling_ratio})" ) msadapter/ms_adapter/torchvision/ops/ps_roi_align.py000066400000000000000000000064731436713554100233670ustar00rootroot00000000000000import ms_adapter.pytorch as torch from ms_adapter.pytorch import nn, Tensor from ms_adapter.pytorch.nn.modules.utils import _pair from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def ps_roi_align( input: Tensor, boxes: Tensor, output_size: int, spatial_scale: float = 1.0, sampling_ratio: int = -1, ) -> Tensor: """ Performs Position-Sensitive Region of Interest (RoI) Align operator mentioned in Light-Head R-CNN. Args: input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element contains ``C`` feature maps of dimensions ``H x W``. boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) format where the regions will be taken from. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. If a single Tensor is passed, then the first column should contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i in the batch. output_size (int or Tuple[int, int]): the size of the output (in bins or pixels) after the pooling is performed, as (height, width). spatial_scale (float): a scaling factor that maps the box coordinates to the input coordinates. For example, if your boxes are defined on the scale of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of the original image), you'll want to set this to 0.5. Default: 1.0 sampling_ratio (int): number of sampling points in the interpolation grid used to compute the output value of each pooled output bin. If > 0, then exactly ``sampling_ratio x sampling_ratio`` sampling points per bin are used. If <= 0, then an adaptive number of grid points are used (computed as ``ceil(roi_width / output_width)``, and likewise for height). Default: -1 Returns: Tensor[K, C / (output_size[0] * output_size[1]), output_size[0], output_size[1]]: The pooled RoIs """ check_roi_boxes_shape(boxes) rois = boxes output_size = _pair(output_size) if not isinstance(rois, torch.Tensor): rois = convert_boxes_to_roi_format(rois) output, _ = torch.ops.torchvision.ps_roi_align( input, rois, spatial_scale, output_size[0], output_size[1], sampling_ratio ) return output class PSRoIAlign(nn.Module): """ See :func:`ps_roi_align`. """ def __init__( self, output_size: int, spatial_scale: float, sampling_ratio: int, ): super().__init__() self.output_size = output_size self.spatial_scale = spatial_scale self.sampling_ratio = sampling_ratio def forward(self, input: Tensor, rois: Tensor) -> Tensor: return ps_roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio) def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(" f"output_size={self.output_size}" f", spatial_scale={self.spatial_scale}" f", sampling_ratio={self.sampling_ratio}" f")" ) return s msadapter/ms_adapter/torchvision/ops/ps_roi_pool.py000066400000000000000000000060101436713554100232310ustar00rootroot00000000000000import mindspore as ms import ms_adapter.pytorch as torch from ms_adapter.pytorch import nn, Tensor from ms_adapter.pytorch.nn.modules.utils import _pair from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def ps_roi_pool( input: Tensor, boxes: Tensor, output_size: int, spatial_scale: float = 1.0, ) -> Tensor: """ Performs Position-Sensitive Region of Interest (RoI) Pool operator described in R-FCN Args: input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element contains ``C`` feature maps of dimensions ``H x W``. boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) format where the regions will be taken from. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. If a single Tensor is passed, then the first column should contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i in the batch. output_size (int or Tuple[int, int]): the size of the output (in bins or pixels) after the pooling is performed, as (height, width). spatial_scale (float): a scaling factor that maps the box coordinates to the input coordinates. For example, if your boxes are defined on the scale of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of the original image), you'll want to set this to 0.5. Default: 1.0 Returns: Tensor[K, C / (output_size[0] * output_size[1]), output_size[0], output_size[1]]: The pooled RoIs. """ check_roi_boxes_shape(boxes) rois = boxes output_size = _pair(output_size) if not isinstance(rois, torch.Tensor): rois = convert_boxes_to_roi_format(rois) # output, _ = torch.ops.torchvision.ps_roi_pool(input, rois, spatial_scale, output_size[0], output_size[1]) input = cast_to_ms_tensor(input) rois = cast_to_ms_tensor(rois) if len(input.shape) != 4: raise ValueError("The input must be Tensor[N, C, H, W]") output_dim = input.shape[1] ps_roi_pool = ms.ops.PSROIPooling(spatial_scale, output_size, output_dim) output = ps_roi_pool(input, rois) return cast_to_adapter_tensor(output) class PSRoIPool(nn.Module): """ See :func:`ps_roi_pool`. """ def __init__(self, output_size: int, spatial_scale: float): super().__init__() self.output_size = output_size self.spatial_scale = spatial_scale def forward(self, input: Tensor, rois: Tensor) -> Tensor: return ps_roi_pool(input, rois, self.output_size, self.spatial_scale) def __repr__(self) -> str: s = f"{self.__class__.__name__}(output_size={self.output_size}, spatial_scale={self.spatial_scale})" return s msadapter/ms_adapter/torchvision/ops/roi_align.py000066400000000000000000000104321436713554100226530ustar00rootroot00000000000000from typing import List, Union import mindspore as ms import ms_adapter.pytorch as torch from ms_adapter.pytorch import nn, Tensor from ms_adapter.pytorch.nn.modules.utils import _pair from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from mindspore.ops._primitive_cache import _get_cache_prim from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def roi_align( input: Tensor, boxes: Union[Tensor, List[Tensor]], output_size, spatial_scale: float = 1.0, sampling_ratio: int = -1, aligned: bool = False, ) -> Tensor: """ Performs Region of Interest (RoI) Align operator with average pooling, as described in Mask R-CNN. Args: input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element contains ``C`` feature maps of dimensions ``H x W``. If the tensor is quantized, we expect a batch size of ``N == 1``. boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) format where the regions will be taken from. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. If a single Tensor is passed, then the first column should contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i in the batch. output_size (int or Tuple[int, int]): the size of the output (in bins or pixels) after the pooling is performed, as (height, width). spatial_scale (float): a scaling factor that maps the box coordinates to the input coordinates. For example, if your boxes are defined on the scale of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of the original image), you'll want to set this to 0.5. Default: 1.0 sampling_ratio (int): number of sampling points in the interpolation grid used to compute the output value of each pooled output bin. If > 0, then exactly ``sampling_ratio x sampling_ratio`` sampling points per bin are used. If <= 0, then an adaptive number of grid points are used (computed as ``ceil(roi_width / output_width)``, and likewise for height). Default: -1 aligned (bool): If False, use the legacy implementation. If True, pixel shift the box coordinates it by -0.5 for a better alignment with the two neighboring pixel indices. This version is used in Detectron2 Returns: Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs. """ check_roi_boxes_shape(boxes) rois = boxes output_size = _pair(output_size) if not isinstance(rois, torch.Tensor): rois = convert_boxes_to_roi_format(rois) roi_end_mode = 0 if aligned: roi_end_mode = 1 if sampling_ratio <= 0: # sampling_ratio = int(np.ceil(rois.shape[1] / output_size[1])) raise NotImplementedError("sampling_ratio is not supported negative number.") roi_align = _get_cache_prim(ms.ops.ROIAlign)(output_size[0], output_size[1], float(spatial_scale), sampling_ratio, roi_end_mode) input = cast_to_ms_tensor(input) rois = cast_to_ms_tensor(rois) output = roi_align(input, rois) return cast_to_adapter_tensor(output) class RoIAlign(nn.Module): """ See :func:`roi_align`. """ def __init__( self, output_size, spatial_scale: float, sampling_ratio: int, aligned: bool = False, ): super().__init__() self.output_size = output_size self.spatial_scale = spatial_scale self.sampling_ratio = sampling_ratio self.aligned = aligned def forward(self, input: Tensor, rois: Tensor) -> Tensor: return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, self.aligned) def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(" f"output_size={self.output_size}" f", spatial_scale={self.spatial_scale}" f", sampling_ratio={self.sampling_ratio}" f", aligned={self.aligned}" f")" ) return s msadapter/ms_adapter/torchvision/ops/roi_pool.py000066400000000000000000000050311436713554100225310ustar00rootroot00000000000000from typing import List, Union import ms_adapter.pytorch as torch from ms_adapter.pytorch import nn, Tensor from ms_adapter.pytorch.nn.modules.utils import _pair from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def roi_pool( input: Tensor, boxes: Union[Tensor, List[Tensor]], output_size, spatial_scale: float = 1.0, ) -> Tensor: """ Performs Region of Interest (RoI) Pool operator described in Fast R-CNN Args: input (Tensor[N, C, H, W]): The input tensor, i.e. a batch with ``N`` elements. Each element contains ``C`` feature maps of dimensions ``H x W``. boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) format where the regions will be taken from. The coordinate must satisfy ``0 <= x1 < x2`` and ``0 <= y1 < y2``. If a single Tensor is passed, then the first column should contain the index of the corresponding element in the batch, i.e. a number in ``[0, N - 1]``. If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i in the batch. output_size (int or Tuple[int, int]): the size of the output after the cropping is performed, as (height, width) spatial_scale (float): a scaling factor that maps the box coordinates to the input coordinates. For example, if your boxes are defined on the scale of a 224x224 image and your input is a 112x112 feature map (resulting from a 0.5x scaling of the original image), you'll want to set this to 0.5. Default: 1.0 Returns: Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs. """ check_roi_boxes_shape(boxes) rois = boxes output_size = _pair(output_size) if not isinstance(rois, torch.Tensor): rois = convert_boxes_to_roi_format(rois) output, _ = torch.ops.torchvision.roi_pool(input, rois, spatial_scale, output_size[0], output_size[1]) return output class RoIPool(nn.Module): """ See :func:`roi_pool`. """ def __init__(self, output_size, spatial_scale: float): super().__init__() self.output_size = output_size self.spatial_scale = spatial_scale def forward(self, input: Tensor, rois: Tensor) -> Tensor: return roi_pool(input, rois, self.output_size, self.spatial_scale) def __repr__(self) -> str: s = f"{self.__class__.__name__}(output_size={self.output_size}, spatial_scale={self.spatial_scale})" return s msadapter/ms_adapter/torchvision/ops/stochastic_depth.py000066400000000000000000000037561436713554100242530ustar00rootroot00000000000000import ms_adapter.pytorch as torch from ms_adapter.pytorch import nn, Tensor def stochastic_depth(input: Tensor, p: float, mode: str, training: bool = True) -> Tensor: """ Implements the Stochastic Depth from `"Deep Networks with Stochastic Depth" `_ used for randomly dropping residual branches of residual architectures. Args: input (Tensor[N, ...]): The input tensor or arbitrary dimensions with the first one being its batch i.e. a batch with ``N`` rows. p (float): probability of the input to be zeroed. mode (str): ``"batch"`` or ``"row"``. ``"batch"`` randomly zeroes the entire input, ``"row"`` zeroes randomly selected rows from the batch. training: apply stochastic depth if is ``True``. Default: ``True`` Returns: Tensor[N, ...]: The randomly zeroed tensor. """ if p < 0.0 or p > 1.0: raise ValueError(f"drop probability has to be between 0 and 1, but got {p}") if mode not in ["batch", "row"]: raise ValueError(f"mode has to be either 'batch' or 'row', but got {mode}") if not training or p == 0.0: return input survival_rate = 1.0 - p if mode == "row": size = [input.shape[0]] + [1] * (input.ndim - 1) else: size = [1] * input.ndim noise = torch.empty(size, dtype=input.dtype, device=input.device) noise = noise.bernoulli_(survival_rate) if survival_rate > 0.0: noise.div_(survival_rate) return input * noise class StochasticDepth(nn.Module): """ See :func:`stochastic_depth`. """ def __init__(self, p: float, mode: str) -> None: super().__init__() self.p = p self.mode = mode def forward(self, input: Tensor) -> Tensor: return stochastic_depth(input, self.p, self.mode, self.training) def __repr__(self) -> str: s = f"{self.__class__.__name__}(p={self.p}, mode={self.mode})" return s msadapter/ms_adapter/torchvision/transforms/000077500000000000000000000000001436713554100217335ustar00rootroot00000000000000msadapter/ms_adapter/torchvision/transforms/__init__.py000066400000000000000000000000651436713554100240450ustar00rootroot00000000000000from .transforms import * from .autoaugment import * msadapter/ms_adapter/torchvision/transforms/_functional_video.py000066400000000000000000000076501436713554100260040ustar00rootroot00000000000000from __future__ import absolute_import import warnings import ms_adapter.pytorch as torch warnings.warn( "The 'torchvision.transforms._functional_video' module is deprecated since 0.12 and will be removed in 0.14. " "Please use the 'torchvision.transforms.functional' module instead." ) def _is_tensor_video_clip(clip): if not torch.is_tensor(clip): raise TypeError("clip should be Tensor. Got %s" % type(clip)) if not clip.ndimension() == 4: raise ValueError("clip should be 4D. Got %dD" % clip.dim()) return True def crop(clip, i, j, h, w): """ Args: clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) """ if len(clip.size()) != 4: raise ValueError("clip should be a 4D tensor") return clip[..., i : i + h, j : j + w] def resize(clip, target_size, interpolation_mode): if len(target_size) != 2: raise ValueError(f"target size should be tuple (height, width), instead got {target_size}") return torch.nn.functional.interpolate(clip, size=target_size, mode=interpolation_mode, align_corners=False) def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"): """ Do spatial cropping and resizing to the video clip Args: clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) i (int): i in (i,j) i.e coordinates of the upper left corner. j (int): j in (i,j) i.e coordinates of the upper left corner. h (int): Height of the cropped region. w (int): Width of the cropped region. size (tuple(int, int)): height and width of resized clip Returns: clip (torch.tensor): Resized and cropped clip. Size is (C, T, H, W) """ if not _is_tensor_video_clip(clip): raise ValueError("clip should be a 4D torch.tensor") clip = crop(clip, i, j, h, w) clip = resize(clip, size, interpolation_mode) return clip def center_crop(clip, crop_size): if not _is_tensor_video_clip(clip): raise ValueError("clip should be a 4D torch.tensor") h, w = clip.size(-2), clip.size(-1) th, tw = crop_size if h < th or w < tw: raise ValueError("height and width must be no smaller than crop_size") i = int(round((h - th) / 2.0)) j = int(round((w - tw) / 2.0)) return crop(clip, i, j, th, tw) def to_tensor(clip): """ Convert tensor data type from uint8 to float, divide value by 255.0 and permute the dimensions of clip tensor Args: clip (torch.tensor, dtype=torch.uint8): Size is (T, H, W, C) Return: clip (torch.tensor, dtype=torch.float): Size is (C, T, H, W) """ _is_tensor_video_clip(clip) if not clip.dtype == torch.uint8: raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype)) return clip.float().permute(3, 0, 1, 2) / 255.0 def normalize(clip, mean, std, inplace=False): """ Args: clip (torch.tensor): Video clip to be normalized. Size is (C, T, H, W) mean (tuple): pixel RGB mean. Size is (3) std (tuple): pixel standard deviation. Size is (3) Returns: normalized clip (torch.tensor): Size is (C, T, H, W) """ if not _is_tensor_video_clip(clip): raise ValueError("clip should be a 4D torch.tensor") if not inplace: clip = clip.clone() mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device) std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device) # clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None]) clip = clip.sub(mean[:, None, None, None]) clip = clip.div(std[:, None, None, None]) return clip def hflip(clip): """ Args: clip (torch.tensor): Video clip to be normalized. Size is (C, T, H, W) Returns: flipped clip (torch.tensor): Size is (C, T, H, W) """ if not _is_tensor_video_clip(clip): raise ValueError("clip should be a 4D torch.tensor") return clip.flip(-1) msadapter/ms_adapter/torchvision/transforms/_pil_constants.py000066400000000000000000000014621436713554100253270ustar00rootroot00000000000000from PIL import Image # See https://pillow.readthedocs.io/en/stable/releasenotes/9.1.0.html#deprecations # TODO: Remove this file once PIL minimal version is >= 9.1 if hasattr(Image, "Resampling"): BICUBIC = Image.Resampling.BICUBIC BILINEAR = Image.Resampling.BILINEAR LINEAR = Image.Resampling.BILINEAR NEAREST = Image.Resampling.NEAREST AFFINE = Image.Transform.AFFINE FLIP_LEFT_RIGHT = Image.Transpose.FLIP_LEFT_RIGHT FLIP_TOP_BOTTOM = Image.Transpose.FLIP_TOP_BOTTOM PERSPECTIVE = Image.Transform.PERSPECTIVE else: BICUBIC = Image.BICUBIC BILINEAR = Image.BILINEAR NEAREST = Image.NEAREST LINEAR = Image.LINEAR AFFINE = Image.AFFINE FLIP_LEFT_RIGHT = Image.FLIP_LEFT_RIGHT FLIP_TOP_BOTTOM = Image.FLIP_TOP_BOTTOM PERSPECTIVE = Image.PERSPECTIVE msadapter/ms_adapter/torchvision/transforms/_presets.py000066400000000000000000000176301436713554100241400ustar00rootroot00000000000000""" This file is part of the private API. Please do not use directly these classes as they will be modified on future versions without warning. The classes should be accessed only via the transforms argument of Weights. """ from __future__ import absolute_import from typing import Optional, Tuple import ms_adapter.pytorch as torch from ms_adapter.pytorch import Tensor, nn from . import functional as F, InterpolationMode __all__ = [ "ObjectDetection", "ImageClassification", "VideoClassification", "SemanticSegmentation", "OpticalFlow", ] class ObjectDetection(nn.Module): def forward(self, img: Tensor) -> Tensor: if not isinstance(img, Tensor): img = F.pil_to_tensor(img) return F.convert_image_dtype(img, torch.float) def __repr__(self) -> str: return self.__class__.__name__ + "()" def describe(self) -> str: return ( "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " "The images are rescaled to ``[0.0, 1.0]``." ) class ImageClassification(nn.Module): def __init__( self, *, crop_size: int, resize_size: int = 256, mean: Tuple[float, ...] = (0.485, 0.456, 0.406), std: Tuple[float, ...] = (0.229, 0.224, 0.225), interpolation: InterpolationMode = InterpolationMode.BILINEAR, ) -> None: super().__init__() self.crop_size = [crop_size] self.resize_size = [resize_size] self.mean = list(mean) self.std = list(std) self.interpolation = interpolation def forward(self, img: Tensor) -> Tensor: img = F.resize(img, self.resize_size, interpolation=self.interpolation) img = F.center_crop(img, self.crop_size) if not isinstance(img, Tensor): img = F.pil_to_tensor(img) img = F.convert_image_dtype(img, torch.float) img = F.normalize(img, mean=self.mean, std=self.std) return img def __repr__(self) -> str: format_string = self.__class__.__name__ + "(" format_string += f"\n crop_size={self.crop_size}" format_string += f"\n resize_size={self.resize_size}" format_string += f"\n mean={self.mean}" format_string += f"\n std={self.std}" format_string += f"\n interpolation={self.interpolation}" format_string += "\n)" return format_string def describe(self) -> str: return ( "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " f"The images are resized to ``resize_size={self.resize_size}`` using ``interpolation={self.interpolation}``, " f"followed by a central crop of ``crop_size={self.crop_size}``. Finally the values are first rescaled to " f"``[0.0, 1.0]`` and then normalized using ``mean={self.mean}`` and ``std={self.std}``." ) class VideoClassification(nn.Module): def __init__( self, *, crop_size: Tuple[int, int], resize_size: Tuple[int, int], mean: Tuple[float, ...] = (0.43216, 0.394666, 0.37645), std: Tuple[float, ...] = (0.22803, 0.22145, 0.216989), interpolation: InterpolationMode = InterpolationMode.BILINEAR, ) -> None: super().__init__() self.crop_size = list(crop_size) self.resize_size = list(resize_size) self.mean = list(mean) self.std = list(std) self.interpolation = interpolation def forward(self, vid: Tensor) -> Tensor: need_squeeze = False if vid.ndim < 5: vid = vid.unsqueeze(dim=0) need_squeeze = True N, T, C, H, W = vid.shape vid = vid.view(-1, C, H, W) vid = F.resize(vid, self.resize_size, interpolation=self.interpolation) vid = F.center_crop(vid, self.crop_size) vid = F.convert_image_dtype(vid, torch.float) vid = F.normalize(vid, mean=self.mean, std=self.std) H, W = self.crop_size vid = vid.view(N, T, C, H, W) vid = vid.permute(0, 2, 1, 3, 4) # (N, T, C, H, W) => (N, C, T, H, W) if need_squeeze: vid = vid.squeeze(dim=0) return vid def __repr__(self) -> str: format_string = self.__class__.__name__ + "(" format_string += f"\n crop_size={self.crop_size}" format_string += f"\n resize_size={self.resize_size}" format_string += f"\n mean={self.mean}" format_string += f"\n std={self.std}" format_string += f"\n interpolation={self.interpolation}" format_string += "\n)" return format_string def describe(self) -> str: return ( "Accepts batched ``(B, T, C, H, W)`` and single ``(T, C, H, W)`` video frame ``torch.Tensor`` objects. " f"The frames are resized to ``resize_size={self.resize_size}`` using ``interpolation={self.interpolation}``, " f"followed by a central crop of ``crop_size={self.crop_size}``. Finally the values are first rescaled to " f"``[0.0, 1.0]`` and then normalized using ``mean={self.mean}`` and ``std={self.std}``. Finally the output " "dimensions are permuted to ``(..., C, T, H, W)`` tensors." ) class SemanticSegmentation(nn.Module): def __init__( self, *, resize_size: Optional[int], mean: Tuple[float, ...] = (0.485, 0.456, 0.406), std: Tuple[float, ...] = (0.229, 0.224, 0.225), interpolation: InterpolationMode = InterpolationMode.BILINEAR, ) -> None: super().__init__() self.resize_size = [resize_size] if resize_size is not None else None self.mean = list(mean) self.std = list(std) self.interpolation = interpolation def forward(self, img: Tensor) -> Tensor: if isinstance(self.resize_size, list): img = F.resize(img, self.resize_size, interpolation=self.interpolation) if not isinstance(img, Tensor): img = F.pil_to_tensor(img) img = F.convert_image_dtype(img, torch.float) img = F.normalize(img, mean=self.mean, std=self.std) return img def __repr__(self) -> str: format_string = self.__class__.__name__ + "(" format_string += f"\n resize_size={self.resize_size}" format_string += f"\n mean={self.mean}" format_string += f"\n std={self.std}" format_string += f"\n interpolation={self.interpolation}" format_string += "\n)" return format_string def describe(self) -> str: return ( "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " f"The images are resized to ``resize_size={self.resize_size}`` using ``interpolation={self.interpolation}``. " f"Finally the values are first rescaled to ``[0.0, 1.0]`` and then normalized using ``mean={self.mean}`` and " f"``std={self.std}``." ) class OpticalFlow(nn.Module): def forward(self, img1: Tensor, img2: Tensor) -> Tuple[Tensor, Tensor]: if not isinstance(img1, Tensor): img1 = F.pil_to_tensor(img1) if not isinstance(img2, Tensor): img2 = F.pil_to_tensor(img2) img1 = F.convert_image_dtype(img1, torch.float) img2 = F.convert_image_dtype(img2, torch.float) # map [0, 1] into [-1, 1] img1 = F.normalize(img1, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) img2 = F.normalize(img2, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) img1 = img1.contiguous() img2 = img2.contiguous() return img1, img2 def __repr__(self) -> str: return self.__class__.__name__ + "()" def describe(self) -> str: return ( "Accepts ``PIL.Image``, batched ``(B, C, H, W)`` and single ``(C, H, W)`` image ``torch.Tensor`` objects. " "The images are rescaled to ``[-1.0, 1.0]``." ) msadapter/ms_adapter/torchvision/transforms/_transforms_video.py000066400000000000000000000116161436713554100260350ustar00rootroot00000000000000#!/usr/bin/env python3 from __future__ import absolute_import import numbers import random import warnings from ms_adapter.torchvision.transforms import ( RandomCrop, RandomResizedCrop, ) from . import _functional_video as F __all__ = [ "RandomCropVideo", "RandomResizedCropVideo", "CenterCropVideo", "NormalizeVideo", "ToTensorVideo", "RandomHorizontalFlipVideo", ] warnings.warn( "The 'torchvision.transforms._transforms_video' module is deprecated since 0.12 and will be removed in 0.14. " "Please use the 'torchvision.transforms' module instead." ) class RandomCropVideo(RandomCrop): def __init__(self, size): if isinstance(size, numbers.Number): self.size = (int(size), int(size)) else: self.size = size def __call__(self, clip): """ Args: clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) Returns: torch.tensor: randomly cropped/resized video clip. size is (C, T, OH, OW) """ i, j, h, w = self.get_params(clip, self.size) return F.crop(clip, i, j, h, w) def __repr__(self) -> str: return f"{self.__class__.__name__}(size={self.size})" class RandomResizedCropVideo(RandomResizedCrop): def __init__( self, size, scale=(0.08, 1.0), ratio=(3.0 / 4.0, 4.0 / 3.0), interpolation_mode="bilinear", ): if isinstance(size, tuple): if len(size) != 2: raise ValueError(f"size should be tuple (height, width), instead got {size}") self.size = size else: self.size = (size, size) self.interpolation_mode = interpolation_mode self.scale = scale self.ratio = ratio def __call__(self, clip): """ Args: clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) Returns: torch.tensor: randomly cropped/resized video clip. size is (C, T, H, W) """ i, j, h, w = self.get_params(clip, self.scale, self.ratio) return F.resized_crop(clip, i, j, h, w, self.size, self.interpolation_mode) def __repr__(self) -> str: return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}, scale={self.scale}, ratio={self.ratio})" class CenterCropVideo: def __init__(self, crop_size): if isinstance(crop_size, numbers.Number): self.crop_size = (int(crop_size), int(crop_size)) else: self.crop_size = crop_size def __call__(self, clip): """ Args: clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W) Returns: torch.tensor: central cropping of video clip. Size is (C, T, crop_size, crop_size) """ return F.center_crop(clip, self.crop_size) def __repr__(self) -> str: return f"{self.__class__.__name__}(crop_size={self.crop_size})" class NormalizeVideo: """ Normalize the video clip by mean subtraction and division by standard deviation Args: mean (3-tuple): pixel RGB mean std (3-tuple): pixel RGB standard deviation inplace (boolean): whether do in-place normalization """ def __init__(self, mean, std, inplace=False): self.mean = mean self.std = std self.inplace = inplace def __call__(self, clip): """ Args: clip (torch.tensor): video clip to be normalized. Size is (C, T, H, W) """ return F.normalize(clip, self.mean, self.std, self.inplace) def __repr__(self) -> str: return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})" class ToTensorVideo: """ Convert tensor data type from uint8 to float, divide value by 255.0 and permute the dimensions of clip tensor """ def __init__(self): pass def __call__(self, clip): """ Args: clip (torch.tensor, dtype=torch.uint8): Size is (T, H, W, C) Return: clip (torch.tensor, dtype=torch.float): Size is (C, T, H, W) """ return F.to_tensor(clip) def __repr__(self) -> str: return self.__class__.__name__ class RandomHorizontalFlipVideo: """ Flip the video clip along the horizonal direction with a given probability Args: p (float): probability of the clip being flipped. Default value is 0.5 """ def __init__(self, p=0.5): self.p = p def __call__(self, clip): """ Args: clip (torch.tensor): Size is (C, T, H, W) Return: clip (torch.tensor): Size is (C, T, H, W) """ if random.random() < self.p: clip = F.hflip(clip) return clip def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" msadapter/ms_adapter/torchvision/transforms/autoaugment.py000066400000000000000000000675171436713554100246560ustar00rootroot00000000000000from __future__ import absolute_import import math from enum import Enum from typing import List, Tuple, Optional, Dict import ms_adapter.pytorch as torch from ms_adapter.pytorch import Tensor from mindspore.ops import constexpr from ms_adapter.pytorch.tensor import cast_to_adapter_tensor from . import functional as F, InterpolationMode __all__ = ["AutoAugmentPolicy", "AutoAugment", "RandAugment", "TrivialAugmentWide", "AugMix"] def _apply_op( img: Tensor, op_name: str, magnitude: float, interpolation: InterpolationMode, fill: Optional[List[float]] ): if op_name == "ShearX": # magnitude should be arctan(magnitude) # official autoaug: (1, level, 0, 0, 1, 0) # https://github.com/tensorflow/models/blob/dd02069717128186b88afa8d857ce57d17957f03/research/autoaugment/augmentation_transforms.py#L290 # compared to # torchvision: (1, tan(level), 0, 0, 1, 0) # https://github.com/pytorch/vision/blob/0c2373d0bba3499e95776e7936e207d8a1676e65/torchvision/transforms/functional.py#L976 img = F.affine( img, angle=0.0, translate=[0, 0], scale=1.0, shear=[math.degrees(math.atan(magnitude)), 0.0], interpolation=interpolation, fill=fill, center=[0, 0], ) elif op_name == "ShearY": # magnitude should be arctan(magnitude) # See above img = F.affine( img, angle=0.0, translate=[0, 0], scale=1.0, shear=[0.0, math.degrees(math.atan(magnitude))], interpolation=interpolation, fill=fill, center=[0, 0], ) elif op_name == "TranslateX": img = F.affine( img, angle=0.0, translate=[int(magnitude), 0], scale=1.0, interpolation=interpolation, shear=[0.0, 0.0], fill=fill, ) elif op_name == "TranslateY": img = F.affine( img, angle=0.0, translate=[0, int(magnitude)], scale=1.0, interpolation=interpolation, shear=[0.0, 0.0], fill=fill, ) elif op_name == "Rotate": img = F.rotate(img, magnitude, interpolation=interpolation, fill=fill) elif op_name == "Brightness": img = F.adjust_brightness(img, 1.0 + magnitude) elif op_name == "Color": img = F.adjust_saturation(img, 1.0 + magnitude) elif op_name == "Contrast": img = F.adjust_contrast(img, 1.0 + magnitude) elif op_name == "Sharpness": img = F.adjust_sharpness(img, 1.0 + magnitude) elif op_name == "Posterize": img = F.posterize(img, int(magnitude)) elif op_name == "Solarize": img = F.solarize(img, magnitude) elif op_name == "AutoContrast": img = F.autocontrast(img) elif op_name == "Equalize": img = F.equalize(img) elif op_name == "Invert": img = F.invert(img) elif op_name == "Identity": pass else: raise ValueError(f"The provided operator {op_name} is not recognized.") return img class AutoAugmentPolicy(Enum): """AutoAugment policies learned on different datasets. Available policies are IMAGENET, CIFAR10 and SVHN. """ IMAGENET = "imagenet" CIFAR10 = "cifar10" SVHN = "svhn" # FIXME: Eliminate copy-pasted code for fill standardization and _augmentation_space() by moving stuff on a base class class AutoAugment(torch.nn.Module): r"""AutoAugment data augmentation method based on `"AutoAugment: Learning Augmentation Strategies from Data" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: policy (AutoAugmentPolicy): Desired policy enum defined by :class:`torchvision.transforms.autoaugment.AutoAugmentPolicy`. Default is ``AutoAugmentPolicy.IMAGENET``. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. """ def __init__( self, policy: AutoAugmentPolicy = AutoAugmentPolicy.IMAGENET, interpolation: InterpolationMode = InterpolationMode.NEAREST, fill: Optional[List[float]] = None, ) -> None: super().__init__() self.policy = policy self.interpolation = interpolation self.fill = fill self.policies = self._get_policies(policy) def _get_policies( self, policy: AutoAugmentPolicy ) -> List[Tuple[Tuple[str, float, Optional[int]], Tuple[str, float, Optional[int]]]]: if policy == AutoAugmentPolicy.IMAGENET: return [ (("Posterize", 0.4, 8), ("Rotate", 0.6, 9)), (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)), (("Equalize", 0.8, None), ("Equalize", 0.6, None)), (("Posterize", 0.6, 7), ("Posterize", 0.6, 6)), (("Equalize", 0.4, None), ("Solarize", 0.2, 4)), (("Equalize", 0.4, None), ("Rotate", 0.8, 8)), (("Solarize", 0.6, 3), ("Equalize", 0.6, None)), (("Posterize", 0.8, 5), ("Equalize", 1.0, None)), (("Rotate", 0.2, 3), ("Solarize", 0.6, 8)), (("Equalize", 0.6, None), ("Posterize", 0.4, 6)), (("Rotate", 0.8, 8), ("Color", 0.4, 0)), (("Rotate", 0.4, 9), ("Equalize", 0.6, None)), (("Equalize", 0.0, None), ("Equalize", 0.8, None)), (("Invert", 0.6, None), ("Equalize", 1.0, None)), (("Color", 0.6, 4), ("Contrast", 1.0, 8)), (("Rotate", 0.8, 8), ("Color", 1.0, 2)), (("Color", 0.8, 8), ("Solarize", 0.8, 7)), (("Sharpness", 0.4, 7), ("Invert", 0.6, None)), (("ShearX", 0.6, 5), ("Equalize", 1.0, None)), (("Color", 0.4, 0), ("Equalize", 0.6, None)), (("Equalize", 0.4, None), ("Solarize", 0.2, 4)), (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)), (("Invert", 0.6, None), ("Equalize", 1.0, None)), (("Color", 0.6, 4), ("Contrast", 1.0, 8)), (("Equalize", 0.8, None), ("Equalize", 0.6, None)), ] elif policy == AutoAugmentPolicy.CIFAR10: return [ (("Invert", 0.1, None), ("Contrast", 0.2, 6)), (("Rotate", 0.7, 2), ("TranslateX", 0.3, 9)), (("Sharpness", 0.8, 1), ("Sharpness", 0.9, 3)), (("ShearY", 0.5, 8), ("TranslateY", 0.7, 9)), (("AutoContrast", 0.5, None), ("Equalize", 0.9, None)), (("ShearY", 0.2, 7), ("Posterize", 0.3, 7)), (("Color", 0.4, 3), ("Brightness", 0.6, 7)), (("Sharpness", 0.3, 9), ("Brightness", 0.7, 9)), (("Equalize", 0.6, None), ("Equalize", 0.5, None)), (("Contrast", 0.6, 7), ("Sharpness", 0.6, 5)), (("Color", 0.7, 7), ("TranslateX", 0.5, 8)), (("Equalize", 0.3, None), ("AutoContrast", 0.4, None)), (("TranslateY", 0.4, 3), ("Sharpness", 0.2, 6)), (("Brightness", 0.9, 6), ("Color", 0.2, 8)), (("Solarize", 0.5, 2), ("Invert", 0.0, None)), (("Equalize", 0.2, None), ("AutoContrast", 0.6, None)), (("Equalize", 0.2, None), ("Equalize", 0.6, None)), (("Color", 0.9, 9), ("Equalize", 0.6, None)), (("AutoContrast", 0.8, None), ("Solarize", 0.2, 8)), (("Brightness", 0.1, 3), ("Color", 0.7, 0)), (("Solarize", 0.4, 5), ("AutoContrast", 0.9, None)), (("TranslateY", 0.9, 9), ("TranslateY", 0.7, 9)), (("AutoContrast", 0.9, None), ("Solarize", 0.8, 3)), (("Equalize", 0.8, None), ("Invert", 0.1, None)), (("TranslateY", 0.7, 9), ("AutoContrast", 0.9, None)), ] elif policy == AutoAugmentPolicy.SVHN: return [ (("ShearX", 0.9, 4), ("Invert", 0.2, None)), (("ShearY", 0.9, 8), ("Invert", 0.7, None)), (("Equalize", 0.6, None), ("Solarize", 0.6, 6)), (("Invert", 0.9, None), ("Equalize", 0.6, None)), (("Equalize", 0.6, None), ("Rotate", 0.9, 3)), (("ShearX", 0.9, 4), ("AutoContrast", 0.8, None)), (("ShearY", 0.9, 8), ("Invert", 0.4, None)), (("ShearY", 0.9, 5), ("Solarize", 0.2, 6)), (("Invert", 0.9, None), ("AutoContrast", 0.8, None)), (("Equalize", 0.6, None), ("Rotate", 0.9, 3)), (("ShearX", 0.9, 4), ("Solarize", 0.3, 3)), (("ShearY", 0.8, 8), ("Invert", 0.7, None)), (("Equalize", 0.9, None), ("TranslateY", 0.6, 6)), (("Invert", 0.9, None), ("Equalize", 0.6, None)), (("Contrast", 0.3, 3), ("Rotate", 0.8, 4)), (("Invert", 0.8, None), ("TranslateY", 0.0, 2)), (("ShearY", 0.7, 6), ("Solarize", 0.4, 8)), (("Invert", 0.6, None), ("Rotate", 0.8, 4)), (("ShearY", 0.3, 7), ("TranslateX", 0.9, 3)), (("ShearX", 0.1, 6), ("Invert", 0.6, None)), (("Solarize", 0.7, 2), ("TranslateY", 0.6, 7)), (("ShearY", 0.8, 4), ("Invert", 0.8, None)), (("ShearX", 0.7, 9), ("TranslateY", 0.8, 3)), (("ShearY", 0.8, 5), ("AutoContrast", 0.7, None)), (("ShearX", 0.7, 2), ("Invert", 0.1, None)), ] else: raise ValueError(f"The provided policy {policy} is not recognized.") def _augmentation_space(self, num_bins: int, image_size: Tuple[int, int]) -> Dict[str, Tuple[Tensor, bool]]: return { # op_name: (magnitudes, signed) "ShearX": (torch.linspace(0.0, 0.3, num_bins), True), "ShearY": (torch.linspace(0.0, 0.3, num_bins), True), "TranslateX": (torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True), "TranslateY": (torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True), "Rotate": (torch.linspace(0.0, 30.0, num_bins), True), "Brightness": (torch.linspace(0.0, 0.9, num_bins), True), "Color": (torch.linspace(0.0, 0.9, num_bins), True), "Contrast": (torch.linspace(0.0, 0.9, num_bins), True), "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True), "Posterize": (8 - (torch.arange(0, num_bins) / ((num_bins - 1) / 4)).round().int(), False), "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), "AutoContrast": (torch.tensor(0.0), False), "Equalize": (torch.tensor(0.0), False), "Invert": (torch.tensor(0.0), False), } @staticmethod def get_params(transform_num: int) -> Tuple[int, Tensor, Tensor]: """Get parameters for autoaugment transformation Returns: params required by the autoaugment transformation """ policy_id = int(torch.randint(0, transform_num, (1,)).item()) probs = torch.rand((2,)) signs = torch.randint(0, 2, (2,)) return policy_id, probs, signs def forward(self, img: Tensor) -> Tensor: """ img (PIL Image or Tensor): Image to be transformed. Returns: PIL Image or Tensor: AutoAugmented image. """ fill = self.fill channels, height, width = F.get_dimensions(img) if isinstance(img, Tensor): if isinstance(fill, (int, float)): fill = [float(fill)] * channels elif fill is not None: fill = [float(f) for f in fill] transform_id, probs, signs = self.get_params(len(self.policies)) op_meta = self._augmentation_space(10, (height, width)) for i, (op_name, p, magnitude_id) in enumerate(self.policies[transform_id]): if probs[i] <= p: magnitudes, signed = op_meta[op_name] magnitude = float(magnitudes[magnitude_id].item()) if magnitude_id is not None else 0.0 if signed and signs[i] == 0: magnitude *= -1.0 img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(policy={self.policy}, fill={self.fill})" class RandAugment(torch.nn.Module): r"""RandAugment data augmentation method based on `"RandAugment: Practical automated data augmentation with a reduced search space" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: num_ops (int): Number of augmentation transformations to apply sequentially. magnitude (int): Magnitude for all the transformations. num_magnitude_bins (int): The number of different magnitude values. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. """ def __init__( self, num_ops: int = 2, magnitude: int = 9, num_magnitude_bins: int = 31, interpolation: InterpolationMode = InterpolationMode.NEAREST, fill: Optional[List[float]] = None, ) -> None: super().__init__() self.num_ops = num_ops self.magnitude = magnitude self.num_magnitude_bins = num_magnitude_bins self.interpolation = interpolation self.fill = fill def _augmentation_space(self, num_bins: int, image_size: Tuple[int, int]) -> Dict[str, Tuple[Tensor, bool]]: return { # op_name: (magnitudes, signed) "Identity": (torch.tensor(0.0), False), "ShearX": (torch.linspace(0.0, 0.3, num_bins), True), "ShearY": (torch.linspace(0.0, 0.3, num_bins), True), "TranslateX": (torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True), "TranslateY": (torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True), "Rotate": (torch.linspace(0.0, 30.0, num_bins), True), "Brightness": (torch.linspace(0.0, 0.9, num_bins), True), "Color": (torch.linspace(0.0, 0.9, num_bins), True), "Contrast": (torch.linspace(0.0, 0.9, num_bins), True), "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True), "Posterize": (8 - (torch.arange(0, num_bins) / ((num_bins - 1) / 4)).round().int(), False), "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), "AutoContrast": (torch.tensor(0.0), False), "Equalize": (torch.tensor(0.0), False), } def forward(self, img: Tensor) -> Tensor: """ img (PIL Image or Tensor): Image to be transformed. Returns: PIL Image or Tensor: Transformed image. """ fill = self.fill channels, height, width = F.get_dimensions(img) if isinstance(img, Tensor): if isinstance(fill, (int, float)): fill = [float(fill)] * channels elif fill is not None: fill = [float(f) for f in fill] op_meta = self._augmentation_space(self.num_magnitude_bins, (height, width)) for _ in range(self.num_ops): op_index = int(torch.randint(0, len(op_meta), (1,)).item()) op_name = list(op_meta.keys())[op_index] magnitudes, signed = op_meta[op_name] magnitude = float(magnitudes[self.magnitude].item()) if magnitudes.ndim > 0 else 0.0 if signed and torch.randint(0, 2, (1,)): magnitude *= -1.0 img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill) return img def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(" f"num_ops={self.num_ops}" f", magnitude={self.magnitude}" f", num_magnitude_bins={self.num_magnitude_bins}" f", interpolation={self.interpolation}" f", fill={self.fill}" f")" ) return s class TrivialAugmentWide(torch.nn.Module): r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: num_magnitude_bins (int): The number of different magnitude values. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. """ def __init__( self, num_magnitude_bins: int = 31, interpolation: InterpolationMode = InterpolationMode.NEAREST, fill: Optional[List[float]] = None, ) -> None: super().__init__() self.num_magnitude_bins = num_magnitude_bins self.interpolation = interpolation self.fill = fill def _augmentation_space(self, num_bins: int) -> Dict[str, Tuple[Tensor, bool]]: return { # op_name: (magnitudes, signed) "Identity": (torch.tensor(0.0), False), "ShearX": (torch.linspace(0.0, 0.99, num_bins), True), "ShearY": (torch.linspace(0.0, 0.99, num_bins), True), "TranslateX": (torch.linspace(0.0, 32.0, num_bins), True), "TranslateY": (torch.linspace(0.0, 32.0, num_bins), True), "Rotate": (torch.linspace(0.0, 135.0, num_bins), True), "Brightness": (torch.linspace(0.0, 0.99, num_bins), True), "Color": (torch.linspace(0.0, 0.99, num_bins), True), "Contrast": (torch.linspace(0.0, 0.99, num_bins), True), "Sharpness": (torch.linspace(0.0, 0.99, num_bins), True), "Posterize": (8 - (torch.arange(0, num_bins) / ((num_bins - 1) / 6)).round().int(), False), "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), "AutoContrast": (torch.tensor(0.0), False), "Equalize": (torch.tensor(0.0), False), } def forward(self, img: Tensor) -> Tensor: """ img (PIL Image or Tensor): Image to be transformed. Returns: PIL Image or Tensor: Transformed image. """ fill = self.fill channels, height, width = F.get_dimensions(img) if isinstance(img, Tensor): if isinstance(fill, (int, float)): fill = [float(fill)] * channels elif fill is not None: fill = [float(f) for f in fill] op_meta = self._augmentation_space(self.num_magnitude_bins) op_index = int(torch.randint(0, len(op_meta), (1,)).item()) op_name = list(op_meta.keys())[op_index] magnitudes, signed = op_meta[op_name] magnitude = ( float(magnitudes[torch.randint(0, len(magnitudes), (1,), dtype=torch.long)].item()) if magnitudes.ndim > 0 else 0.0 ) if signed and torch.randint(0, 2, (1,)): magnitude *= -1.0 return _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill) def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(" f"num_magnitude_bins={self.num_magnitude_bins}" f", interpolation={self.interpolation}" f", fill={self.fill}" f")" ) return s class AugMix(torch.nn.Module): r"""AugMix data augmentation method based on `"AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: severity (int): The severity of base augmentation operators. Default is ``3``. mixture_width (int): The number of augmentation chains. Default is ``3``. chain_depth (int): The depth of augmentation chains. A negative value denotes stochastic depth sampled from the interval [1, 3]. Default is ``-1``. alpha (float): The hyperparameter for the probability distributions. Default is ``1.0``. all_ops (bool): Use all operations (including brightness, contrast, color and sharpness). Default is ``True``. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. """ def __init__( self, severity: int = 3, mixture_width: int = 3, chain_depth: int = -1, alpha: float = 1.0, all_ops: bool = True, interpolation: InterpolationMode = InterpolationMode.BILINEAR, fill: Optional[List[float]] = None, ) -> None: super().__init__() self._PARAMETER_MAX = 10 if not (1 <= severity <= self._PARAMETER_MAX): raise ValueError(f"The severity must be between [1, {self._PARAMETER_MAX}]. Got {severity} instead.") self.severity = severity self.mixture_width = mixture_width self.chain_depth = chain_depth self.alpha = alpha self.all_ops = all_ops self.interpolation = interpolation self.fill = fill def _augmentation_space(self, num_bins: int, image_size: Tuple[int, int]) -> Dict[str, Tuple[Tensor, bool]]: s = { # op_name: (magnitudes, signed) "ShearX": (torch.linspace(0.0, 0.3, num_bins), True), "ShearY": (torch.linspace(0.0, 0.3, num_bins), True), "TranslateX": (torch.linspace(0.0, image_size[1] / 3.0, num_bins), True), "TranslateY": (torch.linspace(0.0, image_size[0] / 3.0, num_bins), True), "Rotate": (torch.linspace(0.0, 30.0, num_bins), True), "Posterize": (4 - (torch.arange(0, num_bins) / ((num_bins - 1) / 4)).round().int(), False), "Solarize": (torch.linspace(255.0, 0.0, num_bins), False), "AutoContrast": (torch.tensor(0.0), False), "Equalize": (torch.tensor(0.0), False), } if self.all_ops: s.update( { "Brightness": (torch.linspace(0.0, 0.9, num_bins), True), "Color": (torch.linspace(0.0, 0.9, num_bins), True), "Contrast": (torch.linspace(0.0, 0.9, num_bins), True), "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True), } ) return s # @torch.jit.unused # @constexpr def _pil_to_tensor(self, img) -> Tensor: return F.pil_to_tensor(img) # @torch.jit.unused # @constexpr def _tensor_to_pil(self, img: Tensor): return F.to_pil_image(img) def _sample_dirichlet(self, params: Tensor) -> Tensor: # Must be on a separate method so that we can overwrite it in tests. return torch._sample_dirichlet(params) def forward(self, orig_img: Tensor) -> Tensor: """ img (PIL Image or Tensor): Image to be transformed. Returns: PIL Image or Tensor: Transformed image. """ fill = self.fill channels, height, width = F.get_dimensions(orig_img) if isinstance(orig_img, Tensor): img = orig_img if isinstance(fill, (int, float)): fill = [float(fill)] * channels elif fill is not None: fill = [float(f) for f in fill] else: img = self._pil_to_tensor(orig_img) op_meta = self._augmentation_space(self._PARAMETER_MAX, (height, width)) orig_dims = list(img.shape) batch = img.view([1] * max(4 - img.ndim, 0) + orig_dims) batch_dims = [batch.size(0)] + [1] * (batch.ndim - 1) # Sample the beta weights for combining the original and augmented image. To get Beta, we use a Dirichlet # with 2 parameters. The 1st column stores the weights of the original and the 2nd the ones of augmented image. m = self._sample_dirichlet( torch.tensor([self.alpha, self.alpha], device=batch.device).expand(batch_dims[0], -1) ) # Sample the mixing weights and combine them with the ones sampled from Beta for the augmented images. combined_weights = self._sample_dirichlet( torch.tensor([self.alpha] * self.mixture_width, device=batch.device).expand(batch_dims[0], -1) ) * m[:, 1].view([batch_dims[0], -1]) mix = m[:, 0].view(batch_dims) * batch for i in range(self.mixture_width): aug = batch depth = self.chain_depth if self.chain_depth > 0 else int(torch.randint(low=1, high=4, size=(1,)).item()) for _ in range(depth): op_index = int(torch.randint(0, len(op_meta), (1,)).item()) op_name = list(op_meta.keys())[op_index] magnitudes, signed = op_meta[op_name] magnitude = ( float(magnitudes[torch.randint(0, self.severity, (1,), dtype=torch.long)].item()) if magnitudes.ndim > 0 else 0.0 ) if signed and torch.randint(0, 2, (1,)): magnitude *= -1.0 aug = _apply_op(aug, op_name, magnitude, interpolation=self.interpolation, fill=fill) mix.add_(combined_weights[:, i].view(batch_dims) * aug) mix = mix.view(orig_dims).to(dtype=img.dtype) if not isinstance(orig_img, Tensor): return self._tensor_to_pil(mix) return mix def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(" f"severity={self.severity}" f", mixture_width={self.mixture_width}" f", chain_depth={self.chain_depth}" f", alpha={self.alpha}" f", all_ops={self.all_ops}" f", interpolation={self.interpolation}" f", fill={self.fill}" f")" ) return s msadapter/ms_adapter/torchvision/transforms/functional.py000066400000000000000000001744461436713554100244670ustar00rootroot00000000000000from __future__ import absolute_import import math import numbers import warnings from enum import Enum from typing import List, Tuple, Any, Optional import numpy as np import ms_adapter.pytorch as torch from PIL import Image from ms_adapter.pytorch import Tensor from mindspore.ops import constexpr try: import accimage except ImportError: accimage = None # from ..utils import _log_api_usage_once from . import functional_pil as F_pil from . import functional_tensor as F_t class InterpolationMode(Enum): """Interpolation modes Available interpolation methods are ``nearest``, ``bilinear``, ``bicubic``, ``box``, ``hamming``, and ``lanczos``. """ NEAREST = "nearest" BILINEAR = "bilinear" BICUBIC = "bicubic" # For PIL compatibility BOX = "box" HAMMING = "hamming" LANCZOS = "lanczos" # TODO: Once torchscript supports Enums with staticmethod # this can be put into InterpolationMode as staticmethod def _interpolation_modes_from_int(i: int) -> InterpolationMode: inverse_modes_mapping = { 0: InterpolationMode.NEAREST, 2: InterpolationMode.BILINEAR, 3: InterpolationMode.BICUBIC, 4: InterpolationMode.BOX, 5: InterpolationMode.HAMMING, 1: InterpolationMode.LANCZOS, } return inverse_modes_mapping[i] pil_modes_mapping = { InterpolationMode.NEAREST: 0, InterpolationMode.BILINEAR: 2, InterpolationMode.BICUBIC: 3, InterpolationMode.BOX: 4, InterpolationMode.HAMMING: 5, InterpolationMode.LANCZOS: 1, } _is_pil_image = F_pil._is_pil_image def get_dimensions(img: Tensor) -> List[int]: """Returns the dimensions of an image as [channels, height, width]. Args: img (PIL Image or Tensor): The image to be checked. Returns: List[int]: The image dimensions. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(get_dimensions) if isinstance(img, torch.Tensor): return F_t.get_dimensions(img) return F_pil.get_dimensions(img) def get_image_size(img: Tensor) -> List[int]: """Returns the size of an image as [width, height]. Args: img (PIL Image or Tensor): The image to be checked. Returns: List[int]: The image size. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(get_image_size) if isinstance(img, torch.Tensor): return F_t.get_image_size(img) return F_pil.get_image_size(img) def get_image_num_channels(img: Tensor) -> int: """Returns the number of channels of an image. Args: img (PIL Image or Tensor): The image to be checked. Returns: int: The number of channels. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(get_image_num_channels) if isinstance(img, torch.Tensor): return F_t.get_image_num_channels(img) return F_pil.get_image_num_channels(img) # @constexpr def _is_numpy(img: Any) -> bool: return isinstance(img, np.ndarray) # @constexpr def _is_numpy_image(img: Any) -> bool: return img.ndim in {2, 3} def to_tensor(pic) -> Tensor: """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. This function does not support torchscript. See :class:`~torchvision.transforms.ToTensor` for more details. Args: pic (PIL Image or numpy.ndarray): Image to be converted to tensor. Returns: Tensor: Converted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(to_tensor) if not (F_pil._is_pil_image(pic) or _is_numpy(pic)): raise TypeError(f"pic should be PIL Image or ndarray. Got {type(pic)}") if _is_numpy(pic) and not _is_numpy_image(pic): raise ValueError(f"pic should be 2/3 dimensional. Got {pic.ndim} dimensions.") # default_float_dtype = torch.get_default_dtype() # TODO: default_float_dtype = torch.float32 if isinstance(pic, np.ndarray): # handle numpy array if pic.ndim == 2: pic = pic[:, :, None] img = torch.from_numpy(pic.transpose((2, 0, 1))).contiguous() # backward compatibility # TODO: if isinstance(img, torch.ByteTensor): if img.dtype == torch.uint8: return img.to(dtype=default_float_dtype).div(255) else: return img if accimage is not None and isinstance(pic, accimage.Image): nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32) pic.copyto(nppic) return torch.from_numpy(nppic).to(dtype=default_float_dtype) # handle PIL Image mode_to_nptype = {"I": np.int32, "I;16": np.int16, "F": np.float32} img = torch.from_numpy(np.array(pic, mode_to_nptype.get(pic.mode, np.uint8), copy=True)) if pic.mode == "1": img = 255 * img img = img.view(pic.size[1], pic.size[0], len(pic.getbands())) # put it from HWC to CHW format img = img.permute((2, 0, 1)).contiguous() # if isinstance(img, torch.ByteTensor): if img.dtype == torch.uint8: return img.to(dtype=default_float_dtype).div(255) else: return img def pil_to_tensor(pic): """Convert a ``PIL Image`` to a tensor of the same type. This function does not support torchscript. See :class:`~torchvision.transforms.PILToTensor` for more details. .. note:: A deep copy of the underlying array is performed. Args: pic (PIL Image): Image to be converted to tensor. Returns: Tensor: Converted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(pil_to_tensor) if not F_pil._is_pil_image(pic): raise TypeError(f"pic should be PIL Image. Got {type(pic)}") if accimage is not None and isinstance(pic, accimage.Image): # accimage format is always uint8 internally, so always return uint8 here nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.uint8) pic.copyto(nppic) return torch.as_tensor(nppic) # handle PIL Image img = torch.as_tensor(np.array(pic, copy=True)) img = img.view(pic.size[1], pic.size[0], len(pic.getbands())) # put it from HWC to CHW format img = img.permute((2, 0, 1)) return img def convert_image_dtype(image: torch.Tensor, dtype = torch.float) -> torch.Tensor: """Convert a tensor image to the given ``dtype`` and scale the values accordingly This function does not support PIL Image. Args: image (torch.Tensor): Image to be converted dtype (torch.dtype): Desired data type of the output Returns: Tensor: Converted image .. note:: When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly. If converted back and forth, this mismatch has no effect. Raises: RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range of the integer ``dtype``. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(convert_image_dtype) if not isinstance(image, torch.Tensor): raise TypeError("Input img should be Tensor Image") return F_t.convert_image_dtype(image, dtype) def to_pil_image(pic, mode=None): """Convert a tensor or an ndarray to PIL Image. This function does not support torchscript. See :class:`~torchvision.transforms.ToPILImage` for more details. Args: pic (Tensor or numpy.ndarray): Image to be converted to PIL Image. mode (`PIL.Image mode`_): color space and pixel depth of input data (optional). .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes Returns: PIL Image: Image converted to PIL Image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(to_pil_image) if not (isinstance(pic, torch.Tensor) or isinstance(pic, np.ndarray)): raise TypeError(f"pic should be Tensor or ndarray. Got {type(pic)}.") elif isinstance(pic, torch.Tensor): if pic.ndimension() not in {2, 3}: raise ValueError(f"pic should be 2/3 dimensional. Got {pic.ndimension()} dimensions.") elif pic.ndimension() == 2: # if 2D image, add channel dimension (CHW) pic = pic.unsqueeze(0) # check number of channels if pic.shape[-3] > 4: raise ValueError(f"pic should not have > 4 channels. Got {pic.shape[-3]} channels.") elif isinstance(pic, np.ndarray): if pic.ndim not in {2, 3}: raise ValueError(f"pic should be 2/3 dimensional. Got {pic.ndim} dimensions.") elif pic.ndim == 2: # if 2D image, add channel dimension (HWC) pic = np.expand_dims(pic, 2) # check number of channels if pic.shape[-1] > 4: raise ValueError(f"pic should not have > 4 channels. Got {pic.shape[-1]} channels.") npimg = pic if isinstance(pic, torch.Tensor): if pic.is_floating_point() and mode != "F": pic = pic.mul(255).byte() npimg = np.transpose(pic.cpu().numpy(), (1, 2, 0)) if not isinstance(npimg, np.ndarray): raise TypeError("Input pic must be a torch.Tensor or NumPy ndarray, not {type(npimg)}") if npimg.shape[2] == 1: expected_mode = None npimg = npimg[:, :, 0] if npimg.dtype == np.uint8: expected_mode = "L" elif npimg.dtype == np.int16: expected_mode = "I;16" elif npimg.dtype == np.int32: expected_mode = "I" elif npimg.dtype == np.float32: expected_mode = "F" if mode is not None and mode != expected_mode: raise ValueError(f"Incorrect mode ({mode}) supplied for input type {np.dtype}. Should be {expected_mode}") mode = expected_mode elif npimg.shape[2] == 2: permitted_2_channel_modes = ["LA"] if mode is not None and mode not in permitted_2_channel_modes: raise ValueError(f"Only modes {permitted_2_channel_modes} are supported for 2D inputs") if mode is None and npimg.dtype == np.uint8: mode = "LA" elif npimg.shape[2] == 4: permitted_4_channel_modes = ["RGBA", "CMYK", "RGBX"] if mode is not None and mode not in permitted_4_channel_modes: raise ValueError(f"Only modes {permitted_4_channel_modes} are supported for 4D inputs") if mode is None and npimg.dtype == np.uint8: mode = "RGBA" else: permitted_3_channel_modes = ["RGB", "YCbCr", "HSV"] if mode is not None and mode not in permitted_3_channel_modes: raise ValueError(f"Only modes {permitted_3_channel_modes} are supported for 3D inputs") if mode is None and npimg.dtype == np.uint8: mode = "RGB" if mode is None: raise TypeError(f"Input type {npimg.dtype} is not supported") return Image.fromarray(npimg, mode=mode) def normalize(tensor: Tensor, mean: List[float], std: List[float], inplace: bool = False) -> Tensor: """Normalize a float tensor image with mean and standard deviation. This transform does not support PIL Image. .. note:: This transform acts out of place by default, i.e., it does not mutates the input tensor. See :class:`~torchvision.transforms.Normalize` for more details. Args: tensor (Tensor): Float tensor image of size (C, H, W) or (B, C, H, W) to be normalized. mean (sequence): Sequence of means for each channel. std (sequence): Sequence of standard deviations for each channel. inplace(bool,optional): Bool to make this operation inplace. Returns: Tensor: Normalized Tensor image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(normalize) if not isinstance(tensor, torch.Tensor): raise TypeError(f"img should be Tensor Image. Got {type(tensor)}") return F_t.normalize(tensor, mean=mean, std=std, inplace=inplace) def resize( img: Tensor, size: List[int], interpolation: InterpolationMode = InterpolationMode.BILINEAR, max_size: Optional[int] = None, antialias: Optional[bool] = None, ) -> Tensor: r"""Resize the input image to the given size. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions .. warning:: The output image might be different depending on its type: when downsampling, the interpolation of PIL images and tensors is slightly different, because PIL applies antialiasing. This may lead to significant differences in the performance of a network. Therefore, it is preferable to train and serve a model with the same input types. See also below the ``antialias`` parameter, which can help making the output of PIL images and tensors closer. Args: img (PIL Image or Tensor): Image to be resized. size (sequence or int): Desired output size. If size is a sequence like (h, w), the output size will be matched to this. If size is an int, the smaller edge of the image will be matched to this number maintaining the aspect ratio. i.e, if height > width, then image will be rescaled to :math:`\left(\text{size} \times \frac{\text{height}}{\text{width}}, \text{size}\right)`. .. note:: In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. max_size (int, optional): The maximum allowed for the longer edge of the resized image: if the longer edge of the image is greater than ``max_size`` after being resized according to ``size``, then the image is resized again so that the longer edge is equal to ``max_size``. As a result, ``size`` might be overruled, i.e the smaller edge may be shorter than ``size``. This is only supported if ``size`` is an int (or a sequence of length 1 in torchscript mode). antialias (bool, optional): antialias flag. If ``img`` is PIL Image, the flag is ignored and anti-alias is always used. If ``img`` is Tensor, the flag is False by default and can be set to True for ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` modes. This can help making the output for PIL images and tensors closer. Returns: PIL Image or Tensor: Resized image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(resize) # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) if not isinstance(interpolation, InterpolationMode): raise TypeError("Argument interpolation should be a InterpolationMode") if not isinstance(img, torch.Tensor): if antialias is not None and not antialias: warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.") pil_interpolation = pil_modes_mapping[interpolation] return F_pil.resize(img, size=size, interpolation=pil_interpolation, max_size=max_size) return F_t.resize(img, size=size, interpolation=interpolation.value, max_size=max_size, antialias=antialias) def pad(img: Tensor, padding: List[int], fill: int = 0, padding_mode: str = "constant") -> Tensor: r"""Pad the given image on all sides with the given "pad" value. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means at most 2 leading dimensions for mode reflect and symmetric, at most 3 leading dimensions for mode edge, and an arbitrary number of leading dimensions for mode constant Args: img (PIL Image or Tensor): Image to be padded. padding (int or sequence): Padding on each border. If a single int is provided this is used to pad all borders. If sequence of length 2 is provided this is the padding on left/right and top/bottom respectively. If a sequence of length 4 is provided this is the padding for the left, top, right and bottom borders respectively. .. note:: In torchscript mode padding as single int is not supported, use a sequence of length 1: ``[padding, ]``. fill (number or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. This value is only used when the padding_mode is constant. Only number is supported for torch Tensor. Only int or tuple value is supported for PIL Image. padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant. - constant: pads with a constant value, this value is specified with fill - edge: pads with the last value at the edge of the image. If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2 - reflect: pads with reflection of image without repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode will result in [3, 2, 1, 2, 3, 4, 3, 2] - symmetric: pads with reflection of image repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode will result in [2, 1, 1, 2, 3, 4, 4, 3] Returns: PIL Image or Tensor: Padded image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(pad) if not isinstance(img, torch.Tensor): return F_pil.pad(img, padding=padding, fill=fill, padding_mode=padding_mode) return F_t.pad(img, padding=padding, fill=fill, padding_mode=padding_mode) def crop(img: Tensor, top: int, left: int, height: int, width: int) -> Tensor: """Crop the given image at specified location and output size. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. If image size is smaller than output size along any edge, image is padded with 0 and then cropped. Args: img (PIL Image or Tensor): Image to be cropped. (0,0) denotes the top left corner of the image. top (int): Vertical component of the top left corner of the crop box. left (int): Horizontal component of the top left corner of the crop box. height (int): Height of the crop box. width (int): Width of the crop box. Returns: PIL Image or Tensor: Cropped image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(crop) if not isinstance(img, torch.Tensor): return F_pil.crop(img, top, left, height, width) return F_t.crop(img, top, left, height, width) def center_crop(img: Tensor, output_size: List[int]) -> Tensor: """Crops the given image at the center. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. If image size is smaller than output size along any edge, image is padded with 0 and then center cropped. Args: img (PIL Image or Tensor): Image to be cropped. output_size (sequence or int): (height, width) of the crop box. If int or sequence with single int, it is used for both directions. Returns: PIL Image or Tensor: Cropped image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(center_crop) if isinstance(output_size, numbers.Number): output_size = (int(output_size), int(output_size)) elif isinstance(output_size, (tuple, list)) and len(output_size) == 1: output_size = (output_size[0], output_size[0]) _, image_height, image_width = get_dimensions(img) crop_height, crop_width = output_size if crop_width > image_width or crop_height > image_height: padding_ltrb = [ (crop_width - image_width) // 2 if crop_width > image_width else 0, (crop_height - image_height) // 2 if crop_height > image_height else 0, (crop_width - image_width + 1) // 2 if crop_width > image_width else 0, (crop_height - image_height + 1) // 2 if crop_height > image_height else 0, ] img = pad(img, padding_ltrb, fill=0) # PIL uses fill value 0 _, image_height, image_width = get_dimensions(img) if crop_width == image_width and crop_height == image_height: return img crop_top = int(round((image_height - crop_height) / 2.0)) crop_left = int(round((image_width - crop_width) / 2.0)) return crop(img, crop_top, crop_left, crop_height, crop_width) def resized_crop( img: Tensor, top: int, left: int, height: int, width: int, size: List[int], interpolation: InterpolationMode = InterpolationMode.BILINEAR, ) -> Tensor: """Crop the given image and resize it to desired size. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions Notably used in :class:`~torchvision.transforms.RandomResizedCrop`. Args: img (PIL Image or Tensor): Image to be cropped. (0,0) denotes the top left corner of the image. top (int): Vertical component of the top left corner of the crop box. left (int): Horizontal component of the top left corner of the crop box. height (int): Height of the crop box. width (int): Width of the crop box. size (sequence or int): Desired output size. Same semantics as ``resize``. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. Returns: PIL Image or Tensor: Cropped image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(resized_crop) img = crop(img, top, left, height, width) img = resize(img, size, interpolation) return img def hflip(img: Tensor) -> Tensor: """Horizontally flip the given image. Args: img (PIL Image or Tensor): Image to be flipped. If img is a Tensor, it is expected to be in [..., H, W] format, where ... means it can have an arbitrary number of leading dimensions. Returns: PIL Image or Tensor: Horizontally flipped image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(hflip) if not isinstance(img, torch.Tensor): return F_pil.hflip(img) return F_t.hflip(img) def _get_perspective_coeffs(startpoints: List[List[int]], endpoints: List[List[int]]) -> List[float]: """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms. In Perspective Transform each pixel (x, y) in the original image gets transformed as, (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) ) Args: startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners ``[top-left, top-right, bottom-right, bottom-left]`` of the original image. endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image. Returns: octuple (a, b, c, d, e, f, g, h) for transforming each pixel. """ a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float) for i, (p1, p2) in enumerate(zip(endpoints, startpoints)): a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]]) a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]]) b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8) # res = torch.linalg.lstsq(a_matrix, b_matrix, driver="gels").solution res, _, _, _ = np.linalg.lstsq(a_matrix.numpy(), b_matrix.numpy()) output: List[float] = res.tolist() return output def perspective( img: Tensor, startpoints: List[List[int]], endpoints: List[List[int]], interpolation: InterpolationMode = InterpolationMode.BILINEAR, fill: Optional[List[float]] = None, ) -> Tensor: """Perform perspective transform of the given image. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. Args: img (PIL Image or Tensor): Image to be transformed. startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners ``[top-left, top-right, bottom-right, bottom-left]`` of the original image. endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. .. note:: In torchscript mode single int/float value is not supported, please use a sequence of length 1: ``[value, ]``. Returns: PIL Image or Tensor: transformed Image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(perspective) coeffs = _get_perspective_coeffs(startpoints, endpoints) # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) if not isinstance(interpolation, InterpolationMode): raise TypeError("Argument interpolation should be a InterpolationMode") if not isinstance(img, torch.Tensor): pil_interpolation = pil_modes_mapping[interpolation] return F_pil.perspective(img, coeffs, interpolation=pil_interpolation, fill=fill) return F_t.perspective(img, coeffs, interpolation=interpolation.value, fill=fill) def vflip(img: Tensor) -> Tensor: """Vertically flip the given image. Args: img (PIL Image or Tensor): Image to be flipped. If img is a Tensor, it is expected to be in [..., H, W] format, where ... means it can have an arbitrary number of leading dimensions. Returns: PIL Image or Tensor: Vertically flipped image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(vflip) if not isinstance(img, torch.Tensor): return F_pil.vflip(img) return F_t.vflip(img) def five_crop(img: Tensor, size: List[int]) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: """Crop the given image into four corners and the central crop. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions .. Note:: This transform returns a tuple of images and there may be a mismatch in the number of inputs and targets your ``Dataset`` returns. Args: img (PIL Image or Tensor): Image to be cropped. size (sequence or int): Desired output size of the crop. If size is an int instead of sequence like (h, w), a square crop (size, size) is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). Returns: tuple: tuple (tl, tr, bl, br, center) Corresponding top left, top right, bottom left, bottom right and center crop. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(five_crop) if isinstance(size, numbers.Number): size = (int(size), int(size)) elif isinstance(size, (tuple, list)) and len(size) == 1: size = (size[0], size[0]) if len(size) != 2: raise ValueError("Please provide only two dimensions (h, w) for size.") _, image_height, image_width = get_dimensions(img) crop_height, crop_width = size if crop_width > image_width or crop_height > image_height: msg = "Requested crop size {} is bigger than input size {}" raise ValueError(msg.format(size, (image_height, image_width))) tl = crop(img, 0, 0, crop_height, crop_width) tr = crop(img, 0, image_width - crop_width, crop_height, crop_width) bl = crop(img, image_height - crop_height, 0, crop_height, crop_width) br = crop(img, image_height - crop_height, image_width - crop_width, crop_height, crop_width) center = center_crop(img, [crop_height, crop_width]) return tl, tr, bl, br, center def ten_crop(img: Tensor, size: List[int], vertical_flip: bool = False) -> List[Tensor]: """Generate ten cropped images from the given image. Crop the given image into four corners and the central crop plus the flipped version of these (horizontal flipping is used by default). If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions .. Note:: This transform returns a tuple of images and there may be a mismatch in the number of inputs and targets your ``Dataset`` returns. Args: img (PIL Image or Tensor): Image to be cropped. size (sequence or int): Desired output size of the crop. If size is an int instead of sequence like (h, w), a square crop (size, size) is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). vertical_flip (bool): Use vertical flipping instead of horizontal Returns: tuple: tuple (tl, tr, bl, br, center, tl_flip, tr_flip, bl_flip, br_flip, center_flip) Corresponding top left, top right, bottom left, bottom right and center crop and same for the flipped image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(ten_crop) if isinstance(size, numbers.Number): size = (int(size), int(size)) elif isinstance(size, (tuple, list)) and len(size) == 1: size = (size[0], size[0]) if len(size) != 2: raise ValueError("Please provide only two dimensions (h, w) for size.") first_five = five_crop(img, size) if vertical_flip: img = vflip(img) else: img = hflip(img) second_five = five_crop(img, size) return first_five + second_five def adjust_brightness(img: Tensor, brightness_factor: float) -> Tensor: """Adjust brightness of an image. Args: img (PIL Image or Tensor): Image to be adjusted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. brightness_factor (float): How much to adjust the brightness. Can be any non negative number. 0 gives a black image, 1 gives the original image while 2 increases the brightness by a factor of 2. Returns: PIL Image or Tensor: Brightness adjusted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(adjust_brightness) if not isinstance(img, torch.Tensor): return F_pil.adjust_brightness(img, brightness_factor) return F_t.adjust_brightness(img, brightness_factor) def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor: """Adjust contrast of an image. Args: img (PIL Image or Tensor): Image to be adjusted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. contrast_factor (float): How much to adjust the contrast. Can be any non negative number. 0 gives a solid gray image, 1 gives the original image while 2 increases the contrast by a factor of 2. Returns: PIL Image or Tensor: Contrast adjusted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(adjust_contrast) if not isinstance(img, torch.Tensor): return F_pil.adjust_contrast(img, contrast_factor) return F_t.adjust_contrast(img, contrast_factor) def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor: """Adjust color saturation of an image. Args: img (PIL Image or Tensor): Image to be adjusted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. saturation_factor (float): How much to adjust the saturation. 0 will give a black and white image, 1 will give the original image while 2 will enhance the saturation by a factor of 2. Returns: PIL Image or Tensor: Saturation adjusted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(adjust_saturation) if not isinstance(img, torch.Tensor): return F_pil.adjust_saturation(img, saturation_factor) return F_t.adjust_saturation(img, saturation_factor) def adjust_hue(img: Tensor, hue_factor: float) -> Tensor: """Adjust hue of an image. The image hue is adjusted by converting the image to HSV and cyclically shifting the intensities in the hue channel (H). The image is then converted back to original image mode. `hue_factor` is the amount of shift in H channel and must be in the interval `[-0.5, 0.5]`. See `Hue`_ for more details. .. _Hue: https://en.wikipedia.org/wiki/Hue Args: img (PIL Image or Tensor): Image to be adjusted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image mode "1", "I", "F" and modes with transparency (alpha channel) are not supported. Note: the pixel values of the input image has to be non-negative for conversion to HSV space; thus it does not work if you normalize your image to an interval with negative values, or use an interpolation that generates negative values before using this function. hue_factor (float): How much to shift the hue channel. Should be in [-0.5, 0.5]. 0.5 and -0.5 give complete reversal of hue channel in HSV space in positive and negative direction respectively. 0 means no shift. Therefore, both -0.5 and 0.5 will give an image with complementary colors while 0 gives the original image. Returns: PIL Image or Tensor: Hue adjusted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(adjust_hue) if not isinstance(img, torch.Tensor): return F_pil.adjust_hue(img, hue_factor) return F_t.adjust_hue(img, hue_factor) def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor: r"""Perform gamma correction on an image. Also known as Power Law Transform. Intensities in RGB mode are adjusted based on the following equation: .. math:: I_{\text{out}} = 255 \times \text{gain} \times \left(\frac{I_{\text{in}}}{255}\right)^{\gamma} See `Gamma Correction`_ for more details. .. _Gamma Correction: https://en.wikipedia.org/wiki/Gamma_correction Args: img (PIL Image or Tensor): PIL Image to be adjusted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image, modes with transparency (alpha channel) are not supported. gamma (float): Non negative real number, same as :math:`\gamma` in the equation. gamma larger than 1 make the shadows darker, while gamma smaller than 1 make dark regions lighter. gain (float): The constant multiplier. Returns: PIL Image or Tensor: Gamma correction adjusted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(adjust_gamma) if not isinstance(img, torch.Tensor): return F_pil.adjust_gamma(img, gamma, gain) return F_t.adjust_gamma(img, gamma, gain) def _get_inverse_affine_matrix( center: List[float], angle: float, translate: List[float], scale: float, shear: List[float], inverted: bool = True ) -> List[float]: # Helper method to compute inverse matrix for affine transformation # Pillow requires inverse affine transformation matrix: # Affine matrix is : M = T * C * RotateScaleShear * C^-1 # # where T is translation matrix: [1, 0, tx | 0, 1, ty | 0, 0, 1] # C is translation matrix to keep center: [1, 0, cx | 0, 1, cy | 0, 0, 1] # RotateScaleShear is rotation with scale and shear matrix # # RotateScaleShear(a, s, (sx, sy)) = # = R(a) * S(s) * SHy(sy) * SHx(sx) # = [ s*cos(a - sy)/cos(sy), s*(-cos(a - sy)*tan(sx)/cos(sy) - sin(a)), 0 ] # [ s*sin(a + sy)/cos(sy), s*(-sin(a - sy)*tan(sx)/cos(sy) + cos(a)), 0 ] # [ 0 , 0 , 1 ] # where R is a rotation matrix, S is a scaling matrix, and SHx and SHy are the shears: # SHx(s) = [1, -tan(s)] and SHy(s) = [1 , 0] # [0, 1 ] [-tan(s), 1] # # Thus, the inverse is M^-1 = C * RotateScaleShear^-1 * C^-1 * T^-1 rot = math.radians(angle) sx = math.radians(shear[0]) sy = math.radians(shear[1]) cx, cy = center tx, ty = translate # RSS without scaling a = math.cos(rot - sy) / math.cos(sy) b = -math.cos(rot - sy) * math.tan(sx) / math.cos(sy) - math.sin(rot) c = math.sin(rot - sy) / math.cos(sy) d = -math.sin(rot - sy) * math.tan(sx) / math.cos(sy) + math.cos(rot) if inverted: # Inverted rotation matrix with scale and shear # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1 matrix = [d, -b, 0.0, -c, a, 0.0] matrix = [x / scale for x in matrix] # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1 matrix[2] += matrix[0] * (-cx - tx) + matrix[1] * (-cy - ty) matrix[5] += matrix[3] * (-cx - tx) + matrix[4] * (-cy - ty) # Apply center translation: C * RSS^-1 * C^-1 * T^-1 matrix[2] += cx matrix[5] += cy else: matrix = [a, b, 0.0, c, d, 0.0] matrix = [x * scale for x in matrix] # Apply inverse of center translation: RSS * C^-1 matrix[2] += matrix[0] * (-cx) + matrix[1] * (-cy) matrix[5] += matrix[3] * (-cx) + matrix[4] * (-cy) # Apply translation and center : T * C * RSS * C^-1 matrix[2] += cx + tx matrix[5] += cy + ty return matrix def rotate( img: Tensor, angle: float, interpolation: InterpolationMode = InterpolationMode.NEAREST, expand: bool = False, center: Optional[List[int]] = None, fill: Optional[List[float]] = None, resample: Optional[int] = None, ) -> Tensor: """Rotate the image by angle. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. Args: img (PIL Image or Tensor): image to be rotated. angle (number): rotation angle value in degrees, counter-clockwise. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. expand (bool, optional): Optional expansion flag. If true, expands the output image to make it large enough to hold the entire rotated image. If false or omitted, make the output image the same size as the input image. Note that the expand flag assumes rotation around the center and no translation. center (sequence, optional): Optional center of rotation. Origin is the upper left corner. Default is the center of the image. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. .. note:: In torchscript mode single int/float value is not supported, please use a sequence of length 1: ``[value, ]``. resample (int, optional): .. warning:: This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``interpolation`` instead. Returns: PIL Image or Tensor: Rotated image. .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(rotate) if resample is not None: warnings.warn( "The parameter 'resample' is deprecated since 0.12 and will be removed 0.14. " "Please use 'interpolation' instead." ) interpolation = _interpolation_modes_from_int(resample) # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) if not isinstance(angle, (int, float)): raise TypeError("Argument angle should be int or float") if center is not None and not isinstance(center, (list, tuple)): raise TypeError("Argument center should be a sequence") if not isinstance(interpolation, InterpolationMode): raise TypeError("Argument interpolation should be a InterpolationMode") if not isinstance(img, torch.Tensor): pil_interpolation = pil_modes_mapping[interpolation] return F_pil.rotate(img, angle=angle, interpolation=pil_interpolation, expand=expand, center=center, fill=fill) center_f = [0.0, 0.0] if center is not None: _, height, width = get_dimensions(img) # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center. center_f = [1.0 * (c - s * 0.5) for c, s in zip(center, [width, height])] # due to current incoherence of rotation angle direction between affine and rotate implementations # we need to set -angle. matrix = _get_inverse_affine_matrix(center_f, -angle, [0.0, 0.0], 1.0, [0.0, 0.0]) return F_t.rotate(img, matrix=matrix, interpolation=interpolation.value, expand=expand, fill=fill) def affine( img: Tensor, angle: float, translate: List[int], scale: float, shear: List[float], interpolation: InterpolationMode = InterpolationMode.NEAREST, fill: Optional[List[float]] = None, resample: Optional[int] = None, fillcolor: Optional[List[float]] = None, center: Optional[List[int]] = None, ) -> Tensor: """Apply affine transformation on the image keeping image center invariant. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. Args: img (PIL Image or Tensor): image to transform. angle (number): rotation angle in degrees between -180 and 180, clockwise direction. translate (sequence of integers): horizontal and vertical translations (post-rotation translation) scale (float): overall scale shear (float or sequence): shear angle value in degrees between -180 to 180, clockwise direction. If a sequence is specified, the first value corresponds to a shear parallel to the x axis, while the second value corresponds to a shear parallel to the y axis. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. .. note:: In torchscript mode single int/float value is not supported, please use a sequence of length 1: ``[value, ]``. fillcolor (sequence or number, optional): .. warning:: This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``fill`` instead. resample (int, optional): .. warning:: This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``interpolation`` instead. center (sequence, optional): Optional center of rotation. Origin is the upper left corner. Default is the center of the image. Returns: PIL Image or Tensor: Transformed image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(affine) if resample is not None: warnings.warn( "The parameter 'resample' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'interpolation' instead." ) interpolation = _interpolation_modes_from_int(resample) # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) if fillcolor is not None: warnings.warn( "The parameter 'fillcolor' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'fill' instead." ) fill = fillcolor if not isinstance(angle, (int, float)): raise TypeError("Argument angle should be int or float") if not isinstance(translate, (list, tuple)): raise TypeError("Argument translate should be a sequence") if len(translate) != 2: raise ValueError("Argument translate should be a sequence of length 2") if scale <= 0.0: raise ValueError("Argument scale should be positive") if not isinstance(shear, (numbers.Number, (list, tuple))): raise TypeError("Shear should be either a single value or a sequence of two values") if not isinstance(interpolation, InterpolationMode): raise TypeError("Argument interpolation should be a InterpolationMode") if isinstance(angle, int): angle = float(angle) if isinstance(translate, tuple): translate = list(translate) if isinstance(shear, numbers.Number): shear = [shear, 0.0] if isinstance(shear, tuple): shear = list(shear) if len(shear) == 1: shear = [shear[0], shear[0]] if len(shear) != 2: raise ValueError(f"Shear should be a sequence containing two values. Got {shear}") if center is not None and not isinstance(center, (list, tuple)): raise TypeError("Argument center should be a sequence") _, height, width = get_dimensions(img) if not isinstance(img, torch.Tensor): # center = (width * 0.5 + 0.5, height * 0.5 + 0.5) # it is visually better to estimate the center without 0.5 offset # otherwise image rotated by 90 degrees is shifted vs output image of torch.rot90 or F_t.affine if center is None: center = [width * 0.5, height * 0.5] matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear) pil_interpolation = pil_modes_mapping[interpolation] return F_pil.affine(img, matrix=matrix, interpolation=pil_interpolation, fill=fill) center_f = [0.0, 0.0] if center is not None: _, height, width = get_dimensions(img) # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center. center_f = [1.0 * (c - s * 0.5) for c, s in zip(center, [width, height])] translate_f = [1.0 * t for t in translate] matrix = _get_inverse_affine_matrix(center_f, angle, translate_f, scale, shear) return F_t.affine(img, matrix=matrix, interpolation=interpolation.value, fill=fill) # @constexpr def to_grayscale(img, num_output_channels=1): """Convert PIL image of any mode (RGB, HSV, LAB, etc) to grayscale version of image. This transform does not support torch Tensor. Args: img (PIL Image): PIL Image to be converted to grayscale. num_output_channels (int): number of channels of the output image. Value can be 1 or 3. Default is 1. Returns: PIL Image: Grayscale version of the image. - if num_output_channels = 1 : returned image is single channel - if num_output_channels = 3 : returned image is 3 channel with r = g = b """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(to_grayscale) if isinstance(img, Image.Image): return F_pil.to_grayscale(img, num_output_channels) raise TypeError("Input should be PIL Image") def rgb_to_grayscale(img: Tensor, num_output_channels: int = 1) -> Tensor: """Convert RGB image to grayscale version of image. If the image is torch Tensor, it is expected to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions Note: Please, note that this method supports only RGB images as input. For inputs in other color spaces, please, consider using meth:`~torchvision.transforms.functional.to_grayscale` with PIL Image. Args: img (PIL Image or Tensor): RGB Image to be converted to grayscale. num_output_channels (int): number of channels of the output image. Value can be 1 or 3. Default, 1. Returns: PIL Image or Tensor: Grayscale version of the image. - if num_output_channels = 1 : returned image is single channel - if num_output_channels = 3 : returned image is 3 channel with r = g = b """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(rgb_to_grayscale) if not isinstance(img, torch.Tensor): return F_pil.to_grayscale(img, num_output_channels) return F_t.rgb_to_grayscale(img, num_output_channels) def erase(img: Tensor, i: int, j: int, h: int, w: int, v: Tensor, inplace: bool = False) -> Tensor: """Erase the input Tensor Image with given value. This transform does not support PIL Image. Args: img (Tensor Image): Tensor image of size (C, H, W) to be erased i (int): i in (i,j) i.e coordinates of the upper left corner. j (int): j in (i,j) i.e coordinates of the upper left corner. h (int): Height of the erased region. w (int): Width of the erased region. v: Erasing value. inplace(bool, optional): For in-place operations. By default is set False. Returns: Tensor Image: Erased image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(erase) if not isinstance(img, torch.Tensor): raise TypeError(f"img should be Tensor Image. Got {type(img)}") return F_t.erase(img, i, j, h, w, v, inplace=inplace) def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Tensor: """Performs Gaussian blurring on the image by given kernel. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. Args: img (PIL Image or Tensor): Image to be blurred kernel_size (sequence of ints or int): Gaussian kernel size. Can be a sequence of integers like ``(kx, ky)`` or a single integer for square kernels. .. note:: In torchscript mode kernel_size as single int is not supported, use a sequence of length 1: ``[ksize, ]``. sigma (sequence of floats or float, optional): Gaussian kernel standard deviation. Can be a sequence of floats like ``(sigma_x, sigma_y)`` or a single float to define the same sigma in both X/Y directions. If None, then it is computed using ``kernel_size`` as ``sigma = 0.3 * ((kernel_size - 1) * 0.5 - 1) + 0.8``. Default, None. .. note:: In torchscript mode sigma as single float is not supported, use a sequence of length 1: ``[sigma, ]``. Returns: PIL Image or Tensor: Gaussian Blurred version of the image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(gaussian_blur) if not isinstance(kernel_size, (int, list, tuple)): raise TypeError(f"kernel_size should be int or a sequence of integers. Got {type(kernel_size)}") if isinstance(kernel_size, int): kernel_size = [kernel_size, kernel_size] if len(kernel_size) != 2: raise ValueError(f"If kernel_size is a sequence its length should be 2. Got {len(kernel_size)}") for ksize in kernel_size: if ksize % 2 == 0 or ksize < 0: raise ValueError(f"kernel_size should have odd and positive integers. Got {kernel_size}") if sigma is None: sigma = [ksize * 0.15 + 0.35 for ksize in kernel_size] if sigma is not None and not isinstance(sigma, (int, float, list, tuple)): raise TypeError(f"sigma should be either float or sequence of floats. Got {type(sigma)}") if isinstance(sigma, (int, float)): sigma = [float(sigma), float(sigma)] if isinstance(sigma, (list, tuple)) and len(sigma) == 1: sigma = [sigma[0], sigma[0]] if len(sigma) != 2: raise ValueError(f"If sigma is a sequence, its length should be 2. Got {len(sigma)}") for s in sigma: if s <= 0.0: raise ValueError(f"sigma should have positive values. Got {sigma}") t_img = img if not isinstance(img, torch.Tensor): if not F_pil._is_pil_image(img): raise TypeError(f"img should be PIL Image or Tensor. Got {type(img)}") t_img = pil_to_tensor(img) output = F_t.gaussian_blur(t_img, kernel_size, sigma) if not isinstance(img, torch.Tensor): output = to_pil_image(output, mode=img.mode) return output def invert(img: Tensor) -> Tensor: """Invert the colors of an RGB/grayscale image. Args: img (PIL Image or Tensor): Image to have its colors inverted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Returns: PIL Image or Tensor: Color inverted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(invert) if not isinstance(img, torch.Tensor): return F_pil.invert(img) return F_t.invert(img) def posterize(img: Tensor, bits: int) -> Tensor: """Posterize an image by reducing the number of bits for each color channel. Args: img (PIL Image or Tensor): Image to have its colors posterized. If img is torch Tensor, it should be of type torch.uint8 and it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". bits (int): The number of bits to keep for each channel (0-8). Returns: PIL Image or Tensor: Posterized image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(posterize) if not (0 <= bits <= 8): raise ValueError(f"The number if bits should be between 0 and 8. Got {bits}") if not isinstance(img, torch.Tensor): return F_pil.posterize(img, bits) return F_t.posterize(img, bits) def solarize(img: Tensor, threshold: float) -> Tensor: """Solarize an RGB/grayscale image by inverting all pixel values above a threshold. Args: img (PIL Image or Tensor): Image to have its colors inverted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". threshold (float): All pixels equal or above this value are inverted. Returns: PIL Image or Tensor: Solarized image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(solarize) if not isinstance(img, torch.Tensor): return F_pil.solarize(img, threshold) return F_t.solarize(img, threshold) def adjust_sharpness(img: Tensor, sharpness_factor: float) -> Tensor: """Adjust the sharpness of an image. Args: img (PIL Image or Tensor): Image to be adjusted. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. sharpness_factor (float): How much to adjust the sharpness. Can be any non negative number. 0 gives a blurred image, 1 gives the original image while 2 increases the sharpness by a factor of 2. Returns: PIL Image or Tensor: Sharpness adjusted image. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(adjust_sharpness) if not isinstance(img, torch.Tensor): return F_pil.adjust_sharpness(img, sharpness_factor) return F_t.adjust_sharpness(img, sharpness_factor) def autocontrast(img: Tensor) -> Tensor: """Maximize contrast of an image by remapping its pixels per channel so that the lowest becomes black and the lightest becomes white. Args: img (PIL Image or Tensor): Image on which autocontrast is applied. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Returns: PIL Image or Tensor: An image that was autocontrasted. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(autocontrast) if not isinstance(img, torch.Tensor): return F_pil.autocontrast(img) return F_t.autocontrast(img) def equalize(img: Tensor) -> Tensor: """Equalize the histogram of an image by applying a non-linear mapping to the input in order to create a uniform distribution of grayscale values in the output. Args: img (PIL Image or Tensor): Image on which equalize is applied. If img is torch Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. The tensor dtype must be ``torch.uint8`` and values are expected to be in ``[0, 255]``. If img is PIL Image, it is expected to be in mode "P", "L" or "RGB". Returns: PIL Image or Tensor: An image that was equalized. """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(equalize) if not isinstance(img, torch.Tensor): return F_pil.equalize(img) return F_t.equalize(img) msadapter/ms_adapter/torchvision/transforms/functional_pil.py000066400000000000000000000316421436713554100253210ustar00rootroot00000000000000from __future__ import absolute_import import numbers from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import numpy as np from PIL import Image, ImageOps, ImageEnhance from typing_extensions import Literal from mindspore.ops import constexpr try: import accimage except ImportError: accimage = None from . import _pil_constants # @constexpr def _is_pil_image(img: Any) -> bool: if accimage is not None: return isinstance(img, (Image.Image, accimage.Image)) else: return isinstance(img, Image.Image) # @constexpr def get_dimensions(img: Any) -> List[int]: if _is_pil_image(img): if hasattr(img, "getbands"): channels = len(img.getbands()) else: channels = img.channels width, height = img.size return [channels, height, width] raise TypeError(f"Unexpected type {type(img)}") # @constexpr def get_image_size(img: Any) -> List[int]: if _is_pil_image(img): return list(img.size) raise TypeError(f"Unexpected type {type(img)}") # @constexpr def get_image_num_channels(img: Any) -> int: if _is_pil_image(img): if hasattr(img, "getbands"): return len(img.getbands()) else: return img.channels raise TypeError(f"Unexpected type {type(img)}") # @constexpr def hflip(img: Image.Image) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return img.transpose(_pil_constants.FLIP_LEFT_RIGHT) # @constexpr def vflip(img: Image.Image) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return img.transpose(_pil_constants.FLIP_TOP_BOTTOM) # @constexpr def adjust_brightness(img: Image.Image, brightness_factor: float) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") enhancer = ImageEnhance.Brightness(img) img = enhancer.enhance(brightness_factor) return img # @constexpr def adjust_contrast(img: Image.Image, contrast_factor: float) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") enhancer = ImageEnhance.Contrast(img) img = enhancer.enhance(contrast_factor) return img # @constexpr def adjust_saturation(img: Image.Image, saturation_factor: float) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") enhancer = ImageEnhance.Color(img) img = enhancer.enhance(saturation_factor) return img # @constexpr def adjust_hue(img: Image.Image, hue_factor: float) -> Image.Image: if not (-0.5 <= hue_factor <= 0.5): raise ValueError(f"hue_factor ({hue_factor}) is not in [-0.5, 0.5].") if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") input_mode = img.mode if input_mode in {"L", "1", "I", "F"}: return img h, s, v = img.convert("HSV").split() np_h = np.array(h, dtype=np.uint8) # uint8 addition take cares of rotation across boundaries with np.errstate(over="ignore"): np_h += np.uint8(hue_factor * 255) h = Image.fromarray(np_h, "L") img = Image.merge("HSV", (h, s, v)).convert(input_mode) return img # @constexpr def adjust_gamma( img: Image.Image, gamma: float, gain: float = 1.0, ) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") if gamma < 0: raise ValueError("Gamma should be a non-negative real number") input_mode = img.mode img = img.convert("RGB") gamma_map = [int((255 + 1 - 1e-3) * gain * pow(ele / 255.0, gamma)) for ele in range(256)] * 3 img = img.point(gamma_map) # use PIL's point-function to accelerate this part img = img.convert(input_mode) return img # @constexpr def pad( img: Image.Image, padding: Union[int, List[int], Tuple[int, ...]], fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0, padding_mode: Literal["constant", "edge", "reflect", "symmetric"] = "constant", ) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") if not isinstance(padding, (numbers.Number, tuple, list)): raise TypeError("Got inappropriate padding arg") if not isinstance(fill, (numbers.Number, tuple, list)): raise TypeError("Got inappropriate fill arg") if not isinstance(padding_mode, str): raise TypeError("Got inappropriate padding_mode arg") if isinstance(padding, list): padding = tuple(padding) if isinstance(padding, tuple) and len(padding) not in [1, 2, 4]: raise ValueError(f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple") if isinstance(padding, tuple) and len(padding) == 1: # Compatibility with `functional_tensor.pad` padding = padding[0] if padding_mode not in ["constant", "edge", "reflect", "symmetric"]: raise ValueError("Padding mode should be either constant, edge, reflect or symmetric") if padding_mode == "constant": opts = _parse_fill(fill, img, name="fill") if img.mode == "P": palette = img.getpalette() image = ImageOps.expand(img, border=padding, **opts) image.putpalette(palette) return image return ImageOps.expand(img, border=padding, **opts) else: if isinstance(padding, int): pad_left = pad_right = pad_top = pad_bottom = padding if isinstance(padding, tuple) and len(padding) == 2: pad_left = pad_right = padding[0] pad_top = pad_bottom = padding[1] if isinstance(padding, tuple) and len(padding) == 4: pad_left = padding[0] pad_top = padding[1] pad_right = padding[2] pad_bottom = padding[3] p = [pad_left, pad_top, pad_right, pad_bottom] cropping = -np.minimum(p, 0) if cropping.any(): crop_left, crop_top, crop_right, crop_bottom = cropping img = img.crop((crop_left, crop_top, img.width - crop_right, img.height - crop_bottom)) pad_left, pad_top, pad_right, pad_bottom = np.maximum(p, 0) if img.mode == "P": palette = img.getpalette() img = np.asarray(img) img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), mode=padding_mode) img = Image.fromarray(img) img.putpalette(palette) return img img = np.asarray(img) # RGB image if len(img.shape) == 3: img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), padding_mode) # Grayscale image if len(img.shape) == 2: img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), padding_mode) return Image.fromarray(img) # @constexpr def crop( img: Image.Image, top: int, left: int, height: int, width: int, ) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return img.crop((left, top, left + width, top + height)) # @constexpr def resize( img: Image.Image, size: Union[Sequence[int], int], interpolation: int = _pil_constants.BILINEAR, max_size: Optional[int] = None, ) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") if not (isinstance(size, int) or (isinstance(size, Sequence) and len(size) in (1, 2))): raise TypeError(f"Got inappropriate size arg: {size}") if isinstance(size, Sequence) and len(size) == 1: size = size[0] if isinstance(size, int): w, h = img.size short, long = (w, h) if w <= h else (h, w) new_short, new_long = size, int(size * long / short) if max_size is not None: if max_size <= size: raise ValueError( f"max_size = {max_size} must be strictly greater than the requested " f"size for the smaller edge size = {size}" ) if new_long > max_size: new_short, new_long = int(max_size * new_short / new_long), max_size new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short) if (w, h) == (new_w, new_h): return img else: return img.resize((new_w, new_h), interpolation) else: if max_size is not None: raise ValueError( "max_size should only be passed if size specifies the length of the smaller edge, " "i.e. size should be an int or a sequence of length 1 in torchscript mode." ) return img.resize(size[::-1], interpolation) # @constexpr def _parse_fill( fill: Optional[Union[float, List[float], Tuple[float, ...]]], img: Image.Image, name: str = "fillcolor", ) -> Dict[str, Optional[Union[float, List[float], Tuple[float, ...]]]]: # Process fill color for affine transforms num_bands = len(img.getbands()) if fill is None: fill = 0 if isinstance(fill, (int, float)) and num_bands > 1: fill = tuple([fill] * num_bands) if isinstance(fill, (list, tuple)): if len(fill) != num_bands: msg = "The number of elements in 'fill' does not match the number of bands of the image ({} != {})" raise ValueError(msg.format(len(fill), num_bands)) fill = tuple(fill) if img.mode != "F": if isinstance(fill, (list, tuple)): fill = tuple(int(x) for x in fill) else: fill = int(fill) return {name: fill} # @constexpr def affine( img: Image.Image, matrix: List[float], interpolation: int = _pil_constants.NEAREST, fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0, ) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") output_size = img.size opts = _parse_fill(fill, img) return img.transform(output_size, _pil_constants.AFFINE, matrix, interpolation, **opts) # @constexpr def rotate( img: Image.Image, angle: float, interpolation: int = _pil_constants.NEAREST, expand: bool = False, center: Optional[Tuple[int, int]] = None, fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0, ) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") opts = _parse_fill(fill, img) return img.rotate(angle, interpolation, expand, center, **opts) # @constexpr def perspective( img: Image.Image, perspective_coeffs: float, interpolation: int = _pil_constants.BICUBIC, fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0, ) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") opts = _parse_fill(fill, img) return img.transform(img.size, _pil_constants.PERSPECTIVE, perspective_coeffs, interpolation, **opts) # @constexpr def to_grayscale(img: Image.Image, num_output_channels: int) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") if num_output_channels == 1: img = img.convert("L") elif num_output_channels == 3: img = img.convert("L") np_img = np.array(img, dtype=np.uint8) np_img = np.dstack([np_img, np_img, np_img]) img = Image.fromarray(np_img, "RGB") else: raise ValueError("num_output_channels should be either 1 or 3") return img # @constexpr def invert(img: Image.Image) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return ImageOps.invert(img) # @constexpr def posterize(img: Image.Image, bits: int) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return ImageOps.posterize(img, bits) # @constexpr def solarize(img: Image.Image, threshold: int) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return ImageOps.solarize(img, threshold) # @constexpr def adjust_sharpness(img: Image.Image, sharpness_factor: float) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") enhancer = ImageEnhance.Sharpness(img) img = enhancer.enhance(sharpness_factor) return img # @constexpr def autocontrast(img: Image.Image) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return ImageOps.autocontrast(img) # @constexpr def equalize(img: Image.Image) -> Image.Image: if not _is_pil_image(img): raise TypeError(f"img should be PIL Image. Got {type(img)}") return ImageOps.equalize(img) msadapter/ms_adapter/torchvision/transforms/functional_tensor.py000066400000000000000000001110571436713554100260460ustar00rootroot00000000000000from __future__ import absolute_import import warnings from typing import Optional, Tuple, List # import torch # from torch import Tensor # from torch.nn.functional import grid_sample, conv2d, interpolate, pad as torch_pad import mindspore as ms import ms_adapter.pytorch as torch from ms_adapter.pytorch import Tensor from ms_adapter.pytorch.nn.functional import grid_sample, conv2d, interpolate, pad as torch_pad from ms_adapter.pytorch.tensor import cast_to_adapter_tensor def _is_tensor_a_torch_image(x: Tensor) -> bool: return x.ndim >= 2 def _assert_image_tensor(img: Tensor) -> None: if not _is_tensor_a_torch_image(img): raise TypeError("Tensor is not a torch image.") def _assert_threshold(img: Tensor, threshold: float) -> None: bound = 1 if img.is_floating_point() else 255 if threshold > bound: raise TypeError("Threshold should be less than bound of img.") def get_dimensions(img: Tensor) -> List[int]: _assert_image_tensor(img) channels = 1 if img.ndim == 2 else img.shape[-3] height, width = img.shape[-2:] return [channels, height, width] def get_image_size(img: Tensor) -> List[int]: # Returns (w, h) of tensor image _assert_image_tensor(img) return [img.shape[-1], img.shape[-2]] def get_image_num_channels(img: Tensor) -> int: _assert_image_tensor(img) if img.ndim == 2: return 1 elif img.ndim > 2: return img.shape[-3] raise TypeError(f"Input ndim should be 2 or more. Got {img.ndim}") def _max_value(dtype) -> int: if dtype == torch.uint8: return 255 elif dtype == torch.int8: return 127 elif dtype == torch.int16: return 32767 elif dtype == torch.int32: return 2147483647 elif dtype == torch.int64: return 9223372036854775807 else: return 1 def _assert_channels(img: Tensor, permitted: List[int]) -> None: c = get_dimensions(img)[0] if c not in permitted: raise TypeError(f"Input image tensor permitted channel values are {permitted}, but found {c}") def convert_image_dtype(image: torch.Tensor, dtype = torch.float) -> torch.Tensor: if image.dtype == dtype: return image if image.is_floating_point(): # TODO: replace with dtype.is_floating_point when torchscript supports it if torch.tensor(0, dtype=dtype).is_floating_point(): return image.to(dtype) # float to int if (image.dtype == torch.float32 and dtype in (torch.int32, torch.int64)) or ( image.dtype == torch.float64 and dtype == torch.int64 ): msg = f"The cast from {image.dtype} to {dtype} cannot be performed safely." raise RuntimeError(msg) # https://github.com/pytorch/vision/pull/2078#issuecomment-612045321 # For data in the range 0-1, (float * 255).to(uint) is only 255 # when float is exactly 1.0. # `max + 1 - epsilon` provides more evenly distributed mapping of # ranges of floats to ints. eps = 1e-3 max_val = float(_max_value(dtype)) result = image.mul(max_val + 1.0 - eps) return result.to(dtype) else: input_max = float(_max_value(image.dtype)) # int to float # TODO: replace with dtype.is_floating_point when torchscript supports it if torch.tensor(0, dtype=dtype).is_floating_point(): image = image.to(dtype) return image / input_max output_max = float(_max_value(dtype)) # int to int if input_max > output_max: # factor should be forced to int for torch jit script # otherwise factor is a float and image // factor can produce different results factor = int((input_max + 1) // (output_max + 1)) image = torch.div(image, factor, rounding_mode="floor") return image.to(dtype) else: # factor should be forced to int for torch jit script # otherwise factor is a float and image * factor can produce different results factor = int((output_max + 1) // (input_max + 1)) image = image.to(dtype) return image * factor def vflip(img: Tensor) -> Tensor: _assert_image_tensor(img) return img.flip(-2) def hflip(img: Tensor) -> Tensor: _assert_image_tensor(img) return img.flip(-1) def crop(img: Tensor, top: int, left: int, height: int, width: int) -> Tensor: _assert_image_tensor(img) _, h, w = get_dimensions(img) right = left + width bottom = top + height if left < 0 or top < 0 or right > w or bottom > h: padding_ltrb = [max(-left, 0), max(-top, 0), max(right - w, 0), max(bottom - h, 0)] return pad(img[..., max(top, 0) : bottom, max(left, 0) : right], padding_ltrb, fill=0) return img[..., top:bottom, left:right] def rgb_to_grayscale(img: Tensor, num_output_channels: int = 1) -> Tensor: if img.ndim < 3: raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") _assert_channels(img, [3]) if num_output_channels not in (1, 3): raise ValueError("num_output_channels should be either 1 or 3") r, g, b = img.unbind(dim=-3) # This implementation closely follows the TF one: # https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/ops/image_ops_impl.py#L2105-L2138 l_img = (0.2989 * r + 0.587 * g + 0.114 * b).to(img.dtype) l_img = l_img.unsqueeze(dim=-3) if num_output_channels == 3: return l_img.expand(img.shape) return l_img def adjust_brightness(img: Tensor, brightness_factor: float) -> Tensor: if brightness_factor < 0: raise ValueError(f"brightness_factor ({brightness_factor}) is not non-negative.") _assert_image_tensor(img) _assert_channels(img, [1, 3]) return _blend(img, torch.zeros_like(img), brightness_factor) def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor: if contrast_factor < 0: raise ValueError(f"contrast_factor ({contrast_factor}) is not non-negative.") _assert_image_tensor(img) _assert_channels(img, [3, 1]) c = get_dimensions(img)[0] dtype = img.dtype if torch.is_floating_point(img) else torch.float32 if c == 3: mean = torch.mean(rgb_to_grayscale(img).to(dtype), dim=(-3, -2, -1), keepdim=True) else: mean = torch.mean(img.to(dtype), dim=(-3, -2, -1), keepdim=True) return _blend(img, mean, contrast_factor) def adjust_hue(img: Tensor, hue_factor: float) -> Tensor: if not (-0.5 <= hue_factor <= 0.5): raise ValueError(f"hue_factor ({hue_factor}) is not in [-0.5, 0.5].") if not (isinstance(img, torch.Tensor)): raise TypeError("Input img should be Tensor image") _assert_image_tensor(img) _assert_channels(img, [1, 3]) if get_dimensions(img)[0] == 1: # Match PIL behaviour return img orig_dtype = img.dtype if img.dtype == torch.uint8: img = img.to(dtype=torch.float32) / 255.0 img = _rgb2hsv(img) h, s, v = img.unbind(dim=-3) h = (h + hue_factor) % 1.0 img = torch.stack((h, s, v), dim=-3) img_hue_adj = _hsv2rgb(img) if orig_dtype == torch.uint8: img_hue_adj = (img_hue_adj * 255.0).to(dtype=orig_dtype) return img_hue_adj def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor: if saturation_factor < 0: raise ValueError(f"saturation_factor ({saturation_factor}) is not non-negative.") _assert_image_tensor(img) _assert_channels(img, [1, 3]) if get_dimensions(img)[0] == 1: # Match PIL behaviour return img return _blend(img, rgb_to_grayscale(img), saturation_factor) def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor: if not isinstance(img, torch.Tensor): raise TypeError("Input img should be a Tensor.") _assert_channels(img, [1, 3]) if gamma < 0: raise ValueError("Gamma should be a non-negative real number") result = img dtype = img.dtype if not torch.is_floating_point(img): result = convert_image_dtype(result, torch.float32) result = (gain * result ** gamma).clamp(0, 1) result = convert_image_dtype(result, dtype) return result def _blend(img1: Tensor, img2: Tensor, ratio: float) -> Tensor: ratio = float(ratio) bound = 1.0 if img1.is_floating_point() else 255.0 return (ratio * img1 + (1.0 - ratio) * img2).clamp(0, bound).to(img1.dtype) # TODO: # def _rgb2hsv(img: Tensor) -> Tensor: # r, g, b = img.unbind(dim=-3) # # # Implementation is based on https://github.com/python-pillow/Pillow/blob/4174d4267616897df3746d315d5a2d0f82c656ee/ # # src/libImaging/Convert.c#L330 # maxc = torch.max(img, dim=-3).values # minc = torch.min(img, dim=-3).values # # # The algorithm erases S and H channel where `maxc = minc`. This avoids NaN # # from happening in the results, because # # + S channel has division by `maxc`, which is zero only if `maxc = minc` # # + H channel has division by `(maxc - minc)`. # # # # Instead of overwriting NaN afterwards, we just prevent it from occuring so # # we don't need to deal with it in case we save the NaN in a buffer in # # backprop, if it is ever supported, but it doesn't hurt to do so. # eqc = maxc == minc # # cr = maxc - minc # # Since `eqc => cr = 0`, replacing denominator with 1 when `eqc` is fine. # ones = torch.ones_like(maxc) # s = cr / torch.where(eqc, ones, maxc) # # Note that `eqc => maxc = minc = r = g = b`. So the following calculation # # of `h` would reduce to `bc - gc + 2 + rc - bc + 4 + rc - bc = 6` so it # # would not matter what values `rc`, `gc`, and `bc` have here, and thus # # replacing denominator with 1 when `eqc` is fine. # cr_divisor = torch.where(eqc, ones, cr) # rc = (maxc - r) / cr_divisor # gc = (maxc - g) / cr_divisor # bc = (maxc - b) / cr_divisor # # hr = (maxc == r) * (bc - gc) # hg = ((maxc == g) & (maxc != r)) * (2.0 + rc - bc) # hb = ((maxc != g) & (maxc != r)) * (4.0 + gc - rc) # h = hr + hg + hb # h = torch.fmod((h / 6.0 + 1.0), 1.0) # return torch.stack((h, s, maxc), dim=-3) # TODO: # def _hsv2rgb(img: Tensor) -> Tensor: # h, s, v = img.unbind(dim=-3) # i = torch.floor(h * 6.0) # f = (h * 6.0) - i # i = i.to(dtype=torch.int32) # # p = torch.clamp((v * (1.0 - s)), 0.0, 1.0) # q = torch.clamp((v * (1.0 - s * f)), 0.0, 1.0) # t = torch.clamp((v * (1.0 - s * (1.0 - f))), 0.0, 1.0) # i = i % 6 # # mask = i.unsqueeze(dim=-3) == torch.arange(0, 6, device=None).view(-1, 1, 1) # # a1 = torch.stack((v, q, p, p, t, v), dim=-3) # a2 = torch.stack((t, v, v, q, p, p), dim=-3) # a3 = torch.stack((p, p, t, v, v, q), dim=-3) # a4 = torch.stack((a1, a2, a3), dim=-4) # # return torch.einsum("...ijk, ...xijk -> ...xjk", mask.to(dtype=img.dtype), a4) def _rgb2hsv(img: Tensor) -> Tensor: img_shape = img.shape img_rank = len(img_shape) if img_rank == 3: img = img.transpose(0, 2) img = img.unsqueeze(dim=0) # TODO: -> ms.ops.RGBToHSV out = ms.ops.operations.image_ops.RGBToHSV()(img) return cast_to_adapter_tensor(out).transpose(3, 1).view(img_shape) if img_rank == 4: img = img.transpose(1, 3) # TODO: -> ms.ops.RGBToHSV out = ms.ops.operations.image_ops.RGBToHSV()(img) return cast_to_adapter_tensor(out).transpose(3, 1) def _hsv2rgb(img: Tensor) -> Tensor: img_shape = img.shape img_rank = len(img_shape) if img_rank == 3: img = img.transpose(0, 2) img = img.unsqueeze(dim=0) out = ms.ops.HSVToRGB()(img) return cast_to_adapter_tensor(out).transpose(3, 1).view(img_shape) if img_rank == 4: img = img.transpose(1, 3) out = ms.ops.HSVToRGB()(img) return cast_to_adapter_tensor(out).transpose(3, 1) def _pad_symmetric(img: Tensor, padding: List[int]) -> Tensor: # padding is left, right, top, bottom # crop if needed if padding[0] < 0 or padding[1] < 0 or padding[2] < 0 or padding[3] < 0: neg_min_padding = [-min(x, 0) for x in padding] crop_left, crop_right, crop_top, crop_bottom = neg_min_padding img = img[..., crop_top : img.shape[-2] - crop_bottom, crop_left : img.shape[-1] - crop_right] padding = [max(x, 0) for x in padding] in_sizes = img.size() _x_indices = [i for i in range(in_sizes[-1])] # [0, 1, 2, 3, ...] left_indices = [i for i in range(padding[0] - 1, -1, -1)] # e.g. [3, 2, 1, 0] right_indices = [-(i + 1) for i in range(padding[1])] # e.g. [-1, -2, -3] temp_x_indices = left_indices + _x_indices + right_indices temp_x_indices = [i if i >= 0 else i + in_sizes[-1] for i in temp_x_indices] x_indices = torch.tensor(temp_x_indices, device=None) _y_indices = [i for i in range(in_sizes[-2])] top_indices = [i for i in range(padding[2] - 1, -1, -1)] bottom_indices = [-(i + 1) for i in range(padding[3])] temp_y_indices = top_indices + _y_indices + bottom_indices temp_y_indices = [i if i >= 0 else i + in_sizes[-2] for i in temp_y_indices] y_indices = torch.tensor(temp_y_indices, device=None) ndim = img.ndim if ndim == 3: return img[:, y_indices[:, None], x_indices[None, :]] elif ndim == 4: return img[:, :, y_indices[:, None], x_indices[None, :]] else: raise RuntimeError("Symmetric padding of N-D tensors are not supported yet") def _parse_pad_padding(padding: List[int]) -> List[int]: if isinstance(padding, int): # if torch.jit.is_scripting(): # # This maybe unreachable # raise ValueError("padding can't be an int while torchscripting, set it as a list [value, ]") pad_left = pad_right = pad_top = pad_bottom = padding elif len(padding) == 1: pad_left = pad_right = pad_top = pad_bottom = padding[0] elif len(padding) == 2: pad_left = pad_right = padding[0] pad_top = pad_bottom = padding[1] else: pad_left = padding[0] pad_top = padding[1] pad_right = padding[2] pad_bottom = padding[3] return [pad_left, pad_right, pad_top, pad_bottom] def pad(img: Tensor, padding: List[int], fill: int = 0, padding_mode: str = "constant") -> Tensor: _assert_image_tensor(img) if not isinstance(padding, (int, tuple, list)): raise TypeError("Got inappropriate padding arg") if not isinstance(fill, (int, float)): raise TypeError("Got inappropriate fill arg") if not isinstance(padding_mode, str): raise TypeError("Got inappropriate padding_mode arg") if isinstance(padding, tuple): padding = list(padding) if isinstance(padding, list) and len(padding) not in [1, 2, 4]: raise ValueError(f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple") if padding_mode not in ["constant", "edge", "reflect", "symmetric"]: raise ValueError("Padding mode should be either constant, edge, reflect or symmetric") p = _parse_pad_padding(padding) if padding_mode == "edge": # remap padding_mode str padding_mode = "replicate" elif padding_mode == "symmetric": # route to another implementation return _pad_symmetric(img, p) need_squeeze = False if img.ndim < 4: img = img.unsqueeze(dim=0) need_squeeze = True out_dtype = img.dtype need_cast = False if (padding_mode != "constant") and img.dtype not in (torch.float32, torch.float64): # Here we temporary cast input tensor to float # until pytorch issue is resolved : # https://github.com/pytorch/pytorch/issues/40763 need_cast = True img = img.to(torch.float32) if padding_mode in ("reflect", "replicate"): img = torch_pad(img, p, mode=padding_mode) else: img = torch_pad(img, p, mode=padding_mode, value=float(fill)) if need_squeeze: img = img.squeeze(dim=0) if need_cast: img = img.to(out_dtype) return img def resize( img: Tensor, size: List[int], interpolation: str = "bilinear", max_size: Optional[int] = None, antialias: Optional[bool] = None, ) -> Tensor: _assert_image_tensor(img) if not isinstance(size, (int, tuple, list)): raise TypeError("Got inappropriate size arg") if not isinstance(interpolation, str): raise TypeError("Got inappropriate interpolation arg") if interpolation not in ["nearest", "bilinear", "bicubic"]: raise ValueError("This interpolation mode is unsupported with Tensor input") if isinstance(size, tuple): size = list(size) if isinstance(size, list): if len(size) not in [1, 2]: raise ValueError( f"Size must be an int or a 1 or 2 element tuple/list, not a {len(size)} element tuple/list" ) if max_size is not None and len(size) != 1: raise ValueError( "max_size should only be passed if size specifies the length of the smaller edge, " "i.e. size should be an int or a sequence of length 1 in torchscript mode." ) if antialias is None: antialias = False if antialias and interpolation not in ["bilinear", "bicubic"]: raise ValueError("Antialias option is supported for bilinear and bicubic interpolation modes only") _, h, w = get_dimensions(img) if isinstance(size, int) or len(size) == 1: # specified size only for the smallest edge short, long = (w, h) if w <= h else (h, w) requested_new_short = size if isinstance(size, int) else size[0] new_short, new_long = requested_new_short, int(requested_new_short * long / short) if max_size is not None: if max_size <= requested_new_short: raise ValueError( f"max_size = {max_size} must be strictly greater than the requested " f"size for the smaller edge size = {size}" ) if new_long > max_size: new_short, new_long = int(max_size * new_short / new_long), max_size new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short) if (w, h) == (new_w, new_h): return img else: # specified both h and w new_w, new_h = size[1], size[0] img, need_cast, need_squeeze, out_dtype = _cast_squeeze_in(img, [torch.float32, torch.float64]) # Define align_corners to avoid warnings align_corners = False if interpolation in ["bilinear", "bicubic"] else None img = interpolate(img, size=[new_h, new_w], mode=interpolation, align_corners=align_corners, antialias=antialias) if interpolation == "bicubic" and out_dtype == torch.uint8: img = img.clamp(min=0, max=255) img = _cast_squeeze_out(img, need_cast=need_cast, need_squeeze=need_squeeze, out_dtype=out_dtype) return img def _assert_grid_transform_inputs( img: Tensor, matrix: Optional[List[float]], interpolation: str, fill: Optional[List[float]], supported_interpolation_modes: List[str], coeffs: Optional[List[float]] = None, ) -> None: if not (isinstance(img, torch.Tensor)): raise TypeError("Input img should be Tensor") _assert_image_tensor(img) if matrix is not None and not isinstance(matrix, list): raise TypeError("Argument matrix should be a list") if matrix is not None and len(matrix) != 6: raise ValueError("Argument matrix should have 6 float values") if coeffs is not None and len(coeffs) != 8: raise ValueError("Argument coeffs should have 8 float values") if fill is not None and not isinstance(fill, (int, float, tuple, list)): warnings.warn("Argument fill should be either int, float, tuple or list") # Check fill num_channels = get_dimensions(img)[0] if isinstance(fill, (tuple, list)) and (len(fill) > 1 and len(fill) != num_channels): msg = ( "The number of elements in 'fill' cannot broadcast to match the number of " "channels of the image ({} != {})" ) raise ValueError(msg.format(len(fill), num_channels)) if interpolation not in supported_interpolation_modes: raise ValueError(f"Interpolation mode '{interpolation}' is unsupported with Tensor input") def _cast_squeeze_in(img: Tensor, req_dtypes): need_squeeze = False # make image NCHW if img.ndim < 4: img = img.unsqueeze(dim=0) need_squeeze = True out_dtype = img.dtype need_cast = False if out_dtype not in req_dtypes: need_cast = True req_dtype = req_dtypes[0] img = img.to(req_dtype) return img, need_cast, need_squeeze, out_dtype def _cast_squeeze_out(img: Tensor, need_cast: bool, need_squeeze: bool, out_dtype) -> Tensor: if need_squeeze: img = img.squeeze(dim=0) if need_cast: if out_dtype in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64): # it is better to round before cast img = torch.round(img) img = img.to(out_dtype) return img def _apply_grid_transform(img: Tensor, grid: Tensor, mode: str, fill: Optional[List[float]]) -> Tensor: img, need_cast, need_squeeze, out_dtype = _cast_squeeze_in(img, [grid.dtype]) if img.shape[0] > 1: # Apply same grid to a batch of images grid = grid.expand(img.shape[0], grid.shape[1], grid.shape[2], grid.shape[3]) # Append a dummy mask for customized fill colors, should be faster than grid_sample() twice if fill is not None: dummy = torch.ones((img.shape[0], 1, img.shape[2], img.shape[3]), dtype=img.dtype, device=None) img = torch.cat((img, dummy), dim=1) img = grid_sample(img, grid, mode=mode, padding_mode="zeros", align_corners=False) # Fill with required color if fill is not None: mask = img[:, -1:, :, :] # N * 1 * H * W img = img[:, :-1, :, :] # N * C * H * W mask = mask.expand_as(img) len_fill = len(fill) if isinstance(fill, (tuple, list)) else 1 fill_img = torch.tensor(fill, dtype=img.dtype, device=None).view(1, len_fill, 1, 1).expand_as(img) if mode == "nearest": mask = mask < 0.5 # img[mask] = fill_img[mask] img = cast_to_adapter_tensor(ms.numpy.where(mask, x=fill_img, y=img)) else: # 'bilinear' img = img * mask + (1.0 - mask) * fill_img img = _cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) return img def _gen_affine_grid( theta: Tensor, w: int, h: int, ow: int, oh: int, ) -> Tensor: # https://github.com/pytorch/pytorch/blob/74b65c32be68b15dc7c9e8bb62459efbfbde33d8/aten/src/ATen/native/ # AffineGridGenerator.cpp#L18 # Difference with AffineGridGenerator is that: # 1) we normalize grid values after applying theta # 2) we can normalize by other image size, such that it covers "extend" option like in PIL.Image.rotate # d = 0.5 # base_grid = torch.empty(1, oh, ow, 3, dtype=theta.dtype, device=theta.device) # x_grid = torch.linspace(-ow * 0.5 + d, ow * 0.5 + d - 1, steps=ow, device=theta.device) # base_grid[..., 0].copy_(x_grid) # y_grid = torch.linspace(-oh * 0.5 + d, oh * 0.5 + d - 1, steps=oh, device=theta.device).unsqueeze_(-1) # base_grid[..., 1].copy_(y_grid) # base_grid[..., 2].fill_(1) d = 0.5 x_grid = torch.linspace(-ow * 0.5 + d, ow * 0.5 + d - 1, steps=ow, device=None) x_grid = ms.ops.broadcast_to(x_grid, (1, oh, ow)) y_grid = torch.linspace(-oh * 0.5 + d, oh * 0.5 + d - 1, steps=oh, device=None) y_grid = y_grid.unsqueeze(-1) y_grid = ms.ops.broadcast_to(y_grid, (1, oh, ow)) z_grid = ms.ops.ones((1, oh, ow), ms.float32) base_grid = ms.ops.stack((x_grid, y_grid, z_grid), 3).to(theta.dtype) rescaled_theta = theta.transpose(1, 2) / torch.tensor([0.5 * w, 0.5 * h], dtype=theta.dtype, device=None) output_grid = base_grid.view(1, oh * ow, 3).bmm(rescaled_theta) output_grid = cast_to_adapter_tensor(output_grid) return output_grid.view(1, oh, ow, 2) def affine( img: Tensor, matrix: List[float], interpolation: str = "nearest", fill: Optional[List[float]] = None ) -> Tensor: _assert_grid_transform_inputs(img, matrix, interpolation, fill, ["nearest", "bilinear"]) dtype = img.dtype if torch.is_floating_point(img) else torch.float32 theta = torch.tensor(matrix, dtype=dtype, device=None).reshape(1, 2, 3) shape = img.shape # grid will be generated on the same device as theta and img grid = _gen_affine_grid(theta, w=shape[-1], h=shape[-2], ow=shape[-1], oh=shape[-2]) return _apply_grid_transform(img, grid, interpolation, fill=fill) def _compute_output_size(matrix: List[float], w: int, h: int) -> Tuple[int, int]: # Inspired of PIL implementation: # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054 # pts are Top-Left, Top-Right, Bottom-Left, Bottom-Right points. # Points are shifted due to affine matrix torch convention about # the center point. Center is (0, 0) for image center pivot point (w * 0.5, h * 0.5) pts = torch.tensor( [ [-0.5 * w, -0.5 * h, 1.0], [-0.5 * w, 0.5 * h, 1.0], [0.5 * w, 0.5 * h, 1.0], [0.5 * w, -0.5 * h, 1.0], ] ) theta = torch.tensor(matrix, dtype=torch.float).view(2, 3) new_pts = torch.matmul(pts, theta.T) min_vals, _ = new_pts.min(dim=0) max_vals, _ = new_pts.max(dim=0) # shift points to [0, w] and [0, h] interval to match PIL results min_vals += torch.tensor((w * 0.5, h * 0.5)) max_vals += torch.tensor((w * 0.5, h * 0.5)) # Truncate precision to 1e-4 to avoid ceil of Xe-15 to 1.0 tol = 1e-4 cmax = torch.ceil((max_vals / tol).trunc() * tol) cmin = torch.floor((min_vals / tol).trunc() * tol) size = cmax - cmin return int(size[0]), int(size[1]) def rotate( img: Tensor, matrix: List[float], interpolation: str = "nearest", expand: bool = False, fill: Optional[List[float]] = None, ) -> Tensor: _assert_grid_transform_inputs(img, matrix, interpolation, fill, ["nearest", "bilinear"]) w, h = img.shape[-1], img.shape[-2] ow, oh = _compute_output_size(matrix, w, h) if expand else (w, h) dtype = img.dtype if torch.is_floating_point(img) else torch.float32 theta = torch.tensor(matrix, dtype=dtype, device=None).reshape(1, 2, 3) # grid will be generated on the same device as theta and img grid = _gen_affine_grid(theta, w=w, h=h, ow=ow, oh=oh) return _apply_grid_transform(img, grid, interpolation, fill=fill) def _perspective_grid(coeffs: List[float], ow: int, oh: int, dtype, device) -> Tensor: # https://github.com/python-pillow/Pillow/blob/4634eafe3c695a014267eefdce830b4a825beed7/ # src/libImaging/Geometry.c#L394 # # x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1) # y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1) # theta1 = torch.tensor( [[[coeffs[0], coeffs[1], coeffs[2]], [coeffs[3], coeffs[4], coeffs[5]]]], dtype=dtype, device=device ) theta2 = torch.tensor([[[coeffs[6], coeffs[7], 1.0], [coeffs[6], coeffs[7], 1.0]]], dtype=dtype, device=device) # d = 0.5 # base_grid = torch.empty(1, oh, ow, 3, dtype=dtype, device=device) # x_grid = torch.linspace(d, ow * 1.0 + d - 1.0, steps=ow, device=device) # base_grid[..., 0].copy_(x_grid) # y_grid = torch.linspace(d, oh * 1.0 + d - 1.0, steps=oh, device=device).unsqueeze_(-1) # base_grid[..., 1].copy_(y_grid) # base_grid[..., 2].fill_(1) d = 0.5 x_grid = torch.linspace(d, ow * 1.0 + d - 1.0, steps=ow, device=device) x_grid = ms.ops.broadcast_to(x_grid, (1, oh, ow)) y_grid = torch.linspace(d, oh * 1.0 + d - 1.0, steps=oh, device=device) y_grid = y_grid.unsqueeze(-1) y_grid = ms.ops.broadcast_to(y_grid, (1, oh, ow)) z_grid = ms.ops.ones((1, oh, ow), ms.float32) base_grid = ms.ops.stack((x_grid, y_grid, z_grid), 3).to(dtype) rescaled_theta1 = theta1.transpose(1, 2) / torch.tensor([0.5 * ow, 0.5 * oh], dtype=dtype, device=device) output_grid1 = base_grid.view(1, oh * ow, 3).bmm(rescaled_theta1) output_grid2 = base_grid.view(1, oh * ow, 3).bmm(theta2.transpose(1, 2)) output_grid = output_grid1 / output_grid2 - 1.0 output_grid = cast_to_adapter_tensor(output_grid) return output_grid.view(1, oh, ow, 2) def perspective( img: Tensor, perspective_coeffs: List[float], interpolation: str = "bilinear", fill: Optional[List[float]] = None ) -> Tensor: if not (isinstance(img, torch.Tensor)): raise TypeError("Input img should be Tensor.") _assert_image_tensor(img) _assert_grid_transform_inputs( img, matrix=None, interpolation=interpolation, fill=fill, supported_interpolation_modes=["nearest", "bilinear"], coeffs=perspective_coeffs, ) ow, oh = img.shape[-1], img.shape[-2] dtype = img.dtype if torch.is_floating_point(img) else torch.float32 grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=None) return _apply_grid_transform(img, grid, interpolation, fill=fill) def _get_gaussian_kernel1d(kernel_size: int, sigma: float) -> Tensor: ksize_half = (kernel_size - 1) * 0.5 x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size) pdf = torch.exp(-0.5 * (x / sigma).pow(2)) kernel1d = pdf / pdf.sum() return kernel1d def _get_gaussian_kernel2d( kernel_size: List[int], sigma: List[float], dtype, device ) -> Tensor: kernel1d_x = _get_gaussian_kernel1d(kernel_size[0], sigma[0]).to(device, dtype=dtype) kernel1d_y = _get_gaussian_kernel1d(kernel_size[1], sigma[1]).to(device, dtype=dtype) kernel2d = torch.mm(kernel1d_y[:, None], kernel1d_x[None, :]) return kernel2d def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: List[float]) -> Tensor: if not (isinstance(img, torch.Tensor)): raise TypeError(f"img should be Tensor. Got {type(img)}") _assert_image_tensor(img) dtype = img.dtype if torch.is_floating_point(img) else torch.float32 kernel = _get_gaussian_kernel2d(kernel_size, sigma, dtype=dtype, device=None) kernel = kernel.expand(img.shape[-3], 1, kernel.shape[0], kernel.shape[1]) img, need_cast, need_squeeze, out_dtype = _cast_squeeze_in( img, [ kernel.dtype, ], ) # padding = (left, right, top, bottom) padding = [kernel_size[0] // 2, kernel_size[0] // 2, kernel_size[1] // 2, kernel_size[1] // 2] img = torch_pad(img, padding, mode="reflect") img = conv2d(img, kernel, groups=img.shape[-3]) img = _cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) return img def invert(img: Tensor) -> Tensor: _assert_image_tensor(img) if img.ndim < 3: raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") _assert_channels(img, [1, 3]) bound = torch.tensor(1 if img.is_floating_point() else 255, dtype=img.dtype, device=None) return bound - img def posterize(img: Tensor, bits: int) -> Tensor: _assert_image_tensor(img) if img.ndim < 3: raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") if img.dtype != torch.uint8: raise TypeError(f"Only torch.uint8 image tensors are supported, but found {img.dtype}") _assert_channels(img, [1, 3]) mask = -int(2 ** (8 - bits)) # JIT-friendly for: ~(2 ** (8 - bits) - 1) return img & mask def solarize(img: Tensor, threshold: float) -> Tensor: _assert_image_tensor(img) if img.ndim < 3: raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") _assert_channels(img, [1, 3]) _assert_threshold(img, threshold) inverted_img = invert(img) out = torch.where(img >= threshold, inverted_img, img) return torch.where(img >= threshold, inverted_img, img) def _blurred_degenerate_image(img: Tensor) -> Tensor: dtype = img.dtype if torch.is_floating_point(img) else torch.float32 kernel = torch.ones((3, 3), dtype=dtype, device=None) kernel[1, 1] = 5.0 kernel /= kernel.sum() kernel = kernel.expand(img.shape[-3], 1, kernel.shape[0], kernel.shape[1]) result_tmp, need_cast, need_squeeze, out_dtype = _cast_squeeze_in( img, [ kernel.dtype, ], ) result_tmp = conv2d(result_tmp, kernel, groups=result_tmp.shape[-3]) result_tmp = _cast_squeeze_out(result_tmp, need_cast, need_squeeze, out_dtype) result = img.clone() # TODO: # result[..., 1:-1, 1:-1] = result_tmp # return result result_dtype = result.dtype result = result.to(torch.float32) result[..., 1:-1, 1:-1] = result_tmp.to(torch.float32) return result.to(result_dtype) def adjust_sharpness(img: Tensor, sharpness_factor: float) -> Tensor: if sharpness_factor < 0: raise ValueError(f"sharpness_factor ({sharpness_factor}) is not non-negative.") _assert_image_tensor(img) _assert_channels(img, [1, 3]) if img.size(-1) <= 2 or img.size(-2) <= 2: return img return _blend(img, _blurred_degenerate_image(img), sharpness_factor) def autocontrast(img: Tensor) -> Tensor: _assert_image_tensor(img) if img.ndim < 3: raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {img.ndim}") _assert_channels(img, [1, 3]) bound = 1.0 if img.is_floating_point() else 255.0 dtype = img.dtype if torch.is_floating_point(img) else torch.float32 minimum = img.amin(dim=(-2, -1), keepdim=True).to(dtype) maximum = img.amax(dim=(-2, -1), keepdim=True).to(dtype) scale = bound / (maximum - minimum) eq_idxs = torch.isfinite(scale).logical_not() minimum[eq_idxs] = 0 scale[eq_idxs] = 1 return ((img - minimum) * scale).clamp(0, bound).to(img.dtype) def _scale_channel(img_chan: Tensor) -> Tensor: # TODO: we should expect bincount to always be faster than histc, but this # isn't always the case. Once # https://github.com/pytorch/pytorch/issues/53194 is fixed, remove the if # block and only use bincount. # if img_chan.is_cuda: # hist = torch.histc(img_chan.to(torch.float32), bins=256, min=0, max=255) # else: # hist = torch.bincount(img_chan.view(-1), minlength=256) hist = torch.bincount(img_chan.view(-1), minlength=256) # TODO: nonzero_hist = hist[hist != 0] mask = hist != 0 nonzero_hist = ms.ops.masked_select(hist, mask) step = torch.div(nonzero_hist[:-1].sum(), 255, rounding_mode="floor") if step == 0: return img_chan lut = torch.div(torch.cumsum(hist, 0) + torch.div(step, 2, rounding_mode="floor"), step, rounding_mode="floor") lut = torch.nn.functional.pad(lut, [1, 0])[:-1].clamp(0, 255) return lut[img_chan.to(torch.int64)].to(torch.uint8) def _equalize_single_image(img: Tensor) -> Tensor: return torch.stack([_scale_channel(img[c]) for c in range(img.size(0))]) def equalize(img: Tensor) -> Tensor: _assert_image_tensor(img) if not (3 <= img.ndim <= 4): raise TypeError(f"Input image tensor should have 3 or 4 dimensions, but found {img.ndim}") if img.dtype != torch.uint8: raise TypeError(f"Only torch.uint8 image tensors are supported, but found {img.dtype}") _assert_channels(img, [1, 3]) if img.ndim == 3: return _equalize_single_image(img) return torch.stack([_equalize_single_image(x) for x in img]) def normalize(tensor: Tensor, mean: List[float], std: List[float], inplace: bool = False) -> Tensor: _assert_image_tensor(tensor) if not tensor.is_floating_point(): raise TypeError(f"Input tensor should be a float tensor. Got {tensor.dtype}.") if tensor.ndim < 3: raise ValueError( f"Expected tensor to be a tensor image of size (..., C, H, W). Got tensor.size() = {tensor.size()}" ) if not inplace: tensor = tensor.clone() dtype = tensor.dtype mean = torch.as_tensor(mean, dtype=dtype, device=None) std = torch.as_tensor(std, dtype=dtype, device=None) if (std == 0).any(): raise ValueError(f"std evaluated to zero after conversion to {dtype}, leading to division by zero.") if mean.ndim == 1: mean = mean.view(-1, 1, 1) if std.ndim == 1: std = std.view(-1, 1, 1) tensor.sub_(mean).div_(std) return tensor def erase(img: Tensor, i: int, j: int, h: int, w: int, v: Tensor, inplace: bool = False) -> Tensor: _assert_image_tensor(img) if not inplace: img = img.clone() # TODO: # img[..., i : i + h, j : j + w] = v # return img img_dtype = img.dtype img = img.to(torch.float32) img[..., i : i + h, j : j + w] = v.to(torch.float32) return img.to(img_dtype) msadapter/ms_adapter/torchvision/transforms/transforms.py000066400000000000000000002377151436713554100245220ustar00rootroot00000000000000from __future__ import absolute_import import math import numbers import random import warnings from collections.abc import Sequence from typing import Tuple, List, Optional import ms_adapter.pytorch as torch from mindspore.ops import constexpr from ms_adapter.pytorch import Tensor try: import accimage except ImportError: accimage = None # from ..utils import _log_api_usage_once from . import functional as F from .functional import InterpolationMode, _interpolation_modes_from_int __all__ = [ "Compose", "ToTensor", "PILToTensor", "ConvertImageDtype", "ToPILImage", "Normalize", "Resize", "CenterCrop", "Pad", "Lambda", "RandomApply", "RandomChoice", "RandomOrder", "RandomCrop", "RandomHorizontalFlip", "RandomVerticalFlip", "RandomResizedCrop", "FiveCrop", "TenCrop", "LinearTransformation", "ColorJitter", "RandomRotation", "RandomAffine", "Grayscale", "RandomGrayscale", "RandomPerspective", "RandomErasing", "GaussianBlur", "InterpolationMode", "RandomInvert", "RandomPosterize", "RandomSolarize", "RandomAdjustSharpness", "RandomAutocontrast", "RandomEqualize", ] class Compose: """Composes several transforms together. This transform does not support torchscript. Please, see the note below. Args: transforms (list of ``Transform`` objects): list of transforms to compose. Example: >>> transforms.Compose([ >>> transforms.CenterCrop(10), >>> transforms.PILToTensor(), >>> transforms.ConvertImageDtype(torch.float), >>> ]) .. note:: In order to script the transformations, please use ``torch.nn.Sequential`` as below. >>> transforms = torch.nn.Sequential( >>> transforms.CenterCrop(10), >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), >>> ) >>> scripted_transforms = torch.jit.script(transforms) Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require `lambda` functions or ``PIL.Image``. """ def __init__(self, transforms): # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(self) self.transforms = transforms def __call__(self, img): for t in self.transforms: img = t(img) return img def __repr__(self) -> str: format_string = self.__class__.__name__ + "(" for t in self.transforms: format_string += "\n" format_string += f" {t}" format_string += "\n)" return format_string class ToTensor: """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. This transform does not support torchscript. Converts a PIL Image or numpy.ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) or if the numpy.ndarray has dtype = np.uint8 In the other cases, tensors are returned without scaling. .. note:: Because the input image is scaled to [0.0, 1.0], this transformation should not be used when transforming target image masks. See the `references`_ for implementing the transforms for image masks. .. _references: https://github.com/pytorch/vision/tree/main/references/segmentation """ # def __init__(self) -> None: # _log_api_usage_once(self) def __call__(self, pic): """ Args: pic (PIL Image or numpy.ndarray): Image to be converted to tensor. Returns: Tensor: Converted image. """ return F.to_tensor(pic) def __repr__(self) -> str: return f"{self.__class__.__name__}()" class PILToTensor: """Convert a ``PIL Image`` to a tensor of the same type. This transform does not support torchscript. Converts a PIL Image (H x W x C) to a Tensor of shape (C x H x W). """ # def __init__(self) -> None: # _log_api_usage_once(self) def __call__(self, pic): """ .. note:: A deep copy of the underlying array is performed. Args: pic (PIL Image): Image to be converted to tensor. Returns: Tensor: Converted image. """ return F.pil_to_tensor(pic) def __repr__(self) -> str: return f"{self.__class__.__name__}()" class ConvertImageDtype(torch.nn.Module): """Convert a tensor image to the given ``dtype`` and scale the values accordingly This function does not support PIL Image. Args: dtype (torch.dtype): Desired data type of the output .. note:: When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly. If converted back and forth, this mismatch has no effect. Raises: RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range of the integer ``dtype``. """ def __init__(self, dtype) -> None: super().__init__() # _log_api_usage_once(self) self.dtype = dtype def forward(self, image): return F.convert_image_dtype(image, self.dtype) class ToPILImage: """Convert a tensor or an ndarray to PIL Image. This transform does not support torchscript. Converts a torch.*Tensor of shape C x H x W or a numpy ndarray of shape H x W x C to a PIL Image while preserving the value range. Args: mode (`PIL.Image mode`_): color space and pixel depth of input data (optional). If ``mode`` is ``None`` (default) there are some assumptions made about the input data: - If the input has 4 channels, the ``mode`` is assumed to be ``RGBA``. - If the input has 3 channels, the ``mode`` is assumed to be ``RGB``. - If the input has 2 channels, the ``mode`` is assumed to be ``LA``. - If the input has 1 channel, the ``mode`` is determined by the data type (i.e ``int``, ``float``, ``short``). .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes """ def __init__(self, mode=None): # _log_api_usage_once(self) self.mode = mode def __call__(self, pic): """ Args: pic (Tensor or numpy.ndarray): Image to be converted to PIL Image. Returns: PIL Image: Image converted to PIL Image. """ return F.to_pil_image(pic, self.mode) def __repr__(self) -> str: format_string = self.__class__.__name__ + "(" if self.mode is not None: format_string += f"mode={self.mode}" format_string += ")" return format_string class Normalize(torch.nn.Module): """Normalize a tensor image with mean and standard deviation. This transform does not support PIL Image. Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n`` channels, this transform will normalize each channel of the input ``torch.*Tensor`` i.e., ``output[channel] = (input[channel] - mean[channel]) / std[channel]`` .. note:: This transform acts out of place, i.e., it does not mutate the input tensor. Args: mean (sequence): Sequence of means for each channel. std (sequence): Sequence of standard deviations for each channel. inplace(bool,optional): Bool to make this operation in-place. """ def __init__(self, mean, std, inplace=False): super().__init__() # _log_api_usage_once(self) self.mean = mean self.std = std self.inplace = inplace def forward(self, tensor: Tensor) -> Tensor: """ Args: tensor (Tensor): Tensor image to be normalized. Returns: Tensor: Normalized Tensor image. """ return F.normalize(tensor, self.mean, self.std, self.inplace) def __repr__(self) -> str: return f"{self.__class__.__name__}(mean={self.mean}, std={self.std})" class Resize(torch.nn.Module): """Resize the input image to the given size. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions .. warning:: The output image might be different depending on its type: when downsampling, the interpolation of PIL images and tensors is slightly different, because PIL applies antialiasing. This may lead to significant differences in the performance of a network. Therefore, it is preferable to train and serve a model with the same input types. See also below the ``antialias`` parameter, which can help making the output of PIL images and tensors closer. Args: size (sequence or int): Desired output size. If size is a sequence like (h, w), output size will be matched to this. If size is an int, smaller edge of the image will be matched to this number. i.e, if height > width, then image will be rescaled to (size * height / width, size). .. note:: In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. max_size (int, optional): The maximum allowed for the longer edge of the resized image: if the longer edge of the image is greater than ``max_size`` after being resized according to ``size``, then the image is resized again so that the longer edge is equal to ``max_size``. As a result, ``size`` might be overruled, i.e the smaller edge may be shorter than ``size``. This is only supported if ``size`` is an int (or a sequence of length 1 in torchscript mode). antialias (bool, optional): antialias flag. If ``img`` is PIL Image, the flag is ignored and anti-alias is always used. If ``img`` is Tensor, the flag is False by default and can be set to True for ``InterpolationMode.BILINEAR`` only mode. This can help making the output for PIL images and tensors closer. .. warning:: There is no autodiff support for ``antialias=True`` option with input ``img`` as Tensor. """ def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None, antialias=None): super().__init__() # _log_api_usage_once(self) if not isinstance(size, (int, Sequence)): raise TypeError(f"Size should be int or sequence. Got {type(size)}") if isinstance(size, Sequence) and len(size) not in (1, 2): raise ValueError("If size is a sequence, it should have 1 or 2 values") self.size = size self.max_size = max_size # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) self.interpolation = interpolation self.antialias = antialias def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be scaled. Returns: PIL Image or Tensor: Rescaled image. """ return F.resize(img, self.size, self.interpolation, self.max_size, self.antialias) def __repr__(self) -> str: detail = f"(size={self.size}, interpolation={self.interpolation.value}, max_size={self.max_size}, antialias={self.antialias})" return f"{self.__class__.__name__}{detail}" class CenterCrop(torch.nn.Module): """Crops the given image at the center. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. If image size is smaller than output size along any edge, image is padded with 0 and then center cropped. Args: size (sequence or int): Desired output size of the crop. If size is an int instead of sequence like (h, w), a square crop (size, size) is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). """ def __init__(self, size): super().__init__() # _log_api_usage_once(self) self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be cropped. Returns: PIL Image or Tensor: Cropped image. """ return F.center_crop(img, self.size) def __repr__(self) -> str: return f"{self.__class__.__name__}(size={self.size})" class Pad(torch.nn.Module): """Pad the given image on all sides with the given "pad" value. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means at most 2 leading dimensions for mode reflect and symmetric, at most 3 leading dimensions for mode edge, and an arbitrary number of leading dimensions for mode constant Args: padding (int or sequence): Padding on each border. If a single int is provided this is used to pad all borders. If sequence of length 2 is provided this is the padding on left/right and top/bottom respectively. If a sequence of length 4 is provided this is the padding for the left, top, right and bottom borders respectively. .. note:: In torchscript mode padding as single int is not supported, use a sequence of length 1: ``[padding, ]``. fill (number or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. This value is only used when the padding_mode is constant. Only number is supported for torch Tensor. Only int or tuple value is supported for PIL Image. padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant. - constant: pads with a constant value, this value is specified with fill - edge: pads with the last value at the edge of the image. If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2 - reflect: pads with reflection of image without repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode will result in [3, 2, 1, 2, 3, 4, 3, 2] - symmetric: pads with reflection of image repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode will result in [2, 1, 1, 2, 3, 4, 4, 3] """ def __init__(self, padding, fill=0, padding_mode="constant"): super().__init__() # _log_api_usage_once(self) if not isinstance(padding, (numbers.Number, tuple, list)): raise TypeError("Got inappropriate padding arg") if not isinstance(fill, (numbers.Number, tuple, list)): raise TypeError("Got inappropriate fill arg") if padding_mode not in ["constant", "edge", "reflect", "symmetric"]: raise ValueError("Padding mode should be either constant, edge, reflect or symmetric") if isinstance(padding, Sequence) and len(padding) not in [1, 2, 4]: raise ValueError( f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple" ) self.padding = padding self.fill = fill self.padding_mode = padding_mode def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be padded. Returns: PIL Image or Tensor: Padded image. """ return F.pad(img, self.padding, self.fill, self.padding_mode) def __repr__(self) -> str: return f"{self.__class__.__name__}(padding={self.padding}, fill={self.fill}, padding_mode={self.padding_mode})" class Lambda: """Apply a user-defined lambda as a transform. This transform does not support torchscript. Args: lambd (function): Lambda/function to be used for transform. """ def __init__(self, lambd): # _log_api_usage_once(self) if not callable(lambd): raise TypeError(f"Argument lambd should be callable, got {repr(type(lambd).__name__)}") self.lambd = lambd def __call__(self, img): return self.lambd(img) def __repr__(self) -> str: return f"{self.__class__.__name__}()" class RandomTransforms: """Base class for a list of transformations with randomness Args: transforms (sequence): list of transformations """ def __init__(self, transforms): # _log_api_usage_once(self) if not isinstance(transforms, Sequence): raise TypeError("Argument transforms should be a sequence") self.transforms = transforms def __call__(self, *args, **kwargs): raise NotImplementedError() def __repr__(self) -> str: format_string = self.__class__.__name__ + "(" for t in self.transforms: format_string += "\n" format_string += f" {t}" format_string += "\n)" return format_string class RandomApply(torch.nn.Module): """Apply randomly a list of transformations with a given probability. .. note:: In order to script the transformation, please use ``torch.nn.ModuleList`` as input instead of list/tuple of transforms as shown below: >>> transforms = transforms.RandomApply(torch.nn.ModuleList([ >>> transforms.ColorJitter(), >>> ]), p=0.3) >>> scripted_transforms = torch.jit.script(transforms) Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require `lambda` functions or ``PIL.Image``. Args: transforms (sequence or torch.nn.Module): list of transformations p (float): probability """ def __init__(self, transforms, p=0.5): super().__init__() # _log_api_usage_once(self) self.transforms = transforms self.p = p def forward(self, img): if self.p < torch.rand(1): return img for t in self.transforms: img = t(img) return img def __repr__(self) -> str: format_string = self.__class__.__name__ + "(" format_string += f"\n p={self.p}" for t in self.transforms: format_string += "\n" format_string += f" {t}" format_string += "\n)" return format_string class RandomOrder(RandomTransforms): """Apply a list of transformations in a random order. This transform does not support torchscript.""" def __call__(self, img): order = list(range(len(self.transforms))) random.shuffle(order) for i in order: img = self.transforms[i](img) return img class RandomChoice(RandomTransforms): """Apply single transformation randomly picked from a list. This transform does not support torchscript.""" def __init__(self, transforms, p=None): super().__init__(transforms) if p is not None and not isinstance(p, Sequence): raise TypeError("Argument p should be a sequence") self.p = p def __call__(self, *args): t = random.choices(self.transforms, weights=self.p)[0] return t(*args) def __repr__(self) -> str: return f"{super().__repr__()}(p={self.p})" class RandomCrop(torch.nn.Module): """Crop the given image at a random location. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions, but if non-constant padding is used, the input is expected to have at most 2 leading dimensions Args: size (sequence or int): Desired output size of the crop. If size is an int instead of sequence like (h, w), a square crop (size, size) is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). padding (int or sequence, optional): Optional padding on each border of the image. Default is None. If a single int is provided this is used to pad all borders. If sequence of length 2 is provided this is the padding on left/right and top/bottom respectively. If a sequence of length 4 is provided this is the padding for the left, top, right and bottom borders respectively. .. note:: In torchscript mode padding as single int is not supported, use a sequence of length 1: ``[padding, ]``. pad_if_needed (boolean): It will pad the image if smaller than the desired size to avoid raising an exception. Since cropping is done after padding, the padding seems to be done at a random offset. fill (number or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. This value is only used when the padding_mode is constant. Only number is supported for torch Tensor. Only int or tuple value is supported for PIL Image. padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant. - constant: pads with a constant value, this value is specified with fill - edge: pads with the last value at the edge of the image. If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2 - reflect: pads with reflection of image without repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode will result in [3, 2, 1, 2, 3, 4, 3, 2] - symmetric: pads with reflection of image repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode will result in [2, 1, 1, 2, 3, 4, 4, 3] """ @staticmethod def get_params(img: Tensor, output_size: Tuple[int, int]) -> Tuple[int, int, int, int]: """Get parameters for ``crop`` for a random crop. Args: img (PIL Image or Tensor): Image to be cropped. output_size (tuple): Expected output size of the crop. Returns: tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. """ _, h, w = F.get_dimensions(img) th, tw = output_size if h + 1 < th or w + 1 < tw: raise ValueError(f"Required crop size {(th, tw)} is larger then input image size {(h, w)}") if w == tw and h == th: return 0, 0, h, w i = torch.randint(0, h - th + 1, size=(1,)).item() j = torch.randint(0, w - tw + 1, size=(1,)).item() return i, j, th, tw def __init__(self, size, padding=None, pad_if_needed=False, fill=0, padding_mode="constant"): super().__init__() # _log_api_usage_once(self) self.size = tuple(_setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")) self.padding = padding self.pad_if_needed = pad_if_needed self.fill = fill self.padding_mode = padding_mode def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be cropped. Returns: PIL Image or Tensor: Cropped image. """ if self.padding is not None: img = F.pad(img, self.padding, self.fill, self.padding_mode) _, height, width = F.get_dimensions(img) # pad the width if needed if self.pad_if_needed and width < self.size[1]: padding = [self.size[1] - width, 0] img = F.pad(img, padding, self.fill, self.padding_mode) # pad the height if needed if self.pad_if_needed and height < self.size[0]: padding = [0, self.size[0] - height] img = F.pad(img, padding, self.fill, self.padding_mode) i, j, h, w = self.get_params(img, self.size) return F.crop(img, i, j, h, w) def __repr__(self) -> str: return f"{self.__class__.__name__}(size={self.size}, padding={self.padding})" class RandomHorizontalFlip(torch.nn.Module): """Horizontally flip the given image randomly with a given probability. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions Args: p (float): probability of the image being flipped. Default value is 0.5 """ def __init__(self, p=0.5): super().__init__() # _log_api_usage_once(self) self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be flipped. Returns: PIL Image or Tensor: Randomly flipped image. """ if torch.rand(1) < self.p: return F.hflip(img) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" class RandomVerticalFlip(torch.nn.Module): """Vertically flip the given image randomly with a given probability. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions Args: p (float): probability of the image being flipped. Default value is 0.5 """ def __init__(self, p=0.5): super().__init__() # _log_api_usage_once(self) self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be flipped. Returns: PIL Image or Tensor: Randomly flipped image. """ if torch.rand(1) < self.p: return F.vflip(img) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" class RandomPerspective(torch.nn.Module): """Performs a random perspective transformation of the given image with a given probability. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. Args: distortion_scale (float): argument to control the degree of distortion and ranges from 0 to 1. Default is 0.5. p (float): probability of the image being transformed. Default is 0.5. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. fill (sequence or number): Pixel fill value for the area outside the transformed image. Default is ``0``. If given a number, the value is used for all bands respectively. """ def __init__(self, distortion_scale=0.5, p=0.5, interpolation=InterpolationMode.BILINEAR, fill=0): super().__init__() # _log_api_usage_once(self) self.p = p # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) self.interpolation = interpolation self.distortion_scale = distortion_scale if fill is None: fill = 0 elif not isinstance(fill, (Sequence, numbers.Number)): raise TypeError("Fill should be either a sequence or a number.") self.fill = fill def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be Perspectively transformed. Returns: PIL Image or Tensor: Randomly transformed image. """ fill = self.fill channels, height, width = F.get_dimensions(img) if isinstance(img, Tensor): if isinstance(fill, (int, float)): fill = [float(fill)] * channels else: fill = [float(f) for f in fill] if torch.rand(1) < self.p: startpoints, endpoints = self.get_params(width, height, self.distortion_scale) return F.perspective(img, startpoints, endpoints, self.interpolation, fill) return img @staticmethod def get_params(width: int, height: int, distortion_scale: float) -> Tuple[List[List[int]], List[List[int]]]: """Get parameters for ``perspective`` for a random perspective transform. Args: width (int): width of the image. height (int): height of the image. distortion_scale (float): argument to control the degree of distortion and ranges from 0 to 1. Returns: List containing [top-left, top-right, bottom-right, bottom-left] of the original image, List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image. """ half_height = height // 2 half_width = width // 2 topleft = [ int(torch.randint(0, int(distortion_scale * half_width) + 1, size=(1,)).item()), int(torch.randint(0, int(distortion_scale * half_height) + 1, size=(1,)).item()), ] topright = [ int(torch.randint(width - int(distortion_scale * half_width) - 1, width, size=(1,)).item()), int(torch.randint(0, int(distortion_scale * half_height) + 1, size=(1,)).item()), ] botright = [ int(torch.randint(width - int(distortion_scale * half_width) - 1, width, size=(1,)).item()), int(torch.randint(height - int(distortion_scale * half_height) - 1, height, size=(1,)).item()), ] botleft = [ int(torch.randint(0, int(distortion_scale * half_width) + 1, size=(1,)).item()), int(torch.randint(height - int(distortion_scale * half_height) - 1, height, size=(1,)).item()), ] startpoints = [[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]] endpoints = [topleft, topright, botright, botleft] return startpoints, endpoints def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" class RandomResizedCrop(torch.nn.Module): """Crop a random portion of image and resize it to a given size. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions A crop of the original image is made: the crop has a random area (H * W) and a random aspect ratio. This crop is finally resized to the given size. This is popularly used to train the Inception networks. Args: size (int or sequence): expected output size of the crop, for each edge. If size is an int instead of sequence like (h, w), a square output size ``(size, size)`` is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). .. note:: In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop, before resizing. The scale is defined with respect to the area of the original image. ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before resizing. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. """ def __init__(self, size, scale=(0.08, 1.0), ratio=(3.0 / 4.0, 4.0 / 3.0), interpolation=InterpolationMode.BILINEAR): super().__init__() # _log_api_usage_once(self) self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") if not isinstance(scale, Sequence): raise TypeError("Scale should be a sequence") if not isinstance(ratio, Sequence): raise TypeError("Ratio should be a sequence") if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): warnings.warn("Scale and ratio should be of kind (min, max)") # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) self.interpolation = interpolation self.scale = scale self.ratio = ratio @staticmethod def get_params(img: Tensor, scale: List[float], ratio: List[float]) -> Tuple[int, int, int, int]: """Get parameters for ``crop`` for a random sized crop. Args: img (PIL Image or Tensor): Input image. scale (list): range of scale of the origin size cropped ratio (list): range of aspect ratio of the origin aspect ratio cropped Returns: tuple: params (i, j, h, w) to be passed to ``crop`` for a random sized crop. """ _, height, width = F.get_dimensions(img) area = height * width log_ratio = torch.log(torch.tensor(ratio)) for _ in range(10): target_area = area * torch.empty(1).uniform_adpater(scale[0], scale[1]).item() aspect_ratio = torch.exp(torch.empty(1).uniform_adpater(log_ratio[0], log_ratio[1])).item() w = int(round(math.sqrt(target_area * aspect_ratio))) h = int(round(math.sqrt(target_area / aspect_ratio))) if 0 < w <= width and 0 < h <= height: i = torch.randint(0, height - h + 1, size=(1,)).item() j = torch.randint(0, width - w + 1, size=(1,)).item() return i, j, h, w # Fallback to central crop in_ratio = float(width) / float(height) if in_ratio < min(ratio): w = width h = int(round(w / min(ratio))) elif in_ratio > max(ratio): h = height w = int(round(h * max(ratio))) else: # whole image w = width h = height i = (height - h) // 2 j = (width - w) // 2 return i, j, h, w def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be cropped and resized. Returns: PIL Image or Tensor: Randomly cropped and resized image. """ i, j, h, w = self.get_params(img, self.scale, self.ratio) return F.resized_crop(img, i, j, h, w, self.size, self.interpolation) def __repr__(self) -> str: interpolate_str = self.interpolation.value format_string = self.__class__.__name__ + f"(size={self.size}" format_string += f", scale={tuple(round(s, 4) for s in self.scale)}" format_string += f", ratio={tuple(round(r, 4) for r in self.ratio)}" format_string += f", interpolation={interpolate_str})" return format_string class FiveCrop(torch.nn.Module): """Crop the given image into four corners and the central crop. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions .. Note:: This transform returns a tuple of images and there may be a mismatch in the number of inputs and targets your Dataset returns. See below for an example of how to deal with this. Args: size (sequence or int): Desired output size of the crop. If size is an ``int`` instead of sequence like (h, w), a square crop of size (size, size) is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). Example: >>> transform = Compose([ >>> FiveCrop(size), # this is a list of PIL Images >>> Lambda(lambda crops: torch.stack([PILToTensor()(crop) for crop in crops])) # returns a 4D tensor >>> ]) >>> #In your test loop you can do the following: >>> input, target = batch # input is a 5d tensor, target is 2d >>> bs, ncrops, c, h, w = input.size() >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops """ def __init__(self, size): super().__init__() # _log_api_usage_once(self) self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be cropped. Returns: tuple of 5 images. Image can be PIL Image or Tensor """ return F.five_crop(img, self.size) def __repr__(self) -> str: return f"{self.__class__.__name__}(size={self.size})" class TenCrop(torch.nn.Module): """Crop the given image into four corners and the central crop plus the flipped version of these (horizontal flipping is used by default). If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions .. Note:: This transform returns a tuple of images and there may be a mismatch in the number of inputs and targets your Dataset returns. See below for an example of how to deal with this. Args: size (sequence or int): Desired output size of the crop. If size is an int instead of sequence like (h, w), a square crop (size, size) is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). vertical_flip (bool): Use vertical flipping instead of horizontal Example: >>> transform = Compose([ >>> TenCrop(size), # this is a list of PIL Images >>> Lambda(lambda crops: torch.stack([PILToTensor()(crop) for crop in crops])) # returns a 4D tensor >>> ]) >>> #In your test loop you can do the following: >>> input, target = batch # input is a 5d tensor, target is 2d >>> bs, ncrops, c, h, w = input.size() >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops """ def __init__(self, size, vertical_flip=False): super().__init__() # _log_api_usage_once(self) self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.") self.vertical_flip = vertical_flip def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be cropped. Returns: tuple of 10 images. Image can be PIL Image or Tensor """ return F.ten_crop(img, self.size, self.vertical_flip) def __repr__(self) -> str: return f"{self.__class__.__name__}(size={self.size}, vertical_flip={self.vertical_flip})" class LinearTransformation(torch.nn.Module): """Transform a tensor image with a square transformation matrix and a mean_vector computed offline. This transform does not support PIL Image. Given transformation_matrix and mean_vector, will flatten the torch.*Tensor and subtract mean_vector from it which is then followed by computing the dot product with the transformation matrix and then reshaping the tensor to its original shape. Applications: whitening transformation: Suppose X is a column vector zero-centered data. Then compute the data covariance matrix [D x D] with torch.mm(X.t(), X), perform SVD on this matrix and pass it as transformation_matrix. Args: transformation_matrix (Tensor): tensor [D x D], D = C x H x W mean_vector (Tensor): tensor [D], D = C x H x W """ def __init__(self, transformation_matrix, mean_vector): super().__init__() # _log_api_usage_once(self) if transformation_matrix.size(0) != transformation_matrix.size(1): raise ValueError( "transformation_matrix should be square. Got " f"{tuple(transformation_matrix.size())} rectangular matrix." ) if mean_vector.size(0) != transformation_matrix.size(0): raise ValueError( f"mean_vector should have the same length {mean_vector.size(0)}" f" as any one of the dimensions of the transformation_matrix [{tuple(transformation_matrix.size())}]" ) # if transformation_matrix.device != mean_vector.device: # raise ValueError( # f"Input tensors should be on the same device. Got {transformation_matrix.device} and {mean_vector.device}" # ) self.transformation_matrix = transformation_matrix self.mean_vector = mean_vector def forward(self, tensor: Tensor) -> Tensor: """ Args: tensor (Tensor): Tensor image to be whitened. Returns: Tensor: Transformed image. """ shape = tensor.shape n = shape[-3] * shape[-2] * shape[-1] if n != self.transformation_matrix.shape[0]: raise ValueError( "Input tensor and transformation matrix have incompatible shape." + f"[{shape[-3]} x {shape[-2]} x {shape[-1]}] != " + f"{self.transformation_matrix.shape[0]}" ) # if tensor.device.type != self.mean_vector.device.type: # raise ValueError( # "Input tensor should be on the same device as transformation matrix and mean vector. " # f"Got {tensor.device} vs {self.mean_vector.device}" # ) flat_tensor = tensor.view(-1, n) - self.mean_vector transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix) tensor = transformed_tensor.view(shape) return tensor def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(transformation_matrix=" f"{self.transformation_matrix.tolist()}" f", mean_vector={self.mean_vector.tolist()})" ) return s class ColorJitter(torch.nn.Module): """Randomly change the brightness, contrast, saturation and hue of an image. If the image is torch Tensor, it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, mode "1", "I", "F" and modes with transparency (alpha channel) are not supported. Args: brightness (float or tuple of float (min, max)): How much to jitter brightness. brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness] or the given [min, max]. Should be non negative numbers. contrast (float or tuple of float (min, max)): How much to jitter contrast. contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast] or the given [min, max]. Should be non negative numbers. saturation (float or tuple of float (min, max)): How much to jitter saturation. saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation] or the given [min, max]. Should be non negative numbers. hue (float or tuple of float (min, max)): How much to jitter hue. hue_factor is chosen uniformly from [-hue, hue] or the given [min, max]. Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5. To jitter hue, the pixel values of the input image has to be non-negative for conversion to HSV space; thus it does not work if you normalize your image to an interval with negative values, or use an interpolation that generates negative values before using this function. """ def __init__(self, brightness=0, contrast=0, saturation=0, hue=0): super().__init__() # _log_api_usage_once(self) self.brightness = self._check_input(brightness, "brightness") self.contrast = self._check_input(contrast, "contrast") self.saturation = self._check_input(saturation, "saturation") self.hue = self._check_input(hue, "hue", center=0, bound=(-0.5, 0.5), clip_first_on_zero=False) def _check_input(self, value, name, center=1, bound=(0, float("inf")), clip_first_on_zero=True): if isinstance(value, numbers.Number): if value < 0: raise ValueError(f"If {name} is a single number, it must be non negative.") value = [center - float(value), center + float(value)] if clip_first_on_zero: value[0] = max(value[0], 0.0) elif isinstance(value, (tuple, list)) and len(value) == 2: if not bound[0] <= value[0] <= value[1] <= bound[1]: raise ValueError(f"{name} values should be between {bound}") else: raise TypeError(f"{name} should be a single number or a list/tuple with length 2.") # if value is 0 or (1., 1.) for brightness/contrast/saturation # or (0., 0.) for hue, do nothing if value[0] == value[1] == center: value = None return value @staticmethod def get_params( brightness: Optional[List[float]], contrast: Optional[List[float]], saturation: Optional[List[float]], hue: Optional[List[float]], ) -> Tuple[Tensor, Optional[float], Optional[float], Optional[float], Optional[float]]: """Get the parameters for the randomized transform to be applied on image. Args: brightness (tuple of float (min, max), optional): The range from which the brightness_factor is chosen uniformly. Pass None to turn off the transformation. contrast (tuple of float (min, max), optional): The range from which the contrast_factor is chosen uniformly. Pass None to turn off the transformation. saturation (tuple of float (min, max), optional): The range from which the saturation_factor is chosen uniformly. Pass None to turn off the transformation. hue (tuple of float (min, max), optional): The range from which the hue_factor is chosen uniformly. Pass None to turn off the transformation. Returns: tuple: The parameters used to apply the randomized transform along with their random order. """ fn_idx = torch.randperm(4) b = None if brightness is None else float(torch.empty(1).uniform_adpater(brightness[0], brightness[1])) c = None if contrast is None else float(torch.empty(1).uniform_adpater(contrast[0], contrast[1])) s = None if saturation is None else float(torch.empty(1).uniform_adpater(saturation[0], saturation[1])) h = None if hue is None else float(torch.empty(1).uniform_adpater(hue[0], hue[1])) return fn_idx, b, c, s, h def forward(self, img): """ Args: img (PIL Image or Tensor): Input image. Returns: PIL Image or Tensor: Color jittered image. """ fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = self.get_params( self.brightness, self.contrast, self.saturation, self.hue ) for fn_id in fn_idx: if fn_id == 0 and brightness_factor is not None: img = F.adjust_brightness(img, brightness_factor) elif fn_id == 1 and contrast_factor is not None: img = F.adjust_contrast(img, contrast_factor) elif fn_id == 2 and saturation_factor is not None: img = F.adjust_saturation(img, saturation_factor) elif fn_id == 3 and hue_factor is not None: img = F.adjust_hue(img, hue_factor) return img def __repr__(self) -> str: s = ( f"{self.__class__.__name__}(" f"brightness={self.brightness}" f", contrast={self.contrast}" f", saturation={self.saturation}" f", hue={self.hue})" ) return s class RandomRotation(torch.nn.Module): """Rotate the image by angle. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. Args: degrees (sequence or number): Range of degrees to select from. If degrees is a number instead of sequence like (min, max), the range of degrees will be (-degrees, +degrees). interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. expand (bool, optional): Optional expansion flag. If true, expands the output to make it large enough to hold the entire rotated image. If false or omitted, make the output image the same size as the input image. Note that the expand flag assumes rotation around the center and no translation. center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner. Default is the center of the image. fill (sequence or number): Pixel fill value for the area outside the rotated image. Default is ``0``. If given a number, the value is used for all bands respectively. resample (int, optional): .. warning:: This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``interpolation`` instead. .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters """ def __init__( self, degrees, interpolation=InterpolationMode.NEAREST, expand=False, center=None, fill=0, resample=None ): super().__init__() # _log_api_usage_once(self) if resample is not None: warnings.warn( "The parameter 'resample' is deprecated since 0.12 and will be removed 0.14. " "Please use 'interpolation' instead." ) interpolation = _interpolation_modes_from_int(resample) # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,)) if center is not None: _check_sequence_input(center, "center", req_sizes=(2,)) self.center = center self.resample = self.interpolation = interpolation self.expand = expand if fill is None: fill = 0 elif not isinstance(fill, (Sequence, numbers.Number)): raise TypeError("Fill should be either a sequence or a number.") self.fill = fill @staticmethod def get_params(degrees: List[float]) -> float: """Get parameters for ``rotate`` for a random rotation. Returns: float: angle parameter to be passed to ``rotate`` for random rotation. """ angle = float(torch.empty(1).uniform_adpater(float(degrees[0]), float(degrees[1])).item()) return angle def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be rotated. Returns: PIL Image or Tensor: Rotated image. """ fill = self.fill channels, _, _ = F.get_dimensions(img) if isinstance(img, Tensor): if isinstance(fill, (int, float)): fill = [float(fill)] * channels else: fill = [float(f) for f in fill] angle = self.get_params(self.degrees) return F.rotate(img, angle, self.resample, self.expand, self.center, fill) def __repr__(self) -> str: interpolate_str = self.interpolation.value format_string = self.__class__.__name__ + f"(degrees={self.degrees}" format_string += f", interpolation={interpolate_str}" format_string += f", expand={self.expand}" if self.center is not None: format_string += f", center={self.center}" if self.fill is not None: format_string += f", fill={self.fill}" format_string += ")" return format_string class RandomAffine(torch.nn.Module): """Random affine transformation of the image keeping center invariant. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions. Args: degrees (sequence or number): Range of degrees to select from. If degrees is a number instead of sequence like (min, max), the range of degrees will be (-degrees, +degrees). Set to 0 to deactivate rotations. translate (tuple, optional): tuple of maximum absolute fraction for horizontal and vertical translations. For example translate=(a, b), then horizontal shift is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default. scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is randomly sampled from the range a <= scale <= b. Will keep original scale by default. shear (sequence or number, optional): Range of degrees to select from. If shear is a number, a shear parallel to the x axis in the range (-shear, +shear) will be applied. Else if shear is a sequence of 2 values a shear parallel to the x axis in the range (shear[0], shear[1]) will be applied. Else if shear is a sequence of 4 values, a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied. Will not apply shear by default. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. fill (sequence or number): Pixel fill value for the area outside the transformed image. Default is ``0``. If given a number, the value is used for all bands respectively. fillcolor (sequence or number, optional): .. warning:: This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``fill`` instead. resample (int, optional): .. warning:: This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``interpolation`` instead. center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner. Default is the center of the image. .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters """ def __init__( self, degrees, translate=None, scale=None, shear=None, interpolation=InterpolationMode.NEAREST, fill=0, fillcolor=None, resample=None, center=None, ): super().__init__() # _log_api_usage_once(self) if resample is not None: warnings.warn( "The parameter 'resample' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'interpolation' instead." ) interpolation = _interpolation_modes_from_int(resample) # Backward compatibility with integer value if isinstance(interpolation, int): warnings.warn( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ) interpolation = _interpolation_modes_from_int(interpolation) if fillcolor is not None: warnings.warn( "The parameter 'fillcolor' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'fill' instead." ) fill = fillcolor self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,)) if translate is not None: _check_sequence_input(translate, "translate", req_sizes=(2,)) for t in translate: if not (0.0 <= t <= 1.0): raise ValueError("translation values should be between 0 and 1") self.translate = translate if scale is not None: _check_sequence_input(scale, "scale", req_sizes=(2,)) for s in scale: if s <= 0: raise ValueError("scale values should be positive") self.scale = scale if shear is not None: self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4)) else: self.shear = shear self.resample = self.interpolation = interpolation if fill is None: fill = 0 elif not isinstance(fill, (Sequence, numbers.Number)): raise TypeError("Fill should be either a sequence or a number.") self.fillcolor = self.fill = fill if center is not None: _check_sequence_input(center, "center", req_sizes=(2,)) self.center = center @staticmethod def get_params( degrees: List[float], translate: Optional[List[float]], scale_ranges: Optional[List[float]], shears: Optional[List[float]], img_size: List[int], ) -> Tuple[float, Tuple[int, int], float, Tuple[float, float]]: """Get parameters for affine transformation Returns: params to be passed to the affine transformation """ angle = float(torch.empty(1).uniform_adpater(float(degrees[0]), float(degrees[1])).item()) if translate is not None: max_dx = float(translate[0] * img_size[0]) max_dy = float(translate[1] * img_size[1]) tx = int(round(torch.empty(1).uniform_adpater(-max_dx, max_dx).item())) ty = int(round(torch.empty(1).uniform_adpater(-max_dy, max_dy).item())) translations = (tx, ty) else: translations = (0, 0) if scale_ranges is not None: scale = float(torch.empty(1).uniform_adpater(scale_ranges[0], scale_ranges[1]).item()) else: scale = 1.0 shear_x = shear_y = 0.0 if shears is not None: shear_x = float(torch.empty(1).uniform_adpater(shears[0], shears[1]).item()) if len(shears) == 4: shear_y = float(torch.empty(1).uniform_adpater(shears[2], shears[3]).item()) shear = (shear_x, shear_y) return angle, translations, scale, shear def forward(self, img): """ img (PIL Image or Tensor): Image to be transformed. Returns: PIL Image or Tensor: Affine transformed image. """ fill = self.fill channels, height, width = F.get_dimensions(img) if isinstance(img, Tensor): if isinstance(fill, (int, float)): fill = [float(fill)] * channels else: fill = [float(f) for f in fill] img_size = [width, height] # flip for keeping BC on get_params call ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img_size) return F.affine(img, *ret, interpolation=self.interpolation, fill=fill, center=self.center) def __repr__(self) -> str: s = f"{self.__class__.__name__}(degrees={self.degrees}" s += f", translate={self.translate}" if self.translate is not None else "" s += f", scale={self.scale}" if self.scale is not None else "" s += f", shear={self.shear}" if self.shear is not None else "" s += f", interpolation={self.interpolation.value}" if self.interpolation != InterpolationMode.NEAREST else "" s += f", fill={self.fill}" if self.fill != 0 else "" s += f", center={self.center}" if self.center is not None else "" s += ")" return s class Grayscale(torch.nn.Module): """Convert image to grayscale. If the image is torch Tensor, it is expected to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions Args: num_output_channels (int): (1 or 3) number of channels desired for output image Returns: PIL Image: Grayscale version of the input. - If ``num_output_channels == 1`` : returned image is single channel - If ``num_output_channels == 3`` : returned image is 3 channel with r == g == b """ def __init__(self, num_output_channels=1): super().__init__() # _log_api_usage_once(self) self.num_output_channels = num_output_channels def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be converted to grayscale. Returns: PIL Image or Tensor: Grayscaled image. """ return F.rgb_to_grayscale(img, num_output_channels=self.num_output_channels) def __repr__(self) -> str: return f"{self.__class__.__name__}(num_output_channels={self.num_output_channels})" class RandomGrayscale(torch.nn.Module): """Randomly convert image to grayscale with a probability of p (default 0.1). If the image is torch Tensor, it is expected to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions Args: p (float): probability that image should be converted to grayscale. Returns: PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged with probability (1-p). - If input image is 1 channel: grayscale version is 1 channel - If input image is 3 channel: grayscale version is 3 channel with r == g == b """ def __init__(self, p=0.1): super().__init__() # _log_api_usage_once(self) self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be converted to grayscale. Returns: PIL Image or Tensor: Randomly grayscaled image. """ num_output_channels, _, _ = F.get_dimensions(img) if torch.rand(1) < self.p: return F.rgb_to_grayscale(img, num_output_channels=num_output_channels) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" class RandomErasing(torch.nn.Module): """Randomly selects a rectangle region in an torch Tensor image and erases its pixels. This transform does not support PIL Image. 'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896 Args: p: probability that the random erasing operation will be performed. scale: range of proportion of erased area against input image. ratio: range of aspect ratio of erased area. value: erasing value. Default is 0. If a single int, it is used to erase all pixels. If a tuple of length 3, it is used to erase R, G, B channels respectively. If a str of 'random', erasing each pixel with random values. inplace: boolean to make this transform inplace. Default set to False. Returns: Erased Image. Example: >>> transform = transforms.Compose([ >>> transforms.RandomHorizontalFlip(), >>> transforms.PILToTensor(), >>> transforms.ConvertImageDtype(torch.float), >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), >>> transforms.RandomErasing(), >>> ]) """ def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False): super().__init__() # _log_api_usage_once(self) if not isinstance(value, (numbers.Number, str, tuple, list)): raise TypeError("Argument value should be either a number or str or a sequence") if isinstance(value, str) and value != "random": raise ValueError("If value is str, it should be 'random'") if not isinstance(scale, (tuple, list)): raise TypeError("Scale should be a sequence") if not isinstance(ratio, (tuple, list)): raise TypeError("Ratio should be a sequence") if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): warnings.warn("Scale and ratio should be of kind (min, max)") if scale[0] < 0 or scale[1] > 1: raise ValueError("Scale should be between 0 and 1") if p < 0 or p > 1: raise ValueError("Random erasing probability should be between 0 and 1") self.p = p self.scale = scale self.ratio = ratio self.value = value self.inplace = inplace @staticmethod def get_params( img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None ) -> Tuple[int, int, int, int, Tensor]: """Get parameters for ``erase`` for a random erasing. Args: img (Tensor): Tensor image to be erased. scale (sequence): range of proportion of erased area against input image. ratio (sequence): range of aspect ratio of erased area. value (list, optional): erasing value. If None, it is interpreted as "random" (erasing each pixel with random values). If ``len(value)`` is 1, it is interpreted as a number, i.e. ``value[0]``. Returns: tuple: params (i, j, h, w, v) to be passed to ``erase`` for random erasing. """ img_c, img_h, img_w = img.shape[-3], img.shape[-2], img.shape[-1] area = img_h * img_w log_ratio = torch.log(torch.tensor(ratio)) for _ in range(10): erase_area = area * torch.empty(1).uniform_adpater(scale[0], scale[1]).item() aspect_ratio = torch.exp(torch.empty(1).uniform_adpater(log_ratio[0], log_ratio[1])).item() h = int(round(math.sqrt(erase_area * aspect_ratio))) w = int(round(math.sqrt(erase_area / aspect_ratio))) if not (h < img_h and w < img_w): continue if value is None: v = torch.empty([img_c, h, w], dtype=torch.float32).normal_() else: v = torch.tensor(value)[:, None, None] i = torch.randint(0, img_h - h + 1, size=(1,)).item() j = torch.randint(0, img_w - w + 1, size=(1,)).item() return i, j, h, w, v # Return original image return 0, 0, img_h, img_w, img def forward(self, img): """ Args: img (Tensor): Tensor image to be erased. Returns: img (Tensor): Erased Tensor image. """ if torch.rand(1) < self.p: # cast self.value to script acceptable type if isinstance(self.value, (int, float)): value = [self.value] elif isinstance(self.value, str): value = None elif isinstance(self.value, tuple): value = list(self.value) else: value = self.value if value is not None and not (len(value) in (1, img.shape[-3])): raise ValueError( "If value is a sequence, it should have either a single value or " f"{img.shape[-3]} (number of input channels)" ) x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=value) return F.erase(img, x, y, h, w, v, self.inplace) return img def __repr__(self) -> str: s = ( f"{self.__class__.__name__}" f"(p={self.p}, " f"scale={self.scale}, " f"ratio={self.ratio}, " f"value={self.value}, " f"inplace={self.inplace})" ) return s class GaussianBlur(torch.nn.Module): """Blurs image with randomly chosen Gaussian blur. If the image is torch Tensor, it is expected to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions. Args: kernel_size (int or sequence): Size of the Gaussian kernel. sigma (float or tuple of float (min, max)): Standard deviation to be used for creating kernel to perform blurring. If float, sigma is fixed. If it is tuple of float (min, max), sigma is chosen uniformly at random to lie in the given range. Returns: PIL Image or Tensor: Gaussian blurred version of the input image. """ def __init__(self, kernel_size, sigma=(0.1, 2.0)): super().__init__() # _log_api_usage_once(self) self.kernel_size = _setup_size(kernel_size, "Kernel size should be a tuple/list of two integers") for ks in self.kernel_size: if ks <= 0 or ks % 2 == 0: raise ValueError("Kernel size value should be an odd and positive number.") if isinstance(sigma, numbers.Number): if sigma <= 0: raise ValueError("If sigma is a single number, it must be positive.") sigma = (sigma, sigma) elif isinstance(sigma, Sequence) and len(sigma) == 2: if not 0.0 < sigma[0] <= sigma[1]: raise ValueError("sigma values should be positive and of the form (min, max).") else: raise ValueError("sigma should be a single number or a list/tuple with length 2.") self.sigma = sigma @staticmethod def get_params(sigma_min: float, sigma_max: float) -> float: """Choose sigma for random gaussian blurring. Args: sigma_min (float): Minimum standard deviation that can be chosen for blurring kernel. sigma_max (float): Maximum standard deviation that can be chosen for blurring kernel. Returns: float: Standard deviation to be passed to calculate kernel for gaussian blurring. """ return torch.empty(1).uniform_adpater(sigma_min, sigma_max).item() def forward(self, img: Tensor) -> Tensor: """ Args: img (PIL Image or Tensor): image to be blurred. Returns: PIL Image or Tensor: Gaussian blurred image """ sigma = self.get_params(self.sigma[0], self.sigma[1]) return F.gaussian_blur(img, self.kernel_size, [sigma, sigma]) def __repr__(self) -> str: s = f"{self.__class__.__name__}(kernel_size={self.kernel_size}, sigma={self.sigma})" return s def _setup_size(size, error_msg): if isinstance(size, numbers.Number): return int(size), int(size) if isinstance(size, Sequence) and len(size) == 1: return size[0], size[0] if len(size) != 2: raise ValueError(error_msg) return size def _check_sequence_input(x, name, req_sizes): msg = req_sizes[0] if len(req_sizes) < 2 else " or ".join([str(s) for s in req_sizes]) if not isinstance(x, Sequence): raise TypeError(f"{name} should be a sequence of length {msg}.") if len(x) not in req_sizes: raise ValueError(f"{name} should be sequence of length {msg}.") def _setup_angle(x, name, req_sizes=(2,)): if isinstance(x, numbers.Number): if x < 0: raise ValueError(f"If {name} is a single number, it must be positive.") x = [-x, x] else: _check_sequence_input(x, name, req_sizes) return [float(d) for d in x] class RandomInvert(torch.nn.Module): """Inverts the colors of the given image randomly with a given probability. If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: p (float): probability of the image being color inverted. Default value is 0.5 """ def __init__(self, p=0.5): super().__init__() # _log_api_usage_once(self) self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be inverted. Returns: PIL Image or Tensor: Randomly color inverted image. """ if torch.rand(1).item() < self.p: return F.invert(img) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" class RandomPosterize(torch.nn.Module): """Posterize the image randomly with a given probability by reducing the number of bits for each color channel. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: bits (int): number of bits to keep for each channel (0-8) p (float): probability of the image being posterized. Default value is 0.5 """ def __init__(self, bits, p=0.5): super().__init__() # _log_api_usage_once(self) self.bits = bits self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be posterized. Returns: PIL Image or Tensor: Randomly posterized image. """ if torch.rand(1).item() < self.p: return F.posterize(img, self.bits) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(bits={self.bits},p={self.p})" class RandomSolarize(torch.nn.Module): """Solarize the image randomly with a given probability by inverting all pixel values above a threshold. If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format, where ... means it can have an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: threshold (float): all pixels equal or above this value are inverted. p (float): probability of the image being solarized. Default value is 0.5 """ def __init__(self, threshold, p=0.5): super().__init__() # _log_api_usage_once(self) self.threshold = threshold self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be solarized. Returns: PIL Image or Tensor: Randomly solarized image. """ if torch.rand(1).item() < self.p: return F.solarize(img, self.threshold) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(threshold={self.threshold},p={self.p})" class RandomAdjustSharpness(torch.nn.Module): """Adjust the sharpness of the image randomly with a given probability. If the image is torch Tensor, it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. Args: sharpness_factor (float): How much to adjust the sharpness. Can be any non negative number. 0 gives a blurred image, 1 gives the original image while 2 increases the sharpness by a factor of 2. p (float): probability of the image being sharpened. Default value is 0.5 """ def __init__(self, sharpness_factor, p=0.5): super().__init__() # _log_api_usage_once(self) self.sharpness_factor = sharpness_factor self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be sharpened. Returns: PIL Image or Tensor: Randomly sharpened image. """ if torch.rand(1).item() < self.p: return F.adjust_sharpness(img, self.sharpness_factor) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(sharpness_factor={self.sharpness_factor},p={self.p})" class RandomAutocontrast(torch.nn.Module): """Autocontrast the pixels of the given image randomly with a given probability. If the image is torch Tensor, it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: p (float): probability of the image being autocontrasted. Default value is 0.5 """ def __init__(self, p=0.5): super().__init__() # _log_api_usage_once(self) self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be autocontrasted. Returns: PIL Image or Tensor: Randomly autocontrasted image. """ if torch.rand(1).item() < self.p: return F.autocontrast(img) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" class RandomEqualize(torch.nn.Module): """Equalize the histogram of the given image randomly with a given probability. If the image is torch Tensor, it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "P", "L" or "RGB". Args: p (float): probability of the image being equalized. Default value is 0.5 """ def __init__(self, p=0.5): super().__init__() # _log_api_usage_once(self) self.p = p def forward(self, img): """ Args: img (PIL Image or Tensor): Image to be equalized. Returns: PIL Image or Tensor: Randomly equalized image. """ if torch.rand(1).item() < self.p: return F.equalize(img) return img def __repr__(self) -> str: return f"{self.__class__.__name__}(p={self.p})" msadapter/ms_adapter/torchvision/utils.py000066400000000000000000000552601436713554100212570ustar00rootroot00000000000000# import math # import pathlib # import warnings # from types import FunctionType # from typing import Any, BinaryIO, List, Optional, Tuple, Union # # import numpy as np # import torch # from PIL import Image, ImageColor, ImageDraw, ImageFont # # __all__ = [ # "make_grid", # "save_image", # "draw_bounding_boxes", # "draw_segmentation_masks", # "draw_keypoints", # "flow_to_image", # ] # # # @torch.no_grad() # def make_grid( # tensor: Union[torch.Tensor, List[torch.Tensor]], # nrow: int = 8, # padding: int = 2, # normalize: bool = False, # value_range: Optional[Tuple[int, int]] = None, # scale_each: bool = False, # pad_value: float = 0.0, # **kwargs, # ) -> torch.Tensor: # """ # Make a grid of images. # # Args: # tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W) # or a list of images all of the same size. # nrow (int, optional): Number of images displayed in each row of the grid. # The final grid size is ``(B / nrow, nrow)``. Default: ``8``. # padding (int, optional): amount of padding. Default: ``2``. # normalize (bool, optional): If True, shift the image to the range (0, 1), # by the min and max values specified by ``value_range``. Default: ``False``. # value_range (tuple, optional): tuple (min, max) where min and max are numbers, # then these numbers are used to normalize the image. By default, min and max # are computed from the tensor. # range (tuple. optional): # .. warning:: # This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``value_range`` # instead. # scale_each (bool, optional): If ``True``, scale each image in the batch of # images separately rather than the (min, max) over all images. Default: ``False``. # pad_value (float, optional): Value for the padded pixels. Default: ``0``. # # Returns: # grid (Tensor): the tensor containing grid of images. # """ # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(make_grid) # if not torch.is_tensor(tensor): # if isinstance(tensor, list): # for t in tensor: # if not torch.is_tensor(t): # raise TypeError(f"tensor or list of tensors expected, got a list containing {type(t)}") # else: # raise TypeError(f"tensor or list of tensors expected, got {type(tensor)}") # # if "range" in kwargs.keys(): # warnings.warn( # "The parameter 'range' is deprecated since 0.12 and will be removed in 0.14. " # "Please use 'value_range' instead." # ) # value_range = kwargs["range"] # # # if list of tensors, convert to a 4D mini-batch Tensor # if isinstance(tensor, list): # tensor = torch.stack(tensor, dim=0) # # if tensor.dim() == 2: # single image H x W # tensor = tensor.unsqueeze(0) # if tensor.dim() == 3: # single image # if tensor.size(0) == 1: # if single-channel, convert to 3-channel # tensor = torch.cat((tensor, tensor, tensor), 0) # tensor = tensor.unsqueeze(0) # # if tensor.dim() == 4 and tensor.size(1) == 1: # single-channel images # tensor = torch.cat((tensor, tensor, tensor), 1) # # if normalize is True: # tensor = tensor.clone() # avoid modifying tensor in-place # if value_range is not None and not isinstance(value_range, tuple): # raise TypeError("value_range has to be a tuple (min, max) if specified. min and max are numbers") # # def norm_ip(img, low, high): # img.clamp_(min=low, max=high) # img.sub_(low).div_(max(high - low, 1e-5)) # # def norm_range(t, value_range): # if value_range is not None: # norm_ip(t, value_range[0], value_range[1]) # else: # norm_ip(t, float(t.min()), float(t.max())) # # if scale_each is True: # for t in tensor: # loop over mini-batch dimension # norm_range(t, value_range) # else: # norm_range(tensor, value_range) # # if not isinstance(tensor, torch.Tensor): # raise TypeError("tensor should be of type torch.Tensor") # if tensor.size(0) == 1: # return tensor.squeeze(0) # # # make the mini-batch of images into a grid # nmaps = tensor.size(0) # xmaps = min(nrow, nmaps) # ymaps = int(math.ceil(float(nmaps) / xmaps)) # height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding) # num_channels = tensor.size(1) # grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value) # k = 0 # for y in range(ymaps): # for x in range(xmaps): # if k >= nmaps: # break # # Tensor.copy_() is a valid method but seems to be missing from the stubs # # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_ # grid.narrow(1, y * height + padding, height - padding).narrow( # type: ignore[attr-defined] # 2, x * width + padding, width - padding # ).copy_(tensor[k]) # k = k + 1 # return grid # # # @torch.no_grad() # def save_image( # tensor: Union[torch.Tensor, List[torch.Tensor]], # fp: Union[str, pathlib.Path, BinaryIO], # format: Optional[str] = None, # **kwargs, # ) -> None: # """ # Save a given Tensor into an image file. # # Args: # tensor (Tensor or list): Image to be saved. If given a mini-batch tensor, # saves the tensor as a grid of images by calling ``make_grid``. # fp (string or file object): A filename or a file object # format(Optional): If omitted, the format to use is determined from the filename extension. # If a file object was used instead of a filename, this parameter should always be used. # **kwargs: Other arguments are documented in ``make_grid``. # """ # # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(save_image) # grid = make_grid(tensor, **kwargs) # # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer # ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() # im = Image.fromarray(ndarr) # im.save(fp, format=format) # # # @torch.no_grad() # def draw_bounding_boxes( # image: torch.Tensor, # boxes: torch.Tensor, # labels: Optional[List[str]] = None, # colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None, # fill: Optional[bool] = False, # width: int = 1, # font: Optional[str] = None, # font_size: Optional[int] = None, # ) -> torch.Tensor: # # """ # Draws bounding boxes on given image. # The values of the input image should be uint8 between 0 and 255. # If fill is True, Resulting Tensor should be saved as PNG image. # # Args: # image (Tensor): Tensor of shape (C x H x W) and dtype uint8. # boxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format. Note that # the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and # `0 <= ymin < ymax < H`. # labels (List[str]): List containing the labels of bounding boxes. # colors (color or list of colors, optional): List containing the colors # of the boxes or single color for all boxes. The color can be represented as # PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``. # By default, random colors are generated for boxes. # fill (bool): If `True` fills the bounding box with specified color. # width (int): Width of bounding box. # font (str): A filename containing a TrueType font. If the file is not found in this filename, the loader may # also search in other directories, such as the `fonts/` directory on Windows or `/Library/Fonts/`, # `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS. # font_size (int): The requested font size in points. # # Returns: # img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted. # """ # # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(draw_bounding_boxes) # if not isinstance(image, torch.Tensor): # raise TypeError(f"Tensor expected, got {type(image)}") # elif image.dtype != torch.uint8: # raise ValueError(f"Tensor uint8 expected, got {image.dtype}") # elif image.dim() != 3: # raise ValueError("Pass individual images, not batches") # elif image.size(0) not in {1, 3}: # raise ValueError("Only grayscale and RGB images are supported") # # num_boxes = boxes.shape[0] # # if num_boxes == 0: # warnings.warn("boxes doesn't contain any box. No box was drawn") # return image # # if labels is None: # labels: Union[List[str], List[None]] = [None] * num_boxes # type: ignore[no-redef] # elif len(labels) != num_boxes: # raise ValueError( # f"Number of boxes ({num_boxes}) and labels ({len(labels)}) mismatch. Please specify labels for each box." # ) # # if colors is None: # colors = _generate_color_palette(num_boxes) # elif isinstance(colors, list): # if len(colors) < num_boxes: # raise ValueError(f"Number of colors ({len(colors)}) is less than number of boxes ({num_boxes}). ") # else: # colors specifies a single color for all boxes # colors = [colors] * num_boxes # # colors = [(ImageColor.getrgb(color) if isinstance(color, str) else color) for color in colors] # # if font is None: # if font_size is not None: # warnings.warn("Argument 'font_size' will be ignored since 'font' is not set.") # txt_font = ImageFont.load_default() # else: # txt_font = ImageFont.truetype(font=font, size=font_size or 10) # # # Handle Grayscale images # if image.size(0) == 1: # image = torch.tile(image, (3, 1, 1)) # # ndarr = image.permute(1, 2, 0).cpu().numpy() # img_to_draw = Image.fromarray(ndarr) # img_boxes = boxes.to(torch.int64).tolist() # # if fill: # draw = ImageDraw.Draw(img_to_draw, "RGBA") # else: # draw = ImageDraw.Draw(img_to_draw) # # for bbox, color, label in zip(img_boxes, colors, labels): # type: ignore[arg-type] # if fill: # fill_color = color + (100,) # draw.rectangle(bbox, width=width, outline=color, fill=fill_color) # else: # draw.rectangle(bbox, width=width, outline=color) # # if label is not None: # margin = width + 1 # draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=color, font=txt_font) # # return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8) # # # @torch.no_grad() # def draw_segmentation_masks( # image: torch.Tensor, # masks: torch.Tensor, # alpha: float = 0.8, # colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None, # ) -> torch.Tensor: # # """ # Draws segmentation masks on given RGB image. # The values of the input image should be uint8 between 0 and 255. # # Args: # image (Tensor): Tensor of shape (3, H, W) and dtype uint8. # masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool. # alpha (float): Float number between 0 and 1 denoting the transparency of the masks. # 0 means full transparency, 1 means no transparency. # colors (color or list of colors, optional): List containing the colors # of the masks or single color for all masks. The color can be represented as # PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``. # By default, random colors are generated for each mask. # # Returns: # img (Tensor[C, H, W]): Image Tensor, with segmentation masks drawn on top. # """ # # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(draw_segmentation_masks) # if not isinstance(image, torch.Tensor): # raise TypeError(f"The image must be a tensor, got {type(image)}") # elif image.dtype != torch.uint8: # raise ValueError(f"The image dtype must be uint8, got {image.dtype}") # elif image.dim() != 3: # raise ValueError("Pass individual images, not batches") # elif image.size()[0] != 3: # raise ValueError("Pass an RGB image. Other Image formats are not supported") # if masks.ndim == 2: # masks = masks[None, :, :] # if masks.ndim != 3: # raise ValueError("masks must be of shape (H, W) or (batch_size, H, W)") # if masks.dtype != torch.bool: # raise ValueError(f"The masks must be of dtype bool. Got {masks.dtype}") # if masks.shape[-2:] != image.shape[-2:]: # raise ValueError("The image and the masks must have the same height and width") # # num_masks = masks.size()[0] # if colors is not None and num_masks > len(colors): # raise ValueError(f"There are more masks ({num_masks}) than colors ({len(colors)})") # # if num_masks == 0: # warnings.warn("masks doesn't contain any mask. No mask was drawn") # return image # # if colors is None: # colors = _generate_color_palette(num_masks) # # if not isinstance(colors, list): # colors = [colors] # if not isinstance(colors[0], (tuple, str)): # raise ValueError("colors must be a tuple or a string, or a list thereof") # if isinstance(colors[0], tuple) and len(colors[0]) != 3: # raise ValueError("It seems that you passed a tuple of colors instead of a list of colors") # # out_dtype = torch.uint8 # # colors_ = [] # for color in colors: # if isinstance(color, str): # color = ImageColor.getrgb(color) # colors_.append(torch.tensor(color, dtype=out_dtype)) # # img_to_draw = image.detach().clone() # # TODO: There might be a way to vectorize this # for mask, color in zip(masks, colors_): # img_to_draw[:, mask] = color[:, None] # # out = image * (1 - alpha) + img_to_draw * alpha # return out.to(out_dtype) # # # @torch.no_grad() # def draw_keypoints( # image: torch.Tensor, # keypoints: torch.Tensor, # connectivity: Optional[List[Tuple[int, int]]] = None, # colors: Optional[Union[str, Tuple[int, int, int]]] = None, # radius: int = 2, # width: int = 3, # ) -> torch.Tensor: # # """ # Draws Keypoints on given RGB image. # The values of the input image should be uint8 between 0 and 255. # # Args: # image (Tensor): Tensor of shape (3, H, W) and dtype uint8. # keypoints (Tensor): Tensor of shape (num_instances, K, 2) the K keypoints location for each of the N instances, # in the format [x, y]. # connectivity (List[Tuple[int, int]]]): A List of tuple where, # each tuple contains pair of keypoints to be connected. # colors (str, Tuple): The color can be represented as # PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``. # radius (int): Integer denoting radius of keypoint. # width (int): Integer denoting width of line connecting keypoints. # # Returns: # img (Tensor[C, H, W]): Image Tensor of dtype uint8 with keypoints drawn. # """ # # if not torch.jit.is_scripting() and not torch.jit.is_tracing(): # _log_api_usage_once(draw_keypoints) # if not isinstance(image, torch.Tensor): # raise TypeError(f"The image must be a tensor, got {type(image)}") # elif image.dtype != torch.uint8: # raise ValueError(f"The image dtype must be uint8, got {image.dtype}") # elif image.dim() != 3: # raise ValueError("Pass individual images, not batches") # elif image.size()[0] != 3: # raise ValueError("Pass an RGB image. Other Image formats are not supported") # # if keypoints.ndim != 3: # raise ValueError("keypoints must be of shape (num_instances, K, 2)") # # ndarr = image.permute(1, 2, 0).cpu().numpy() # img_to_draw = Image.fromarray(ndarr) # draw = ImageDraw.Draw(img_to_draw) # img_kpts = keypoints.to(torch.int64).tolist() # # for kpt_id, kpt_inst in enumerate(img_kpts): # for inst_id, kpt in enumerate(kpt_inst): # x1 = kpt[0] - radius # x2 = kpt[0] + radius # y1 = kpt[1] - radius # y2 = kpt[1] + radius # draw.ellipse([x1, y1, x2, y2], fill=colors, outline=None, width=0) # # if connectivity: # for connection in connectivity: # start_pt_x = kpt_inst[connection[0]][0] # start_pt_y = kpt_inst[connection[0]][1] # # end_pt_x = kpt_inst[connection[1]][0] # end_pt_y = kpt_inst[connection[1]][1] # # draw.line( # ((start_pt_x, start_pt_y), (end_pt_x, end_pt_y)), # width=width, # ) # # return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8) # # # # Flow visualization code adapted from https://github.com/tomrunia/OpticalFlow_Visualization # @torch.no_grad() # def flow_to_image(flow: torch.Tensor) -> torch.Tensor: # # """ # Converts a flow to an RGB image. # # Args: # flow (Tensor): Flow of shape (N, 2, H, W) or (2, H, W) and dtype torch.float. # # Returns: # img (Tensor): Image Tensor of dtype uint8 where each color corresponds # to a given flow direction. Shape is (N, 3, H, W) or (3, H, W) depending on the input. # """ # # if flow.dtype != torch.float: # raise ValueError(f"Flow should be of dtype torch.float, got {flow.dtype}.") # # orig_shape = flow.shape # if flow.ndim == 3: # flow = flow[None] # Add batch dim # # if flow.ndim != 4 or flow.shape[1] != 2: # raise ValueError(f"Input flow should have shape (2, H, W) or (N, 2, H, W), got {orig_shape}.") # # max_norm = torch.sum(flow ** 2, dim=1).sqrt().max() # epsilon = torch.finfo((flow).dtype).eps # normalized_flow = flow / (max_norm + epsilon) # img = _normalized_flow_to_image(normalized_flow) # # if len(orig_shape) == 3: # img = img[0] # Remove batch dim # return img # # # @torch.no_grad() # def _normalized_flow_to_image(normalized_flow: torch.Tensor) -> torch.Tensor: # # """ # Converts a batch of normalized flow to an RGB image. # # Args: # normalized_flow (torch.Tensor): Normalized flow tensor of shape (N, 2, H, W) # Returns: # img (Tensor(N, 3, H, W)): Flow visualization image of dtype uint8. # """ # # N, _, H, W = normalized_flow.shape # device = normalized_flow.device # flow_image = torch.zeros((N, 3, H, W), dtype=torch.uint8, device=device) # colorwheel = _make_colorwheel().to(device) # shape [55x3] # num_cols = colorwheel.shape[0] # norm = torch.sum(normalized_flow ** 2, dim=1).sqrt() # a = torch.atan2(-normalized_flow[:, 1, :, :], -normalized_flow[:, 0, :, :]) / torch.pi # fk = (a + 1) / 2 * (num_cols - 1) # k0 = torch.floor(fk).to(torch.long) # k1 = k0 + 1 # k1[k1 == num_cols] = 0 # f = fk - k0 # # for c in range(colorwheel.shape[1]): # tmp = colorwheel[:, c] # col0 = tmp[k0] / 255.0 # col1 = tmp[k1] / 255.0 # col = (1 - f) * col0 + f * col1 # col = 1 - norm * (1 - col) # flow_image[:, c, :, :] = torch.floor(255 * col) # return flow_image # # # def _make_colorwheel() -> torch.Tensor: # """ # Generates a color wheel for optical flow visualization as presented in: # Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007) # URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf. # # Returns: # colorwheel (Tensor[55, 3]): Colorwheel Tensor. # """ # # RY = 15 # YG = 6 # GC = 4 # CB = 11 # BM = 13 # MR = 6 # # ncols = RY + YG + GC + CB + BM + MR # colorwheel = torch.zeros((ncols, 3)) # col = 0 # # # RY # colorwheel[0:RY, 0] = 255 # colorwheel[0:RY, 1] = torch.floor(255 * torch.arange(0, RY) / RY) # col = col + RY # # YG # colorwheel[col : col + YG, 0] = 255 - torch.floor(255 * torch.arange(0, YG) / YG) # colorwheel[col : col + YG, 1] = 255 # col = col + YG # # GC # colorwheel[col : col + GC, 1] = 255 # colorwheel[col : col + GC, 2] = torch.floor(255 * torch.arange(0, GC) / GC) # col = col + GC # # CB # colorwheel[col : col + CB, 1] = 255 - torch.floor(255 * torch.arange(CB) / CB) # colorwheel[col : col + CB, 2] = 255 # col = col + CB # # BM # colorwheel[col : col + BM, 2] = 255 # colorwheel[col : col + BM, 0] = torch.floor(255 * torch.arange(0, BM) / BM) # col = col + BM # # MR # colorwheel[col : col + MR, 2] = 255 - torch.floor(255 * torch.arange(MR) / MR) # colorwheel[col : col + MR, 0] = 255 # return colorwheel # # # def _generate_color_palette(num_objects: int): # palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1]) # return [tuple((i * palette) % 255) for i in range(num_objects)] # # # def _log_api_usage_once(obj: Any) -> None: # # """ # Logs API usage(module and name) within an organization. # In a large ecosystem, it's often useful to track the PyTorch and # TorchVision APIs usage. This API provides the similar functionality to the # logging module in the Python stdlib. It can be used for debugging purpose # to log which methods are used and by default it is inactive, unless the user # manually subscribes a logger via the `SetAPIUsageLogger method `_. # Please note it is triggered only once for the same API call within a process. # It does not collect any data from open-source users since it is no-op by default. # For more information, please refer to # * PyTorch note: https://pytorch.org/docs/stable/notes/large_scale_deployments.html#api-usage-logging; # * Logging policy: https://github.com/pytorch/vision/issues/5052; # # Args: # obj (class instance or method): an object to extract info from. # """ # if not obj.__module__.startswith("torchvision"): # return # name = obj.__class__.__name__ # if isinstance(obj, FunctionType): # name = obj.__name__ # torch._C._log_api_usage_once(f"{obj.__module__}.{name}") msadapter/ms_adapter/utils.py000066400000000000000000000052471436713554100167100ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms from mindspore import context from mindspore.ops import constexpr def unsupported_attr(attr): """ To mark the attribute that is not currently supported. """ return attr @constexpr def pynative_mode_condition(): return context.get_context("mode") == context.PYNATIVE_MODE @constexpr def graph_mode_condition(): return context.get_context("mode") == context.GRAPH_MODE @constexpr def get_backend(): return context.get_context("device_target") @constexpr def is_under_gpu_context(): return get_backend() == 'GPU' @constexpr def is_under_ascend_context(): return get_backend() == 'Ascend' _AscendGenernalConvertDict = { ms.float64: ms.float32, ms.int8: ms.float16, ms.int16: ms.float16, ms.int32: ms.float32, ms.int64: ms.float32, ms.uint8: ms.float16, ms.uint16: ms.float32, ms.uint32: ms.float32, ms.uint64: ms.float32, ms.bool_: ms.float16, ms.double: ms.float32, } def _ascend_tensor_general_cast(input, conver_dicts={}): """ Example: >>> import ms_adapter.pytorch as torch >>> from ms_adapter.utils import _ascend_tensor_general_cast >>> a = torch.tensor(2) >>> print(a.dtype) Int64 >>> b = _ascend_tensor_general_cast(a) >>> print(b.dtype) Float32 >>> c = _ascend_tensor_general_cast(a, conver_dicts={torch.int64: torch.int32}) >>> print(b.dtype) Int32 """ value = conver_dicts.get(input.dtype) if value: return input.astype(value) _to_dtype = _AscendGenernalConvertDict.get(input.dtype) if _to_dtype: return input.astype(_to_dtype) return input @constexpr def _infer_size(shape, numel): if len(shape) == 1 and isinstance(shape[0], tuple): shape = shape[0] dim = None newsize = 1 for i, d in enumerate(shape): if d == -1: if dim is not None: raise RuntimeError("only one dimension can be inferred") dim = i elif d >= 0: newsize *= d else: raise RuntimeError(f"invalid shape dimension {d}") if not (numel == newsize or (dim is not None and newsize > 0 and numel % newsize == 0)): raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {numel}") if dim is not None: if newsize == 0: raise RuntimeError(f"cannot reshape tensor fo 0 elements into shape {shape} because the " "unspecified dimension size -1 can be any value and is ambiguous.") shape = list(shape) shape[dim] = numel // newsize return tuple(shape)msadapter/pylint_check.sh000066400000000000000000000007241436713554100160620ustar00rootroot00000000000000#!/bin/bash set -e if ! type pylint >/dev/null 2>&1; then pip install pylint==2.8.2 echo "pylint is now installed." else PYLINT_VERSION="$(pylint --version)" echo ${PYLINT_VERSION} if [[ "$PYLINT_VERSION" =~ "pylint 2.8.2" ]]; then echo "pylint is already installed." else pip install pylint==2.8.2 echo "pylint is now installed." fi fi echo "pylint check is starting." pylint ./ms_adapter/pytorch --ignore=data msadapter/pylintrc000066400000000000000000000071151436713554100146420ustar00rootroot00000000000000[MASTER] # Profiled execution. profile=no # Pickle collected data for later comparisons. persistent=yes [MESSAGES CONTROL] # Enable the message, report, category or checker with the given id(s). You can # either give multiple identifier separated by comma (,) or put this option # multiple time. See also the "--disable" option for examples. enable=indexing-exception,old-raise-syntax # Disable the message, report, category or checker with the given id(s). You # can either give multiple identifiers separated by comma (,) or put this # option multiple times (only on the command line, not in the configuration # file where it should appear only once).You can also use "--disable=all" to # disable everything first and then reenable specific checks. For example, if # you want to run only the similarities checker, you can use "--disable=all # --enable=similarities". If you want to run only the classes checker, but have # no Warning level messages displayed, use"--disable=all --enable=classes # --disable=W" disable=design, similarities, no-self-use, attribute-defined-outside-init, locally-disabled, star-args, pointless-except, bad-option-value, global-statement, fixme, suppressed-message, useless-suppression, locally-enabled, no-member, no-name-in-module, unsubscriptable-object, unbalanced-tuple-unpacking, not-context-manager, missing-docstring, too-many-locals, too-many-arguments, too-many-lines, redefined-builtin, too-many-function-args, invalid-overridden-method, wildcard-import, self-cls-assignment, super-with-arguments, not-callable, no-else-return, invalid-name, abstract-method, arguments-differ, redefined-outer-name, global-variable-undefined, super-init-not-called, no-else-raise, protected-access, consider-using-f-string, unexpected-keyword-arg, consider-using-enumerate, no-value-for-parameter, [REPORTS] # Set the output format. Available formats are text, parseable, colorized, msvs # (visual studio) and html. You can also give a reporter class, eg # mypackage.mymodule.MyReporterClass. output-format=text # Put messages in a separate file for each module / package specified on the # command line instead of printing them on stdout. Reports (if any) will be # written in a file name "pylint_global.[txt|html]". files-output=no # Tells whether to display a full report or only the messages reports=no [FORMAT] # Maximum number of characters on a single line. max-line-length=120 # String used as indentation unit. This is usually " " (4 spaces) or "\t" (1 # tab). # 4 spaces (Default: 2 spaces) indent-string=' ' [CLASSES] # List of interface methods to ignore, separated by a comma. This is used for # instance to not check methods defines in Zope's Interface base class. ignore-iface-methods=isImplementedBy,deferred,extends,names,namesAndDescriptions,queryDescriptionFor,getBases,getDescriptionFor,getDoc,getName,getTaggedValue,getTaggedValueTags,isEqualOrExtendedBy,setTaggedValue,isImplementedByInstancesOf,adaptWith,is_implemented_by # List of method names used to declare (i.e. assign) instance attributes. defining-attr-methods=__init__,__new__,setUp # List of valid names for the first argument in a class method. valid-classmethod-first-arg=cls,class_ # List of valid names for the first argument in a metaclass class method. valid-metaclass-classmethod-first-arg=mcsmsadapter/requirements.txt000066400000000000000000000000511436713554100163270ustar00rootroot00000000000000# mindspore>=2.0.0 # mindspore-gpu>=2.0.0msadapter/run.sh000066400000000000000000000002321436713554100142040ustar00rootroot00000000000000#!/bin/bash set -e # install MSAdapter python setup.py install # run uint test bash ./testing/ut/runtest.sh # code pylint check bash ./pylint_check.sh msadapter/setup.py000066400000000000000000000050461436713554100145660ustar00rootroot00000000000000#! /usr/bin/python # -*- coding: utf-8 -*- from __future__ import absolute_import from setuptools import setup, find_packages import os, codecs MAJOR = 0 MINOR = 1 PATCH = 0 PRE_RELEASE = '' # Use the following formatting: (major, minor, patch, prerelease) VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE) if os.path.exists('README.rst'): # codec is used for consistent encoding long_description = codecs.open( os.path.join(os.path.abspath(os.path.dirname(__file__)), 'README.rst'), 'r', 'utf-8' ).read() else: long_description = "MSAdapter is a toolkit for support the pytorch model runing on ascend.\n" long_description += "Usage: import ms_adapter.pytorch as ms \n" long_description += "OpenI: https://openi.pcl.ac.cn/OpenI/MSAdapter \n" long_description += "Email: pcl.openi@pcl.ac.cn" def req_file(filename, folder=''): with open(os.path.join(folder, filename)) as f: content = f.readlines() return [x.strip() for x in content] setup( name="ms_adapter", version='.'.join(map(str, VERSION[:3])) + ''.join(VERSION[3:]), author="Peng Cheng Lab, HUAWEI", author_email="pcl.openi@pcl.ac.cn", description="MSAdapter is a toolkit for support the pytorch model runing on ascend.", long_description=long_description, long_description_content_type="text/plain", url="https://openi.pcl.ac.cn/OpenI/MSAdapter ", packages=find_packages(), install_requires=req_file("requirements.txt"), classifiers=[ # How mature is this project? Common values are # 1 - Planning 2 - Pre-Alpha 3 - Alpha 4 - Beta 5 - Production/Stable 6 - Mature 7 - Inactive 'Development Status :: 3 - Alpha', # Specify the Python versions you support here. 'Programming Language :: Python :: 3 :: Only', "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", # Indicate who your project is intended for "Intended Audience :: Developers", "Intended Audience :: Science/Research", "Intended Audience :: Information Technology", # Pick your license as you wish (should match "license" above) "License :: OSI Approved :: Apache Software License", # Indicate what your project relates to 'Topic :: Scientific/Engineering :: Artificial Intelligence', 'Topic :: Software Development :: Libraries', 'Topic :: Software Development :: Libraries :: Python Modules', # Additionnal Settings 'Operating System :: OS Independent', ], license='Apache 2.0', ) msadapter/testing/000077500000000000000000000000001436713554100145245ustar00rootroot00000000000000msadapter/testing/ut/000077500000000000000000000000001436713554100151545ustar00rootroot00000000000000msadapter/testing/ut/pytorch/000077500000000000000000000000001436713554100166445ustar00rootroot00000000000000msadapter/testing/ut/pytorch/autograd/000077500000000000000000000000001436713554100204525ustar00rootroot00000000000000msadapter/testing/ut/pytorch/autograd/test_autograd_function.py000066400000000000000000000037241436713554100256040ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.PYNATIVE_MODE) def adapter_autograd_function(): class Net(ms_torch.nn.Module): def __init__(self): super(Net, self).__init__() def forward(self, x, y): out = ms_torch.matmul(x, y) return out # bprop: https://www.mindspore.cn/tutorials/experts/zh-CN/r1.9/network/custom_cell_reverse.html def bprop(self, x, y, out, dout): dx = x + 1 dy = y + 1 return dx, dy x = ms_torch.tensor([[0.5, 0.6, 0.4]], dtype=ms_torch.float32) y = ms_torch.tensor([[0.01], [0.2], [3.3]], dtype=ms_torch.float32) net = Net() out = net(x, y) grad_out = ms.grad(net, grad_position=(0, 1))(x, y) return out, grad_out def torch_autograd_function(): class Net(torch.autograd.Function): @staticmethod def forward(ctx, x, y): result = torch.matmul(x, y) ctx.save_for_backward(x, y) return result @staticmethod def backward(ctx, grad_output): x, y = ctx.saved_tensors dx = x + 1 dy = y + 1 return dx, dy x = torch.tensor([[0.5, 0.6, 0.4]], dtype=torch.float32, requires_grad=True) y = torch.tensor([[0.01], [0.2], [3.3]], dtype=torch.float32, requires_grad=True) out = Net.apply(x, y) out.backward() grad_out = x.grad, y.grad return out, grad_out def test_autograd_funciton(): ms_out, ms_grad_out = adapter_autograd_function() pt_out, pt_grad_out = torch_autograd_function() assert np.allclose(ms_out.asnumpy(), pt_out.detach().numpy()) assert np.allclose(ms_grad_out[0].asnumpy(), pt_grad_out[0].numpy()) assert np.allclose(ms_grad_out[1].asnumpy(), pt_grad_out[1].numpy()) if __name__ == '__main__': test_autograd_funciton() msadapter/testing/ut/pytorch/common/000077500000000000000000000000001436713554100201345ustar00rootroot00000000000000msadapter/testing/ut/pytorch/common/test_inplace.py000066400000000000000000000045231436713554100231640ustar00rootroot00000000000000import numpy as np import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_torch import ms_adapter.pytorch.nn as nn def test_tensor_inplace(): x = ms_torch.tensor([1.0, 2.0]) y = -1.0 try: x1 = ms_torch.tensor([1.0, 2.0]) context.set_context(mode=ms.GRAPH_MODE) out = x1.add_(y) # assert (id(out) != id(x)) assert (id(out) == id(x1)) except RuntimeError as e: assert "is not encouraged to use in MindSpore static graph mode" in str(e) context.set_context(mode=ms.PYNATIVE_MODE) out = x.add_(y) assert (id(out) == id(x)) assert np.allclose(out.numpy(), [0., 1.0]) def test_nn_inplace(): input = ms_torch.randn(2) try: context.set_context(mode=ms.GRAPH_MODE) nn_relu = nn.ReLU(inplace=True) output = nn_relu(input) assert (id(output) != id(input)) except ValueError as e: assert "please set inplace=False and use return value instead" in str(e) context.set_context(mode=ms.PYNATIVE_MODE) nn_relu = nn.ReLU(inplace=True) output = nn_relu(input) assert (id(output) == id(input)) assert np.allclose(output.numpy(), input.numpy()) def test_out_param(): x = ms_torch.tensor([10.0, 2.0]) out = ms_torch.tensor([1.0, 1.0]) try: context.set_context(mode=ms.GRAPH_MODE) output = ms_torch.log(x, out=out) assert (id(output) != id(out)) except ValueError as e: assert "please set out=None and use return value instead" in str(e) context.set_context(mode=ms.PYNATIVE_MODE) output = ms_torch.log(x, out=out) assert (id(out) == id(output)) assert np.allclose(output.numpy(), out.numpy()) def test_inplace_param(): x = ms_torch.tensor([-1, -2, 0, 2, 1], dtype=ms_torch.float16) try: context.set_context(mode=ms.GRAPH_MODE) output = nn.functional.relu6(x, inplace=True) assert (id(output) != id(x)) except ValueError as e: assert "please set inplace=False and use return value instead" in str(e) context.set_context(mode=ms.PYNATIVE_MODE) output = nn.functional.relu6(x, inplace=True) assert (id(x) == id(output)) assert np.allclose(output.numpy(), x.numpy()) if __name__ == '__main__': test_tensor_inplace() test_nn_inplace() test_out_param() test_inplace_param() msadapter/testing/ut/pytorch/data/000077500000000000000000000000001436713554100175555ustar00rootroot00000000000000msadapter/testing/ut/pytorch/data/__init__.py000066400000000000000000000000001436713554100216540ustar00rootroot00000000000000msadapter/testing/ut/pytorch/data/test_dataloader.py000066400000000000000000001466671436713554100233120ustar00rootroot00000000000000# Owner(s): ["module: dataloader"] import math import sys import errno import os import ctypes import faulthandler import gc import time import signal import unittest import warnings import multiprocessing as mp from ms_adapter.pytorch.utils.data import ( ChainDataset, ConcatDataset, DataLoader, # DataLoader2, Dataset, IterableDataset, # IterDataPipe, Subset, TensorDataset, # communication, _utils ) import ms_adapter.pytorch as torch import mindspore as ms from ms_adapter.pytorch._utils import ExceptionWrapper from ms_adapter.pytorch.utils.data.dataset import random_split IS_WINDOWS = sys.platform == "win32" NO_MULTIPROCESSING_SPAWN = os.environ.get('NO_MULTIPROCESSING_SPAWN', '0') == '1' IS_CI = bool(os.getenv('CI')) try: import psutil HAS_PSUTIL = True except ImportError: HAS_PSUTIL = False err_msg = ("psutil not found. Some critical data loader tests relying on it " "(e.g., TestDataLoader.test_proper_exit) will not run.") if IS_CI: raise ImportError(err_msg) from None else: warnings.warn(err_msg) try: import dill # XXX: By default, dill writes the Pickler dispatch table to inject its # own logic there. This globally affects the behavior of the standard library # pickler for any user who transitively depends on this module! # Undo this extension to avoid altering the behavior of the pickler globally. dill.extend(use_dill=False) HAS_DILL = True except ImportError: HAS_DILL = False skipIfNoDill = unittest.skipIf(not HAS_DILL, "no dill") try: import numpy as np HAS_NUMPY = True except ImportError: HAS_NUMPY = False skipIfNoNumpy = unittest.skipIf(not HAS_NUMPY, "no NumPy") JOIN_TIMEOUT = 60.0 # seconds supported_multiprocessing_contexts = [None] + list(mp.get_all_start_methods()) np.random.seed(42) ms.set_seed(42) def test_splits_have_correct_size(): splits = random_split([1, 2, 3, 4, 5, 6], [2, 4]) assert np.allclose(len(splits), 2) assert np.allclose(len(splits[0]), 2) assert np.allclose(len(splits[1]), 4) splits = random_split([1, 2, 3, 4, 5, 6], [0.5, 0.5]) assert np.allclose(len(splits), 2) assert np.allclose(len(splits[0]), 3) assert np.allclose(len(splits[1]), 3) # Odd size splits assert np.allclose( len(random_split(range(3), [0.5, 0.5])), 2 ) # Odd sized round-robin splits splits = random_split(range(106), [0.1, 0.2, 0.3, 0.4],) assert np.allclose(len(splits[0]), 11) assert np.allclose(len(splits[1]), 22) assert np.allclose(len(splits[2]), 31) assert np.allclose(len(splits[3]), 42) def test_splits_are_mutually_exclusive(): data = [5, 2, 3, 4, 1, 6] splits = random_split(data, [2, 4]) all_values = [] all_values.extend(list(splits[0])) all_values.extend(list(splits[1])) data.sort() all_values.sort() assert np.allclose(data, all_values) splits = random_split(data, [0.33, 0.67]) all_values = [] all_values.extend(list(splits[0])) all_values.extend(list(splits[1])) data.sort() all_values.sort() assert np.allclose(data, all_values) data = [1, 2, 3, 4] splits = random_split(data, [0.25, 0.75]) all_values = [] all_values.extend(list(splits[0])) all_values.extend(list(splits[1])) data.sort() all_values.sort() assert np.allclose(data, all_values) def test_splits_indexing_type(): class CustomDataset(): def __init__(self, custom_list): self.data = custom_list def __getitem__(self, key): if type(key) == type(0): return self.data[key] else: raise TypeError("Type do not match.") def __len__(self): return len(self.data) x = [1, 2, 3, 4, 5] dataset = CustomDataset(x) dataset = random_split(dataset, [5])[0] data_loader = DataLoader(dataset) for batch in data_loader: pass # fractional splitting dataset = CustomDataset(x) dataset = random_split(dataset, [1.0])[0] data_loader = DataLoader(dataset) for batch in data_loader: pass def test_slicing_of_subset_of_dataset(): # Testing slicing a subset initialized with a dataset dataset = TensorDataset(torch.tensor([1, 2, 3, 4, 5])) subset_of_dataset = Subset(dataset, [0, 1, 2, 3, 4]) assert np.allclose(subset_of_dataset[:][0].asnumpy(), dataset[:][0].asnumpy()) assert np.allclose(subset_of_dataset[1:2][0].asnumpy(), dataset[1:2][0].asnumpy()) assert np.allclose(subset_of_dataset[0:-1:2][0].asnumpy(), dataset[0:-1:2][0].asnumpy()) # Testing slicing of subset from random split subset1, subset2 = random_split(dataset, [3, 2]) assert np.allclose(subset1[:][0].asnumpy(), dataset[subset1.indices[:]][0].asnumpy()) assert np.allclose(subset1[0:2][0].asnumpy(), dataset[subset1.indices[0:2]][0].asnumpy()) assert np.allclose(subset1[0:-1:2][0].asnumpy(), dataset[subset1.indices[0:-1:2]][0].asnumpy()) def test_slicing_of_subset_of_subset(): # Testing slicing a subset initialized with a subset dataset = TensorDataset(torch.tensor([1, 2, 3, 4, 5])) subset_of_dataset = Subset(dataset, [0, 1, 2, 3, 4]) subset_of_subset = Subset(subset_of_dataset, [0, 1, 2, 3, 4]) assert np.allclose(subset_of_subset[:][0].asnumpy(), dataset[:][0].asnumpy()) assert np.allclose(subset_of_subset[0:2][0].asnumpy(), dataset[0:2][0].asnumpy()) assert np.allclose(subset_of_subset[0:-1:2][0].asnumpy(), dataset[0:-1:2][0].asnumpy()) # Testing slicing of subset of subset from random split subset1, subset2 = random_split(dataset, [4, 1]) subset_of_subset1, subset_of_subset2 = random_split(subset1, [3, 1]) idx = [subset1.indices[i] for i in subset_of_subset1.indices] assert np.allclose(subset_of_subset1[:][0].asnumpy(), dataset[idx[:]][0].asnumpy()) assert np.allclose(subset_of_subset1[0:2][0].asnumpy(), dataset[idx[0:2]][0].asnumpy()) assert np.allclose(subset_of_subset1[0:-1:2][0].asnumpy(), dataset[idx[0:-1:2]][0].asnumpy()) class CountingDataset(Dataset): def __init__(self, n): super(CountingDataset, self).__init__() self.n = n def __getitem__(self, i): return i def __len__(self): return self.n class CountingIterableDataset(IterableDataset): def __init__(self, n): super(CountingIterableDataset, self).__init__() self.n = n def __iter__(self): return iter(range(self.n)) def __len__(self): return self.n def test_len(): source = TensorDataset(torch.randn(15, 10, 2, 3, 4, 5), torch.randperm(15)) assert np.allclose(len(source), 15) def test_getitem(): t = torch.randn(15, 10, 2, 3, 4, 5) l = torch.randn(15, 10) source = TensorDataset(t, l) for i in range(15): assert np.allclose(t[i].asnumpy(), source[i][0].asnumpy()) assert np.allclose(l[i].asnumpy(), source[i][1].asnumpy()) def test_getitem_1d(): t = torch.randn(15) l = torch.randn(15) source = TensorDataset(t, l) for i in range(15): assert np.allclose(t[i].asnumpy(), source[i][0].asnumpy()) assert np.allclose(l[i].asnumpy(), source[i][1].asnumpy()) def test_single_tensor(): t = torch.randn(5, 10) source = TensorDataset(t) assert np.allclose(len(source), 5) for i in range(5): assert np.allclose(t[i].asnumpy(), source[i][0].asnumpy()) def test_many_tensors(): t0 = torch.randn(5, 10, 2, 3, 4, 5) t1 = torch.randn(5, 10) t2 = torch.randn(5, 10, 2, 5) t3 = torch.randn(5, 10, 3, 7) source = TensorDataset(t0, t1, t2, t3) assert np.allclose(len(source), 5) for i in range(5): assert np.allclose(t0[i].asnumpy(), source[i][0].asnumpy()) assert np.allclose(t1[i].asnumpy(), source[i][1].asnumpy()) assert np.allclose(t2[i].asnumpy(), source[i][2].asnumpy()) assert np.allclose(t3[i].asnumpy(), source[i][3].asnumpy()) def test_concat_two_singletons(): result = ConcatDataset([[0], [1]]) assert np.allclose(2, len(result)) assert np.allclose(0, result[0]) assert np.allclose(1, result[1]) def test_concat_two_non_singletons(): result = ConcatDataset([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]) assert np.allclose(10, len(result)) assert np.allclose(0, result[0]) assert np.allclose(5, result[5]) def test_concat_two_non_singletons_with_empty(): # Adding an empty dataset somewhere is correctly handled result = ConcatDataset([[0, 1, 2, 3, 4], [], [5, 6, 7, 8, 9]]) assert np.allclose(10, len(result)) assert np.allclose(0, result[0]) assert np.allclose(5, result[5]) def test_add_dataset(): d1 = TensorDataset(torch.randn(7, 3, 28, 28), torch.randn(7)) d2 = TensorDataset(torch.randn(7, 3, 28, 28), torch.randn(7)) d3 = TensorDataset(torch.randn(7, 3, 28, 28), torch.randn(7)) result = d1 + d2 + d3 assert np.allclose(21, len(result)) assert np.allclose(0, (d1[0][0] - result[0][0]).abs().sum().asnumpy()) assert np.allclose(0, (d2[0][0] - result[7][0]).abs().sum().asnumpy()) assert np.allclose(0, (d3[0][0] - result[14][0]).abs().sum().asnumpy()) # takes in dummy var so this can also be used as a `worker_init_fn` def set_faulthander_if_available(_=None): faulthandler.enable(sys.__stderr__) if not IS_WINDOWS: # windows does not have faulthandler.register # chain=False prevents the default behavior of killing the process faulthandler.register(signal.SIGUSR1, file=sys.__stderr__, chain=False) set_faulthander_if_available() # Process `pid` must have called `set_faulthander_if_available` def print_traces_of_all_threads(pid): if not IS_WINDOWS: # use the custom signal if available os.kill(pid, signal.SIGUSR1) else: # otherwise we can still use the handler given by faulthandler.enable() # at the cost of killing the process. os.kill(pid, signal.SIGSEGV) # wait in parent process to give subprocess some time to print time.sleep(5) class ErrorTrackingProcess(mp.Process): # Why no *args? # py2 doesn't support def fn(x, *args, key=val, **kwargs) # Setting disable_stderr=True may generate a lot of unrelated error outputs # but could be helpful for debugging. def __init__(self, disable_stderr=True, **kwargs): super(ErrorTrackingProcess, self).__init__(**kwargs) self._pconn, self._cconn = mp.Pipe() self._exception = None self.disable_stderr = disable_stderr def run(self): set_faulthander_if_available() if self.disable_stderr: # Disable polluting stderr with errors that are supposed to happen. with open(os.devnull, 'w') as devnull: os.dup2(devnull.fileno(), sys.stderr.fileno()) try: super(ErrorTrackingProcess, self).run() self._cconn.send(None) except Exception: self._cconn.send(ExceptionWrapper(sys.exc_info())) raise def print_traces_of_all_threads(self): assert self.is_alive(), "can only use print_traces_of_all_threads if the process is alive" assert not self.disable_stderr, "do not disable stderr if you use print_traces_of_all_threads" # On platforms without `SIGUSR1`, `set_faulthander_if_available` sets # `faulthandler.enable()`, and `print_traces_of_all_threads` may kill # the process. So let's poll the exception first _ = self.exception print_traces_of_all_threads(self.pid) @property def exception(self): if self._pconn.poll(): self._exception = self._pconn.recv() if self._exception is None: return None else: return self._exception.exc_type(self._exception.exc_msg) # ESRCH means that os.kill can't finds alive proc def send_signal(self, signum, ignore_ESRCH=False): try: os.kill(self.pid, signum) except OSError as e: if not ignore_ESRCH or e.errno != errno.ESRCH: raise class ErrorDataset(Dataset): def __init__(self, size): self.size = size def __len__(self): return self.size class SegfaultDataset(Dataset): def __init__(self, size): self.size = size def __getitem__(self, idx): return ctypes.string_at(0) def __len__(self): return self.size class SleepDataset(Dataset): def __init__(self, size, sleep_sec): self.size = size self.sleep_sec = sleep_sec self.sleeped = False def __getitem__(self, idx): if not self.sleeped: time.sleep(self.sleep_sec) self.sleeped = True return idx def __len__(self): return self.size class SeedDataset(Dataset): def __init__(self, size): self.size = size def __getitem__(self, idx): return ms.get_seed() def __len__(self): return self.size class WorkerSpecificIterableDataset(IterableDataset): def __init__(self, sizes_for_all_workers): self.sizes_for_all_workers = sizes_for_all_workers def __iter__(self): worker_info = torch.utils.data.get_worker_info() assert worker_info is not None return iter(range(self.sizes_for_all_workers[worker_info.id])) def __len__(self): return sum(self.sizes_for_all_workers) class SynchronizedDataset(Dataset): def __init__(self, size, batch_size, num_workers): assert size >= num_workers * batch_size self.count = mp.Value('i', 0, lock=True) self.barrier = mp.Semaphore(0) self.num_workers = num_workers self.size = size def sync_once(self): with self.count.get_lock(): self.count.value += 1 if self.count.value == self.num_workers: self.barrier.release() self.barrier.acquire() self.barrier.release() def __getitem__(self, idx): raise NotImplementedError def __len__(self): return self.size class EmptyTensorDataset(torch.utils.data.Dataset): def __init__(self, len): self.len = len def __len__(self): return self.len def __getitem__(self, any): return torch.empty(0) class SynchronizedSeedDataset(SynchronizedDataset): def __getitem__(self, idx): self.sync_once() return ms.get_seed() def _test_timeout(persistent_workers): dataset = SleepDataset(10, 3) dataloader = DataLoader(dataset, batch_size=2, num_workers=1, timeout=1, persistent_workers=persistent_workers) _ = next(iter(dataloader)) def _test_timeout_pin_memory(persistent_workers): dataset = SleepDataset(10, 3) dataloader = DataLoader(dataset, batch_size=2, num_workers=1, timeout=1, pin_memory=True, persistent_workers=persistent_workers) _ = next(iter(dataloader)) def _test_large_sampler_indices(persistent_workers): dataloader = torch.utils.data.DataLoader( EmptyTensorDataset(10000000), batch_size=40960, persistent_workers=persistent_workers, num_workers=1) it = iter(dataloader) for x in it: assert x.numel() == 0 raise RuntimeError('My Error') def disable_stderr(worker_id): r""" Avoids printing "ERROR: Unexpected segmentation fault encountered in worker." from workers. Since worker signal handler prints with low-level write(), this has to be done on OS level via dup. This is used as worker_init_fn for test_segfault. """ sys.stderr.flush() # flush library buffers that dup2 knows nothing about # Can't use a with-block because otherwise the fd will be closed when this # function ends. with open(os.devnull, 'w') as devnull: os.dup2(devnull.fileno(), sys.stderr.fileno()) def _test_segfault(): dataset = SegfaultDataset(10) dataloader = DataLoader(dataset, batch_size=2, num_workers=1, worker_init_fn=disable_stderr) _ = next(iter(dataloader)) def _test_no_segfault(): dataset = [1, 2, 3] num_threads = torch.get_num_threads() if num_threads < 4: torch.set_num_threads(4) else: torch.set_num_threads(num_threads) mp_ctx = torch.multiprocessing.get_context(method='fork') dataloader = DataLoader(dataset, num_workers=1, worker_init_fn=disable_stderr, multiprocessing_context=mp_ctx) _ = next(iter(dataloader)) class TestProperExitDataset(Dataset): def __init__(self, size, error_event): self.size = size self.error_event = error_event def __len__(self): return self.size def __getitem__(self, idx): worker_info = torch.utils.data.get_worker_info() if self.error_event is not None and self.error_event.is_set() and \ worker_info.id == worker_info.num_workers - 1: # only error in the last worker raise RuntimeError('Worker error') return torch.tensor([idx]) class TestProperExitIterableDataset(IterableDataset): def __init__(self, size, error_event): self.error_event = error_event self.size = size self.remaining = size def __len__(self): return self.size def __iter__(self): return self def __next__(self): worker_info = torch.utils.data.get_worker_info() if self.error_event is not None and self.error_event.is_set() and \ worker_info.id == worker_info.num_workers - 1: # only error in the last worker raise RuntimeError('Worker error') self.remaining -= 1 if self.remaining < 0: raise StopIteration return torch.tensor(-1000) # See TestDataLoader.test_proper_exit for usage def _test_proper_exit(is_iterable_dataset, use_workers, pin_memory, exit_method, hold_iter_reference, loader_setup_event, tester_setup_event, persistent_workers): num_workers = 2 if use_workers else 0 if exit_method == 'worker_error' or exit_method == 'worker_kill': assert use_workers is True if exit_method == 'worker_error': worker_error_event = mp.Event() else: worker_error_event = None if is_iterable_dataset: ds = TestProperExitIterableDataset(7, worker_error_event) else: ds = TestProperExitDataset(12, worker_error_event) loader = DataLoader(ds, batch_size=1, shuffle=False, num_workers=num_workers, pin_memory=pin_memory, worker_init_fn=set_faulthander_if_available, persistent_workers=persistent_workers) error_it = 2 if use_workers: # 2 is the magical per-worker prefetch number... # FIXME: change this after the number becomes configurable. if is_iterable_dataset: assert len(ds) * num_workers > (error_it + 2 + 1) else: assert len(loader) > (error_it + 2 + 1) * num_workers else: if is_iterable_dataset: assert len(ds) > error_it + 1 else: assert len(loader) > error_it + 1 it = iter(loader) if use_workers: workers = it._workers def kill_pid(pid): psutil_p = psutil.Process(pid) psutil_p.kill() psutil_p.wait(JOIN_TIMEOUT) assert not psutil_p.is_running() for i, _ in enumerate(it): if i == 0: if not hold_iter_reference: del it del loader loader_setup_event.set() tester_setup_event.wait() # ensure that the workers are still alive if use_workers: for w in workers: assert w.is_alive() if worker_error_event is not None: worker_error_event.set() if i == error_it: if exit_method == 'loader_error': raise RuntimeError('Loader error') elif exit_method == 'loader_kill': kill_pid(os.getpid()) elif exit_method == 'worker_kill': kill_pid(workers[-1].pid) # kill last worker if not hold_iter_reference: # Tries to trigger the __del__ clean-up rather than the automatic # exiting of daemonic children. Technically it should be automatically # triggered, but I don't want to rely on the implementation detail of # Python gc. gc.collect() class TestWorkerInfoDataset(SynchronizedDataset): def __getitem__(self, idx): self.sync_once() return torch.tensor(self.value) # Should be used as worker_init_fn with TestWorkerInfoDataset. # See _test_get_worker_info below for usage. def _test_worker_info_init_fn(worker_id): worker_info = torch.utils.data.get_worker_info() assert worker_id == worker_info.id, "worker_init_fn and worker_info should have consistent id" assert worker_id < worker_info.num_workers, "worker_init_fn and worker_info should have valid id" assert worker_info.seed == ms.get_seed(), "worker_init_fn and worker_info should have consistent seed" dataset = worker_info.dataset assert isinstance(dataset, TestWorkerInfoDataset), "worker_info should have correct dataset copy" assert not hasattr(dataset, 'value'), "worker_info should have correct dataset copy" # test that WorkerInfo attributes are read-only try: worker_info.id = 3999 except RuntimeError as e: assert str(e) == "Cannot assign attributes to WorkerInfo objects" try: worker_info.a = 3 except RuntimeError as e: assert str(e) == "Cannot assign attributes to WorkerInfo objects" for k in ['id', 'num_workers', 'seed', 'dataset']: assert "{}=".format(k) in repr(worker_info) dataset.value = [worker_id, os.getpid()] def _test_get_worker_info(): # get_worker_info returns None in main proc assert torch.utils.data.get_worker_info() is None num_workers = 2 batch_size = 2 dataset = TestWorkerInfoDataset(6, batch_size, num_workers) dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, worker_init_fn=_test_worker_info_init_fn) it = iter(dataloader) data = [] for d in it: data.append(d) worker_pids = [w.pid for w in it._workers] data = torch.cat(data, 0) for d in data: # each `d` is a [worker_id, worker_pid] pair, which is set in # _test_worker_info_init_fn assert d[1] == worker_pids[d[0]] # get_worker_info returns None in main proc after data loading assert torch.utils.data.get_worker_info() is None # main proc dataset was never assigned this attribute assert not hasattr(dataset, 'value') try: _ = dataset[0] except AttributeError: return raise RuntimeError('Expected AttributeError') # test custom init function def init_fn(worker_id): ms.set_seed(12345) # used with test_error_in_init class ErrorIterableDataset(IterableDataset): def __iter__(self): raise RuntimeError("Error in __iter__") # used with test_error_in_init def error_worker_init_fn(_): raise RuntimeError("Error in worker_init_fn") class BulkLoadingDataset(Dataset): def __init__(self, length): self.length = length def __getitem__(self, indices): assert isinstance(indices, (list, tuple)) return torch.tensor(indices) def __len__(self): return self.length class BulkLoadingSampler(torch.utils.data.Sampler): def __init__(self, dataset, batch_size): self.dataset = dataset self.batch_size = batch_size def __iter__(self): for x in torch.randperm(len(self.dataset)).split(self.batch_size): yield x.asnumpy().tolist() def __len__(self): return int(math.ceil(len(self.dataset) / float(self.batch_size))) class TestMultiEpochDataset(IterableDataset): def __init__(self, length): self.length = length def __iter__(self): worker_info = torch.utils.data.get_worker_info() assert worker_info is not None worker_id = worker_info.id for idx in range(self.length // worker_info.num_workers): yield worker_id def __len__(self): return self.length class CustomList(list): pass class CustomDict(dict): pass def row_processor(row): return np.add(row, 1) def filter_len(row): return len(row) == 4 data = torch.randn(100, 2, 3, 5) labels = torch.randperm(50).repeat(2) dataset = TensorDataset(data, labels) persistent_workers = False def _get_data_loader(dataset, **kwargs): persistent_workers = kwargs.get('persistent_workers', False) if persistent_workers and kwargs.get('num_workers', 0) == 0: persistent_workers = False kwargs['persistent_workers'] = persistent_workers return DataLoader(dataset, **kwargs) def _test_sequential(loader): batch_size = loader.batch_size if batch_size is None: for idx, (sample, target) in enumerate(loader): assert np.allclose(sample.asnumpy(), data[idx].asnumpy()) assert np.allclose(target.asnumpy(), labels[idx].asnumpy()) assert np.allclose(idx, len(dataset) - 1) else: for i, (sample, target) in enumerate(loader): idx = i * batch_size assert np.allclose(sample.asnumpy(), data[idx:idx + batch_size].asnumpy()) assert np.allclose(target.asnumpy(), labels[idx:idx + batch_size].asnumpy()) assert np.allclose(i, math.floor((len(dataset) - 1) / batch_size)) def _test_shuffle(loader): found_data = {i: 0 for i in range(data.shape[0])} found_labels = {i: 0 for i in range(labels.shape[0])} batch_size = loader.batch_size if batch_size is None: for i, (batch_samples, batch_targets) in enumerate(loader): sample, target = (batch_samples, batch_targets) for data_point_idx, data_point in enumerate(data): res = ms.ops.equal(data_point, sample) res = res.all() if res: found_data[data_point_idx] += 1 break assert np.allclose(target.asnumpy(), labels[data_point_idx].asnumpy()) found_labels[data_point_idx] += 1 assert np.allclose(sum(found_data.values()), (i + 1)) assert np.allclose(sum(found_labels.values()), (i + 1)) assert np.allclose(i, (len(dataset) - 1)) else: for i, (batch_samples, batch_targets) in enumerate(loader): for sample, target in zip(batch_samples, batch_targets): for data_point_idx, data_point in enumerate(data): res = ms.ops.equal(data_point, sample) res = res.all() if res: found_data[data_point_idx] += 1 break assert np.allclose(target.asnumpy(), labels[data_point_idx].asnumpy()) found_labels[data_point_idx] += 1 assert np.allclose(sum(found_data.values()), (i + 1) * batch_size) assert np.allclose(sum(found_labels.values()), (i + 1) * batch_size) assert np.allclose(i, math.floor((len(dataset) - 1) / batch_size)) def _test_error(loader): it = iter(loader) errors = 0 while True: try: next(it) except NotImplementedError: errors += 1 except StopIteration: assert np.allclose(errors, math.ceil(float(len(loader.dataset)) / loader.batch_size)) return def test_sequential_nonbatch(): _test_sequential(_get_data_loader(dataset, batch_size=None)) def test_sequential_batch(): _test_sequential(_get_data_loader(dataset)) _test_sequential(_get_data_loader(dataset, batch_size=2)) def test_growing_dataset(): dataset = [torch.ones(4) for _ in range(4)] dataloader_seq = _get_data_loader(dataset, shuffle=False) dataloader_shuffle = _get_data_loader(dataset, shuffle=True) dataset.append(torch.ones(4)) assert np.allclose(len(dataloader_seq), 5) assert np.allclose(len(dataloader_shuffle), 5) # def test_multiple_dataloaders(): # for multiprocessing_context in supported_multiprocessing_contexts: # loader1_it = iter(_get_data_loader(dataset, num_workers=0)) # loader2_it = iter(_get_data_loader(dataset, num_workers=0, multiprocessing_context=multiprocessing_context)) # next(loader1_it) # next(loader1_it) # next(loader2_it) # next(loader2_it) # next(loader1_it) # next(loader2_it) # del loader1_it # del loader2_it def test_builtin_collection_conversion(): for coll_ty in (list, tuple): # map-style dataset dataset = CountingDataset(20) # no auto-batching fetched = coll_ty(_get_data_loader(dataset, batch_size=None, num_workers=0)) assert np.allclose(fetched, coll_ty(range(20))) # auto-batching fetched = coll_ty(_get_data_loader(dataset, batch_size=2, num_workers=0)) sample = coll_ty(torch.tensor([i, i + 1]) for i in range(0, 20, 2)) for i in range(len(fetched)): assert np.allclose(fetched[i].asnumpy(), sample[i].asnumpy()) # iterable-style dataset dataset = CountingIterableDataset(20) # no auto-batching fetched = coll_ty(_get_data_loader(dataset, batch_size=None, num_workers=0)) assert np.allclose(fetched, coll_ty(range(20))) fetched = coll_ty(_get_data_loader(dataset, batch_size=2, num_workers=0)) sample = coll_ty(torch.tensor([i, i + 1]) for i in range(0, 20, 2)) for i in range(len(fetched)): assert np.allclose(fetched[i].asnumpy(), sample[i].asnumpy()) def test_iterable_style_dataset(): # [no auto-batching] single process loading dataset = CountingIterableDataset(20) dataloader = _get_data_loader(dataset, batch_size=None) fetched = list(dataloader) assert np.allclose(len(fetched), 20) for i, d in enumerate(fetched): # non-batched should not convert ints into tensors assert np.allclose(d, i) # DataLoader should match len of the iterable-style dataset (if implemented) assert np.allclose(len(dataloader), len(dataset)) # [no auto-batching] multiprocessing loading # num_workers = 3 # sizes_for_all_workers = [0, 4, 20] # expected = sorted(sum((list(range(s)) for s in sizes_for_all_workers), [])) # assert len(sizes_for_all_workers) == num_workers, 'invalid test case' # for prefetch_factor in [2, 3, 4]: # dataset = WorkerSpecificIterableDataset(sizes_for_all_workers) # dataloader = _get_data_loader(dataset, num_workers=num_workers, batch_size=None, # worker_init_fn=set_faulthander_if_available, # prefetch_factor=prefetch_factor) # dataloader_iter = iter(dataloader) # fetched = sorted(dataloader_iter) # for a, b in zip(fetched, expected): # # non-batched should not convert ints into tensors # assert np.allclose(a, b) # # DataLoader should match len of the iterable-style dataset (if implemented) # assert np.allclose(len(dataloader), len(dataset)) # # When loading more than len(dataset) data, after accessing len(dataloader), # # we should get a warning. See NOTE [ IterableDataset and __len__ ]. # dataset = CountingIterableDataset(20) # dataloader = _get_data_loader(dataset, num_workers=num_workers, # worker_init_fn=set_faulthander_if_available, # prefetch_factor=prefetch_factor) # it = iter(dataloader) # assert np.allclose(len(dataloader), len(dataset)) # assert np.allclose(len(dataloader), 20) # it = iter(dataloader) # # [no auto-batching] test that workers exit gracefully # workers = dataloader_iter._workers # del dataloader_iter # del dataloader # try: # for w in workers: # w.join(JOIN_TIMEOUT) # assert np.allclose(w.exitcode, 0) # finally: # for w in workers: # w.terminate() # # # [auto-batching] single process loading # dataset = CountingIterableDataset(20) # fetched = list(_get_data_loader(dataset, batch_size=7)) # assert np.allclose(len(fetched), 3) # assert np.allclose(fetched[0].asnumpy().tolist(), list(range(7))) # assert np.allclose(fetched[1].asnumpy().tolist(), list(range(7, 14))) # assert np.allclose(fetched[2].asnumpy().tolist(), list(range(14, 20))) # # # [auto-batching] multiprocessing loading # num_workers = 3 # sizes_for_all_workers = [0, 4, 20] # expected = sorted(sum((list(range(s)) for s in sizes_for_all_workers), [])) # assert len(sizes_for_all_workers) == num_workers, 'invalid test case' # for prefetch_factor in [2, 3, 4]: # dataset = WorkerSpecificIterableDataset(sizes_for_all_workers) # # worker 0 should return 0 batches # # worker 1 should return 1 batches # # worker 2 should return 3 batches # dataloader = _get_data_loader(dataset, num_workers=num_workers, batch_size=7, prefetch_factor=prefetch_factor) # dataloader_iter = iter(dataloader) # fetched = list(dataloader_iter) # assert np.allclose(len(fetched), 4) # fetched = set(tuple(t.tolist()) for t in fetched) # assert np.allclose(fetched, tuple(range(4))) # assert np.allclose(fetched, tuple(range(7))) # assert np.allclose(fetched, tuple(range(7, 14))) # assert np.allclose(fetched, tuple(range(14, 20))) # # # [auto-batching] test that workers exit gracefully # workers = dataloader_iter._workers # del dataloader_iter # del dataloader # try: # for w in workers: # w.join(JOIN_TIMEOUT) # assert np.allclose(w.exitcode, 0) # finally: # for w in workers: # w.terminate() # # [auto-batching & drop_last] single process loading # dataset = CountingIterableDataset(20) # fetched = list(_get_data_loader(dataset, batch_size=7, drop_last=True)) # assert np.allclose(len(fetched), 2) # assert np.allclose(fetched[0].tolist(), list(range(7))) # assert np.allclose(fetched[1].tolist(), list(range(7, 14))) # # # [auto-batching & drop_last] multiprocessing loading # num_workers = 3 # sizes_for_all_workers = [0, 4, 20] # expected = sorted(sum((list(range(s)) for s in sizes_for_all_workers), [])) # assert len(sizes_for_all_workers) == num_workers, 'invalid test case' # for prefetch_factor in [2, 3, 4]: # dataset = WorkerSpecificIterableDataset(sizes_for_all_workers) # # worker 0 should return 0 batches # # worker 1 should return 1 batches # # worker 2 should return 3 batches # dataloader = _get_data_loader(dataset, num_workers=num_workers, batch_size=7, drop_last=True, # worker_init_fn=set_faulthander_if_available, # prefetch_factor=prefetch_factor) # dataloader_iter = iter(dataloader) # fetched = list(dataloader_iter) # assert np.allclose(len(fetched), 2) # fetched = set(tuple(t.tolist()) for t in fetched) # assert np.allclose(fetched, {tuple(range(7)), tuple(range(7, 14))}) # # # [auto-batching & drop_last] test that workers exit gracefully # workers = dataloader_iter._workers # del dataloader_iter # del dataloader # try: # for w in workers: # w.join(JOIN_TIMEOUT) # assert np.allclose(w.exitcode, 0) # finally: # for w in workers: # w.terminate() def test_chain_iterable_style_dataset(): dataset1 = CountingIterableDataset(20) dataset2 = CountingIterableDataset(15) expected = list(range(20)) + list(range(15)) for num_workers in [0]: #TODO test num_workers > 0 for chained_dataset in [dataset1 + dataset2, ChainDataset([dataset1, dataset2])]: fetched = list(_get_data_loader(chained_dataset, num_workers=num_workers)) assert np.allclose(len(fetched), len(expected)) for e, d in zip(expected, fetched): assert np.allclose(e, d.asnumpy()) # def test_multiprocessing_contexts(): # reference = [ # torch.arange(0, 3), # torch.arange(3, 6), # torch.arange(6, 9), # torch.arange(9, 11), # ] # counting_ds_n = 11 # dl_common_args = dict(num_workers=1, batch_size=3) # for ctx in supported_multiprocessing_contexts: # ds_cls = CountingDataset # sample = list(_get_data_loader(ds_cls(counting_ds_n), multiprocessing_context=ctx, **dl_common_args)) # for i in range(4): # assert np.allclose(reference[i].asnumpy(), sample[i].asnumpy()) # if ctx is not None: # # test ctx object # ctx = mp.get_context(ctx) # sample = list(_get_data_loader(ds_cls(counting_ds_n), multiprocessing_context=ctx, **dl_common_args)) # for i in range(4): # assert np.allclose(reference[i].asnumpy(), sample[i].asnumpy()) # def test_worker_seed(): # num_workers = 6 # batch_size = 1 # dataset = SynchronizedSeedDataset(num_workers, batch_size, num_workers) # dataloader = _get_data_loader(dataset, batch_size=batch_size, num_workers=num_workers) # seeds = set() # for batch in dataloader: # seeds.add(batch[0]) # assert np.allclose(len(seeds), num_workers) # def test_worker_seed_reproducibility(): np.random.seed(42) ms.set_seed(42) def get_dataloader(): return DataLoader(dataset, batch_size=batch_size, num_workers=num_workers) num_workers = 0 # TODO test num_workers > 0 batch_size = 1 dataset = SynchronizedSeedDataset(num_workers, batch_size, num_workers) first = list(int(batch) for batch in get_dataloader()) second = list(int(batch) for batch in get_dataloader()) for i in range(len(first)): assert np.allclose(first[i], second[i]) # def test_multi_epochs_reproducibility(): # num_workers = 3 # TODO test num_workers > 0 # batch_size = 10 # num_epochs = 3 # # dataset = TestMultiEpochDataset(batch_size * num_workers) # dataloader = _get_data_loader(dataset, batch_size=batch_size, # shuffle=False, num_workers=num_workers) # # for ind in range(num_epochs): # for batch_idx, sample in enumerate(dataloader): # np.allclose(sample.asnumpy().tolist(), [batch_idx % num_workers] * batch_size) def test_get_worker_info(): p = ErrorTrackingProcess(target=_test_get_worker_info) p.start() p.join(JOIN_TIMEOUT) try: np.allclose(p.exitcode, 0) finally: p.terminate() def test_shuffle(): _test_shuffle(_get_data_loader(dataset, shuffle=True)) def test_shuffle_batch_none(): _test_shuffle(DataLoader(dataset, batch_size=None, shuffle=True)) def test_shuffle_batch(): _test_shuffle(_get_data_loader(dataset, batch_size=2, shuffle=True)) def test_shuffle_reproducibility(): np.random.seed(42) #TODO test num_workers > 0 ms.set_seed(42) first = list(DataLoader(dataset, shuffle=True, num_workers=0)) np.random.seed(42) second = list(DataLoader(dataset, shuffle=True, num_workers=0)) length = len(first) for i in range(length): for j in range(len(first[i])): assert np.allclose(first[i][j].asnumpy(), second[i][j].asnumpy()) # def test_sequential_workers(): # _test_sequential(_get_data_loader(dataset, num_workers=4)) # # def test_seqential_batch_workers(): # _test_sequential(_get_data_loader(dataset, batch_size=2, num_workers=4)) # # def test_seqential_batch_workers_prefetch(): # _test_sequential(DataLoader(dataset, batch_size=2, num_workers=4, prefetch_factor=3)) # # def test_shuffle_workers(): # _test_shuffle(_get_data_loader(dataset, shuffle=True, num_workers=4)) # # def test_shuffle_batch_workers(): # _test_shuffle(_get_data_loader(dataset, batch_size=2, shuffle=True, num_workers=4)) # # def test_shuffle_batch_workers_prefetch(): # _test_shuffle(DataLoader(dataset, batch_size=2, shuffle=True, num_workers=4, prefetch_factor=3)) def test_random_sampler(): from collections import Counter from ms_adapter.pytorch.utils.data import RandomSampler def sample_stat(sampler, num_samples): counts = Counter(sampler) count_repeated = sum(val > 1 for val in counts.values()) return (count_repeated, min(counts.keys()), max(counts.keys()), sum(counts.values())) # test sample with replacement n = len(dataset) + 1 # ensure at least one sample is drawn more than once sampler_with_replacement = RandomSampler(dataset, replacement=True, num_samples=n) count_repeated, minval, maxval, count_total = sample_stat(sampler_with_replacement, n) assert np.allclose(count_total, n) # test sample without replacement and without specified num_samples sampler_without_replacement = RandomSampler(dataset) count_repeated, minval, maxval, count_total = sample_stat(sampler_without_replacement, len(dataset)) assert np.allclose(count_repeated, 0) assert np.allclose(minval , 0) assert np.allclose(maxval , len(dataset) - 1) assert np.allclose(count_total , len(dataset)) # test sample without replacement and with specified num_samples n = len(dataset) * 2 sampler_without_replacement = RandomSampler(dataset, num_samples=n) count_repeated, minval, maxval, count_total = sample_stat(sampler_without_replacement, len(dataset)) assert np.allclose(count_repeated, len(dataset)) assert np.allclose(minval, 0) assert np.allclose(maxval, len(dataset) - 1) assert np.allclose(count_total, n) n = len(dataset) - 1 sampler_without_replacement = RandomSampler(dataset, num_samples=n) count_repeated, minval, maxval, count_total = sample_stat(sampler_without_replacement, len(dataset)) assert np.allclose(count_total, n) n = len(dataset) + 1 sampler_without_replacement = RandomSampler(dataset, num_samples=n) count_repeated, minval, maxval, count_total = sample_stat(sampler_without_replacement, len(dataset)) assert np.allclose(count_repeated , 1) assert np.allclose(minval,0) assert np.allclose(maxval, len(dataset) - 1) assert np.allclose(count_total, n) def test_random_sampler_len_with_replacement(): from ms_adapter.pytorch.utils.data import RandomSampler # add 5 extra samples num_samples = len(dataset) + 5 sampler = RandomSampler(dataset, replacement=True, num_samples=num_samples) # test len method assert np.allclose(num_samples, len(sampler)) # test with iteration count_num_samples = sum(1 for _ in sampler) assert np.allclose(num_samples, count_num_samples) # test with dataloader, batch_size = 1 batch_size = 1 count_num_samples_in_data_loader = len(_get_data_loader( dataset, batch_size=batch_size, sampler=sampler)) assert np.allclose(num_samples, count_num_samples_in_data_loader) # test with dataloader, batch_size = 6 batch_size = 6 count_num_samples_in_data_loader = len(_get_data_loader( dataset, batch_size=batch_size, sampler=sampler)) assert np.allclose(int(math.ceil(float(num_samples) / batch_size)), count_num_samples_in_data_loader) def test_random_sampler_len_without_replacement(): from ms_adapter.pytorch.utils.data import RandomSampler # add 5 extra samples num_samples = len(dataset) + 5 sampler = RandomSampler(dataset, replacement=False, num_samples=num_samples) # test len method assert np.allclose(num_samples, len(sampler)) # test with iteration count_num_samples = sum(1 for _ in sampler) assert np.allclose(num_samples, count_num_samples) # test with dataloader, batch_size = 1 batch_size = 1 count_num_samples_in_data_loader = len(_get_data_loader( dataset, batch_size=batch_size, sampler=sampler)) assert np.allclose(num_samples, count_num_samples_in_data_loader) # test with dataloader, batch_size = 6 batch_size = 6 count_num_samples_in_data_loader = len(_get_data_loader( dataset, batch_size=batch_size, sampler=sampler)) assert np.allclose(num_samples // batch_size + (num_samples % batch_size > 0), count_num_samples_in_data_loader) def test_duplicating_data_with_drop_last(): # TODO: Ascend not support concat empty tensor if ms.get_context('device_target') == 'Ascend': warnings.warn("test_dataloder.py: test_duplicating_data_with_drop_last not run on Ascend.") return from ms_adapter.pytorch.utils.data.distributed import DistributedSampler num_processes = 4 num_batches = 9 data_set = torch.Tensor(list(range(num_batches))).astype(torch.int64) scanned_data = torch.Tensor([]).astype(torch.int64) for i in range(num_processes): s = DistributedSampler(data_set, num_processes, i) d_loader = _get_data_loader(data_set, batch_size=int(num_batches / num_processes), drop_last=True, sampler=s) for data in d_loader: scanned_data = torch.cat((scanned_data, data), 0) unique_data, _ = ms.ops.unique(scanned_data) assert np.allclose(scanned_data.shape, unique_data.shape) def test_sampler_reproducibility(): from ms_adapter.pytorch.utils.data import RandomSampler, WeightedRandomSampler, SubsetRandomSampler weights = [0.1, 0.9, 0.4, 0.7, 3.0, 0.6] np.random.seed(42) ms.set_seed(42) first = list(RandomSampler(dataset, num_samples=5, replacement=True)) np.random.seed(42) ms.set_seed(42) second = list(RandomSampler(dataset, num_samples=5, replacement=True)) assert np.allclose(first, second) np.random.seed(42) ms.set_seed(42) first =list(RandomSampler(dataset, replacement=False)) np.random.seed(42) ms.set_seed(42) second =list(RandomSampler(dataset, replacement=False)) assert np.allclose(first, second) np.random.seed(42) ms.set_seed(42) first = list(WeightedRandomSampler(weights, num_samples=5, replacement=True)) np.random.seed(42) ms.set_seed(42) second = list(WeightedRandomSampler(weights, num_samples=5, replacement=True)) np.allclose(first, second) np.random.seed(42) ms.set_seed(42) first = list(WeightedRandomSampler(weights, num_samples=5, replacement=False)) np.random.seed(42) ms.set_seed(42) second = list(WeightedRandomSampler(weights, num_samples=5, replacement=False)) np.allclose(first, second) np.random.seed(42) ms.set_seed(42) first = list(SubsetRandomSampler(range(10))) np.random.seed(42) ms.set_seed(42) second = list(SubsetRandomSampler(range(10))) assert np.allclose(first, second) def _test_sampler(**kwargs): indices = range(2, 12) # using a regular iterable dl = _get_data_loader(dataset, sampler=indices, batch_size=2, **kwargs) assert np.allclose(len(dl), 5) for i, (input, _target) in enumerate(dl): assert np.allclose(len(input), 2) assert np.allclose(input.asnumpy(), data[i * 2 + 2:i * 2 + 4].asnumpy()) # def test_sampler(): # _test_sampler() # _test_sampler(num_workers=4) # if not NO_MULTIPROCESSING_SPAWN: # _test_batch_sampler(num_workers=4, multiprocessing_context='spawn') def _test_batch_sampler( **kwargs): # [(0, 1), (2, 3, 4), (5, 6), (7, 8, 9), ...] batches = [] # using a regular iterable for i in range(0, 20, 5): batches.append(tuple(range(i, i + 2))) batches.append(tuple(range(i + 2, i + 5))) dl = _get_data_loader(dataset, batch_sampler=batches, **kwargs) assert np.allclose(len(dl), 8) for i, (input, _target) in enumerate(dl): if i % 2 == 0: offset = i * 5 // 2 assert np.allclose(len(input), 2) assert np.allclose(input.asnumpy(), data[offset:offset + 2].asnumpy()) else: offset = i * 5 // 2 assert np.allclose(len(input), 3) assert np.allclose(input.asnumpy(), data[offset:offset + 3].asnumpy()) # def test_batch_sampler(): # _test_batch_sampler() # _test_batch_sampler(num_workers=4) # if not NO_MULTIPROCESSING_SPAWN: # _test_batch_sampler(num_workers=4, multiprocessing_context='spawn') def test_error(): _test_error(_get_data_loader(ErrorDataset(100), batch_size=2, shuffle=True)) def test_error_workers(): _test_error(_get_data_loader(ErrorDataset(41), batch_size=2, shuffle=True, num_workers=4)) def test_len_1(): def check_len(dl, expected): assert np.allclose(len(dl), expected) n = 0 for _ in dl: n += 1 assert np.allclose(n, expected) check_len(dataset, 100) check_len(_get_data_loader(dataset, batch_size=2), 50) check_len(_get_data_loader(dataset, batch_size=3), 34) def test_iterabledataset_len(): class IterableDataset(torch.utils.data.IterableDataset): def __len__(self): return 10 def __iter__(self): return iter(range(10)) iterable_loader = DataLoader(IterableDataset(), batch_size=1) assert np.allclose(len(iterable_loader), 10) iterable_loader = DataLoader(IterableDataset(), batch_size=1, drop_last=True) assert np.allclose(len(iterable_loader), 10) iterable_loader = DataLoader(IterableDataset(), batch_size=2) assert np.allclose(len(iterable_loader), 5) iterable_loader = DataLoader(IterableDataset(), batch_size=2, drop_last=True) assert np.allclose(len(iterable_loader), 5) iterable_loader = DataLoader(IterableDataset(), batch_size=3) assert np.allclose(len(iterable_loader), 4) iterable_loader = DataLoader(IterableDataset(), batch_size=3, drop_last=True) assert np.allclose(len(iterable_loader), 3) def test_default_convert_mapping_keep_type(): data = CustomDict({"a": 1, "b": 2}) converted = _utils.collate.default_convert(data) assert converted == data def test_default_convert_sequence_keep_type(): data = CustomList([1, 2, 3]) converted = _utils.collate.default_convert(data) assert np.allclose(converted, data) def test_default_convert_sequence_dont_keep_type(): data = range(2) converted = _utils.collate.default_convert(data) assert np.allclose(converted, [0, 1]) def test_default_collate_dtype(): arr = [1, 2, -1] collated = _utils.collate.default_collate(arr) assert np.allclose(collated.asnumpy(), torch.tensor(arr).asnumpy()) assert collated.dtype == torch.int64 arr = [1.1, 2.3, -0.9] collated = _utils.collate.default_collate(arr) assert np.allclose(collated.asnumpy(), torch.tensor(arr).asnumpy()) assert collated.dtype == torch.float64 arr = [True, False] collated = _utils.collate.default_collate(arr) assert np.allclose(collated.asnumpy(), torch.tensor(arr).asnumpy()) assert collated.dtype == torch.bool_ # Should be a no-op arr = ['a', 'b', 'c'] assert arr == _utils.collate.default_collate(arr) def test_default_collate_mapping_keep_type(): batch = [CustomDict({"a": 1, "b": 2}), CustomDict({"a": 3, "b": 4})] collated = _utils.collate.default_collate(batch) expected = CustomDict({"a": torch.tensor([1, 3]), "b": torch.tensor([2, 4])}) def test_default_collate_sequence_keep_type(): batch = [CustomList([1, 2, 3]), CustomList([4, 5, 6])] collated = _utils.collate.default_collate(batch) expected = CustomList([ torch.tensor([1, 4]), torch.tensor([2, 5]), torch.tensor([3, 6]), ]) for i in range(len(expected)): assert np.allclose(collated[i].asnumpy(), expected[i].asnumpy()) msadapter/testing/ut/pytorch/data/test_sampler.py000066400000000000000000000055331436713554100226370ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch.utils.data.sampler import ( BatchSampler, RandomSampler, Sampler, SequentialSampler, SubsetRandomSampler, WeightedRandomSampler, ) def test_BatchSampler(): index_1 = list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False)) index_2 = list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True)) len_1 = len(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False)) len_2 = len(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True)) print(index_1) print(index_2) print(len_1) print(len_2) def test_SequentialSampler(): index_1 = list(SequentialSampler(range(10))) len_1 = len(SequentialSampler(range(10))) print(index_1) print(len_1) def test_WeightedRandomSampler(): index_1 = list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True)) index_2 = list(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False)) len_1 = len(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True)) len_2 = len(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False)) print(index_1) print(index_2) print(len_1) print(len_2) def test_RandomSampler(): index_1 = list(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=True)) len_1 = len(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=True)) print(index_1) print(len_1) index_2 = list(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=True, num_samples=3)) len_2 = len(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=True, num_samples=3)) print(index_2) print(len_2) index_3 = list(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=False)) len_3 = len(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=False)) print(index_3) print(len_3) index_4 = list(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=False, num_samples=3)) len_4 = len(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=False, num_samples=3)) print(index_4) print(len_4) index_5 = list(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=True, num_samples=9)) len_5 = len(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=True, num_samples=9)) print(index_5) print(len_5) index_6 = list(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=False, num_samples=9)) len_6 = len(RandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], replacement=False, num_samples=9)) print(index_6) print(len_6) def test_SubsetRandomSampler(): index_1 = list(SubsetRandomSampler([1, 9, 4, 7, 3, 6])) len_1 = len(SubsetRandomSampler([1, 9, 4, 7, 3, 6])) print(index_1) print(len_1)msadapter/testing/ut/pytorch/functional/000077500000000000000000000000001436713554100210065ustar00rootroot00000000000000msadapter/testing/ut/pytorch/functional/test_activation.py000066400000000000000000000104221436713554100245570ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms from mindspore import Tensor import torch.nn.functional as t_fun import torch import numpy as np from mindspore import context import ms_adapter.pytorch as ms_torch import ms_adapter.pytorch.nn.functional as msa_fun context.set_context(mode=ms.PYNATIVE_MODE) def test_rrelu(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.rrelu(torch_input, inplace=True) ms_input = Tensor(data) ms_output = msa_fun.rrelu(ms_input, inplace=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_selu(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.selu(torch_input) ms_input = Tensor(data) ms_output = msa_fun.selu(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_celu(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.celu(torch_input) ms_input = Tensor(data) ms_output = msa_fun.celu(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_gelu(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.gelu(torch_input) ms_input = Tensor(data) ms_output = msa_fun.gelu(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) def test_mish(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.mish(torch_input) ms_input = Tensor(data) ms_output = msa_fun.mish(ms_input) if context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_softshrink(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.softshrink(torch_input) ms_input = Tensor(data) ms_output = msa_fun.softshrink(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_relu(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.relu(torch_input) ms_input = Tensor(data) ms_output = msa_fun.relu(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_hardtanh(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.hardtanh(torch_input) ms_input = Tensor(data) ms_output = msa_fun.hardtanh(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_hardswish(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.hardswish(torch_input) ms_input = Tensor(data) ms_output = msa_fun.hardswish(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_relu6(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.relu6(torch_input) ms_input = Tensor(data) ms_output = msa_fun.relu6(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_leaky_relu(): data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = t_fun.leaky_relu(torch_input) ms_input = Tensor(data) ms_output = msa_fun.leaky_relu(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_prelu(): input = np.random.rand(1, 3, 32, 32).astype(np.float32) weight = np.random.rand(1).astype(np.float32) torch_input = torch.tensor(input) torch_weight = torch.tensor(weight) torch_output = t_fun.prelu(torch_input, torch_weight) ms_input = ms_torch.tensor(input) ms_weight = ms_torch.tensor(weight) ms_output = msa_fun.prelu(ms_input, ms_weight) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) if __name__ == '__main__': test_rrelu() test_selu() test_celu() test_gelu() test_mish() test_softshrink() test_relu() test_hardtanh() test_prelu() msadapter/testing/ut/pytorch/functional/test_arange.py000066400000000000000000000012651436713554100236600ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_arange1(): ms_result = ms_torch.arange(1, 10) torch_result = torch.arange(1, 10) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_arange2(): ms_result = ms_torch.arange(1.0, 3, 0.5) torch_result = torch.arange(1.0, 3, 0.5) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_arange1() test_arange2() msadapter/testing/ut/pytorch/functional/test_blas_and_lapack.py000066400000000000000000000047131436713554100255020ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) #context.set_context(mode=ms.GRAPH_MODE) def test_matmul1(): np_1 = np.array([1, 2, 3, 4]).astype(np.int64) np_2 = np.array([5, 6, 7, 8]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.matmul(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.matmul(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_matmul2(): np_1 = np.array([1, 2]).astype(np.float32) np_2 = np.array([[5, 6], [7, 8]]).astype(np.float32) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.matmul(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.matmul(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_matmul3(): np_1 = np.arange(0, 36).reshape(3, 1, 6, 2).astype(np.int32) np_2 = np.arange(0, 40).reshape(5, 2, 4).astype(np.int32) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.matmul(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.matmul(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_matmul4(): np_1 = np.arange(0, 12).reshape(3, 1, 2, 2).astype(np.int32) np_2 = np.arange(0, 20).reshape(5, 2, 2).astype(np.int32) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.matmul(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.matmul(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_matmul1() test_matmul2() test_matmul3() test_matmul4() msadapter/testing/ut/pytorch/functional/test_cat.py000066400000000000000000000043331436713554100231710ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_cat1(): ms_tensor = ms_torch.tensor([1, 2, 3]) ms_result = ms_torch.cat((ms_tensor, ms_tensor), dim=0) torch_tensor = torch.tensor([1, 2, 3]) torch_result = torch.cat((torch_tensor, torch_tensor), dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_cat2(): ms_tensor = ms_torch.tensor([[1, 2, 3], [1, 2, 3]]) ms_result = ms_torch.cat((ms_tensor, ms_tensor), dim=1) torch_tensor = torch.tensor([[1, 2, 3], [1, 2, 3]]) torch_result = torch.cat((torch_tensor, torch_tensor), dim=1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_cat3(): ms_tensor = ms_torch.tensor([[1, 2, 3], [1, 2, 3]]) ms_result = ms_torch.cat([ms_tensor, ms_tensor], dim=0) torch_tensor = torch.tensor([[1, 2, 3], [1, 2, 3]]) torch_result = torch.cat([torch_tensor, torch_tensor], dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_concat1(): ms_tensor = ms_torch.tensor([1, 2, 3]) ms_result = ms_torch.concat((ms_tensor, ms_tensor), dim=0) torch_tensor = torch.tensor([1, 2, 3]) torch_result = torch.concat((torch_tensor, torch_tensor), dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_concat2(): ms_tensor = ms_torch.tensor([[1, 2, 3], [1, 2, 3]]) ms_result = ms_torch.concat((ms_tensor, ms_tensor), dim=1) torch_tensor = torch.tensor([[1, 2, 3], [1, 2, 3]]) torch_result = torch.concat((torch_tensor, torch_tensor), dim=1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_concat3(): ms_tensor = ms_torch.tensor([[1, 2, 3], [1, 2, 3]]) ms_result = ms_torch.concat([ms_tensor, ms_tensor], dim=0) torch_tensor = torch.tensor([[1, 2, 3], [1, 2, 3]]) torch_result = torch.concat([torch_tensor, torch_tensor], dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) if __name__ == '__main__': test_cat1() test_cat2() test_cat3() test_concat1() test_concat2() test_concat3() msadapter/testing/ut/pytorch/functional/test_chunk.py000066400000000000000000000012151436713554100235260ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_chunk1(): np_1 = np.arange(12) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.chunk(ms_tensor, 6) torch_tensor = torch.tensor(np_1) torch_result = torch.chunk(torch_tensor, 6) for i in range(6): assert np.allclose(ms_result[i].asnumpy(), torch_result[i].numpy()) assert ms_result[i].asnumpy().dtype == torch_result[i].numpy().dtype if __name__ == '__main__': test_chunk1() msadapter/testing/ut/pytorch/functional/test_compare.py000066400000000000000000000216661436713554100240600ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) #context.set_context(mode=ms.GRAPH_MODE) def test_ge1(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[4, 2]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.ge(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.ge(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_ge2(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.half) ms_tensor_1 = ms_torch.tensor(np_1) ms_out = ms_torch.ge(ms_tensor_1, 3.0) torch_tensor_1 = torch.tensor(np_1) torch_out = torch.ge(torch_tensor_1, 3.0) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_gt1(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[4, 2]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.gt(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.gt(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_gt2(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.half) ms_tensor_1 = ms_torch.tensor(np_1) ms_out = ms_torch.gt(ms_tensor_1, 3.0) torch_tensor_1 = torch.tensor(np_1) torch_out = torch.gt(torch_tensor_1, 3.0) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_le1(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[4, 2]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.le(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.le(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_le2(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.half) ms_tensor_1 = ms_torch.tensor(np_1) ms_out = ms_torch.le(ms_tensor_1, 3.0) torch_tensor_1 = torch.tensor(np_1) torch_out = torch.le(torch_tensor_1, 3.0) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_lt1(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[4, 2]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_torch.lt(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch.lt(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_lt2(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.half) ms_tensor_1 = ms_torch.tensor(np_1) ms_out = ms_torch.lt(ms_tensor_1, 3.0) torch_tensor_1 = torch.tensor(np_1) torch_out = torch.lt(torch_tensor_1, 3.0) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_maximum_minimum(): np_array = np.random.randn(1, 2, 3, 4).astype(np.half) np_other = np.random.randn(1, 2, 3, 4).astype(np.double) torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out1 = torch.maximum(torch_tensor, torch_other) torch_out2 = torch.minimum(torch_tensor, torch_other) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out1 = ms_torch.maximum(ms_tensor, ms_other) ms_out2 = ms_torch.minimum(ms_tensor, ms_other) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_eq(): np_array = np.random.rand(1, 2, 3, 4) * 5 np_other = np.random.rand(1, 2, 3, 4) * 5 np_array = np_array.astype(np.int32) np_other = np_other.astype(np.int32) torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out = torch.eq(torch_tensor, torch_other) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out = ms_torch.eq(ms_tensor, ms_other) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_ne(): x = np.random.randn(2, 3, 4) y = np.random.randn(3, 4) ms_x = ms_torch.tensor(x) ms_y = ms_torch.tensor(y) ms_out = ms_torch.ne(ms_x, ms_y) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch.ne(torch_x, torch_y) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_equal(): torch_out = torch.equal(torch.tensor([1, 2]), torch.tensor([1, 2])) ms_out = ms_torch.equal(ms_torch.tensor([1, 2]), ms_torch.tensor([1, 2])) assert ms_out == torch_out assert type(ms_out) == type(torch_out) torch_out = torch.equal(torch.tensor([1., 1.]), torch.tensor([1.])) ms_out = ms_torch.equal(ms_torch.tensor([1., 1.]), ms_torch.tensor([1.])) assert ms_out == torch_out assert type(ms_out) == type(torch_out) def test_greater_equal(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.float16) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.float16) ms_out = ms_torch.ge(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.float16) torch_tensor_2 = torch.tensor(np_2, dtype=torch.float16) torch_out = torch.ge(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_greater(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.float64) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.float64) ms_out = ms_torch.ge(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.float64) torch_tensor_2 = torch.tensor(np_2, dtype=torch.float64) torch_out = torch.ge(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_less_equal(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.double) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.double) ms_out = ms_torch.less_equal(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.double) torch_tensor_2 = torch.tensor(np_2, dtype=torch.double) torch_out = torch.less_equal(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_less(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.long) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.long) ms_out = ms_torch.ge(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.long) torch_tensor_2 = torch.tensor(np_2, dtype=torch.long) torch_out = torch.ge(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_not_equal(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.int32) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.int32) ms_out = ms_torch.ge(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.int32) torch_tensor_2 = torch.tensor(np_2, dtype=torch.int32) torch_out = torch.ge(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_ge1() test_ge2() test_gt1() test_gt2() test_le1() test_le2() test_lt1() test_lt2() test_maximum_minimum() test_eq() test_ne() test_equal() test_greater_equal() test_greater() test_less_equal() test_less() test_not_equal() msadapter/testing/ut/pytorch/functional/test_corrcoef.py000066400000000000000000000034711436713554100242260ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_corrcoef1(): np_1 = np.random.randn(3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.corrcoef(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.corrcoef(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_corrcoef2(): np_1 = np.random.randn(8).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.corrcoef(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.corrcoef(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_corrcoef3(): np_1 = np.random.randint(0, 10, 5).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.corrcoef(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.corrcoef(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_corrcoef4(): np_1 = np.random.randint(0, 10, (3, 5)).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.corrcoef(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.corrcoef(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == "__main__": test_corrcoef1() test_corrcoef2() test_corrcoef3() test_corrcoef4() msadapter/testing/ut/pytorch/functional/test_cov.py000066400000000000000000000046001436713554100232060ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_cov1(): np_1 = np.random.randn(8).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cov(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.cov(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cov2(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cov(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.cov(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cov3(): np_1 = np.random.randint(0, 10, 5).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cov(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.cov(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cov4(): np_1 = np.random.randint(0, 10, (3, 4)).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cov(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.cov(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cov5(): np_1 = np.random.randn(3, 4).astype(np.float32) fw = np.random.randint(1, 10, (4,)) aw = np.random.rand(4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_fw = ms_torch.tensor(fw) ms_aw = ms_torch.tensor(aw) ms_result = ms_torch.cov(ms_tensor, fweights=ms_fw, aweights=ms_aw) torch_tensor = torch.tensor(np_1) torch_fw = torch.tensor(fw) torch_aw = torch.tensor(aw) torch_result = torch.cov(torch_tensor, fweights=torch_fw, aweights=torch_aw) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == "__main__": test_cov1() test_cov2() test_cov3() test_cov4() test_cov5() msadapter/testing/ut/pytorch/functional/test_cross.py000066400000000000000000000025071436713554100235540ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_cross1(): np_1 = np.random.randn(4, 3).astype(np.float32) np_2 = np.random.randn(4, 3).astype(np.float32) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.cross(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.cross(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cross2(): np_1 = np.random.randn(4, 3).astype(np.int32) np_2 = np.random.randn(4, 3).astype(np.int32) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.cross(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.cross(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == "__main__": test_cross1() test_cross2() msadapter/testing/ut/pytorch/functional/test_cummax.py000066400000000000000000000050331436713554100237120ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_cummax1(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummax(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummax(torch_tensor, dim=0) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cummax2(): np_1 = np.random.randn(3, 4).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummax(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummax(torch_tensor, dim=0) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cummax3(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummax(ms_tensor, dim=1) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummax(torch_tensor, dim=1) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cummax4(): np_1 = np.random.randn(3, 4).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummax(ms_tensor, dim=1) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummax(torch_tensor, dim=1) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype if __name__ == "__main__": test_cummax1() test_cummax2() test_cummax3() test_cummax4() msadapter/testing/ut/pytorch/functional/test_cummin.py000066400000000000000000000050331436713554100237100ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_cummin1(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummin(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummin(torch_tensor, dim=0) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cummin2(): np_1 = np.random.randn(3, 4).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummin(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummin(torch_tensor, dim=0) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cummin3(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummin(ms_tensor, dim=1) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummin(torch_tensor, dim=1) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cummin4(): np_1 = np.random.randn(3, 4).astype(np.int32) ms_tensor = ms_torch.tensor(np_1) ms_result_1, ms_result_2 = ms_torch.cummin(ms_tensor, dim=1) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch.cummin(torch_tensor, dim=1) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype if __name__ == "__main__": test_cummin1() test_cummin2() test_cummin3() test_cummin4() msadapter/testing/ut/pytorch/functional/test_cumprod.py000066400000000000000000000027001436713554100240670ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_cumprod1(): np_1 = np.random.randn(3, 4).astype(np.float64) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cumprod(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result = torch.cumprod(torch_tensor, dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cumprod2(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cumprod(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result = torch.cumprod(torch_tensor, dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cumprod3(): np_1 = np.random.randn(3, 4).astype(np.int64) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cumprod(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result = torch.cumprod(torch_tensor, dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == "__main__": test_cumprod1() test_cumprod2() test_cumprod3() msadapter/testing/ut/pytorch/functional/test_cumsum.py000066400000000000000000000026571436713554100237420ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_cumsum1(): np_1 = np.random.randn(3, 4).astype(np.float64) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cumsum(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result = torch.cumsum(torch_tensor, dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cumsum2(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cumsum(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result = torch.cumsum(torch_tensor, dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_cumsum3(): np_1 = np.random.randn(3, 4).astype(np.int64) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.cumsum(ms_tensor, dim=0) torch_tensor = torch.tensor(np_1) torch_result = torch.cumsum(torch_tensor, dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == "__main__": test_cumsum1() test_cumsum2() test_cumsum3() msadapter/testing/ut/pytorch/functional/test_diag.py000066400000000000000000000062331436713554100233270ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_diag1(): np_1 = np.array([1, 2, 3]) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diag(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diag(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diag2(): np_1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diag(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diag(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diag3(): np_1 = np.array([1, 2, 3]) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diag(ms_tensor_1, 1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diag(torch_tensor_1, 1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagonal1(): np_1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diagonal(ms_tensor_1, 0) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diagonal(torch_tensor_1, 0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagonal2(): np_1 = np.random.randn(2, 5, 4, 2) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diagonal(ms_tensor_1, offset=-1, dim1=1, dim2=2) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diagonal(torch_tensor_1, offset=-1, dim1=1, dim2=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) # TODO: mindspore return float32 != torch return float64 # assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagonal3(): np_1 = np.random.randn(3, 3) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diagonal(ms_tensor_1, 0) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diagonal(torch_tensor_1, 0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagonal4(): np_1 = np.random.randn(4, 5, 6, 7) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diagonal(ms_tensor_1, offset=2, dim1=1, dim2=2) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diagonal(torch_tensor_1, offset=2, dim1=1, dim2=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) #TODO: mindspore return float32 != torch return float64 # assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_diag1() test_diag2() test_diag3() test_diagonal1() test_diagonal2() test_diagonal3() test_diagonal4() msadapter/testing/ut/pytorch/functional/test_diagflat.py000066400000000000000000000053001436713554100241700ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_diagflat1(): np_1 = np.random.randn(8).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.diagflat(ms_tensor, offset=-1) torch_tensor = torch.tensor(np_1) torch_result = torch.diagflat(torch_tensor, offset=-1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagflat2(): np_1 = np.random.randn(8).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.diagflat(ms_tensor, offset=0) torch_tensor = torch.tensor(np_1) torch_result = torch.diagflat(torch_tensor, offset=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagflat3(): np_1 = np.random.randn(8).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.diagflat(ms_tensor, offset=2) torch_tensor = torch.tensor(np_1) torch_result = torch.diagflat(torch_tensor, offset=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagflat4(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.diagflat(ms_tensor, offset=-1) torch_tensor = torch.tensor(np_1) torch_result = torch.diagflat(torch_tensor, offset=-1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagflat5(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.diagflat(ms_tensor, offset=0) torch_tensor = torch.tensor(np_1) torch_result = torch.diagflat(torch_tensor, offset=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagflat6(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.diagflat(ms_tensor, offset=2) torch_tensor = torch.tensor(np_1) torch_result = torch.diagflat(torch_tensor, offset=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == "__main__": test_diagflat1() test_diagflat2() test_diagflat3() test_diagflat4() test_diagflat5() test_diagflat6() msadapter/testing/ut/pytorch/functional/test_diff.py000066400000000000000000000047701436713554100233370ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_diff1(): np_1 = np.random.randn(4, 5, 6, 7) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diff(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diff(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-6) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diff2(): np_1 = np.random.randn(3, 4, 5) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diff(ms_tensor_1, dim=0) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diff(torch_tensor_1, dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-6) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diff3(): np_1 = np.random.randn(4, 5) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.diff(ms_tensor_1, dim=1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.diff(torch_tensor_1, dim=1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-6) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diff4(): np_1 = np.random.randn(7, 8) np_2 = np.random.randn(7, 8) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.diff(ms_tensor_1, dim=-1, append=ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.diff(torch_tensor_1, dim=-1, append=torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-6) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diff5(): np_1 = np.random.randn(7, 8) np_2 = np.random.randn(7, 8) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.diff(ms_tensor_1, n=3, dim=-1, append=ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.diff(torch_tensor_1, n=3, dim=-1, append=torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-6) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_diff1() test_diff2() test_diff3() test_diff4() test_diff5() msadapter/testing/ut/pytorch/functional/test_div.py000066400000000000000000000050231436713554100232010ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_div1(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[1.0, 2],[3, 4]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.div(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.div(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_div2(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int32) np_2 = np.array([[1, 2],[3, 4]]).astype(np.int32) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.div(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.div(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype # TODO # mindspore.ops.trunc has some problem on code gate version of mindspore ''' def test_div3(): np_1 = np.array([[3, 4],[3, 4]]) np_2 = np.array([[2.0, 2.0],[2.0, 2.0]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.div(ms_tensor_1, ms_tensor_2, rounding_mode='trunc') torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.div(torch_tensor_1, torch_tensor_2, rounding_mode='trunc') assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype ''' def test_div4(): np_1 = np.array([[1, 2],[3, 4]]) np_2 = np.array([[2, 2],[2, 2]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.div(ms_tensor_1, ms_tensor_2, rounding_mode='floor') torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.div(torch_tensor_1, torch_tensor_2, rounding_mode='floor') assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_div1() test_div2() test_div3() test_div4() msadapter/testing/ut/pytorch/functional/test_empty.py000066400000000000000000000010601436713554100235520ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_empty1(): ms_result = ms_torch.empty(2, 3, dtype=ms_torch.float32) torch_result = torch.empty(2, 3, dtype=torch.float32) def test_empty2(): ms_result = ms_torch.empty((2, 3), dtype=ms_torch.float32) torch_result = torch.empty((2, 3), dtype=torch.float32) if __name__ == '__main__': test_empty1() test_empty2() msadapter/testing/ut/pytorch/functional/test_eye.py000066400000000000000000000012461436713554100232040ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_eye1(): ms_result = ms_torch.eye(n=2, m=3, dtype=ms_torch.float32) torch_result = torch.eye(n=2, m=3, dtype=torch.float32) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_eye2(): ms_result = ms_torch.eye(n=2, dtype=ms_torch.float32) torch_result = torch.eye(n=2, dtype=torch.float32) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) if __name__ == '__main__': test_eye1() test_eye2() msadapter/testing/ut/pytorch/functional/test_flatten.py000066400000000000000000000023251436713554100240560ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_flatten1(): np_1 = np.ones((2,3,4,5)).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.flatten(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.flatten(torch_tensor) assert ms_result.asnumpy().shape == torch_result.numpy().shape assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-4) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_flatten2(): np_1 = np.ones((2,3,4,5)).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.flatten(ms_tensor, start_dim=1, end_dim=-2) torch_tensor = torch.tensor(np_1) torch_result = torch.flatten(torch_tensor, start_dim=1, end_dim=-2) assert ms_result.asnumpy().shape == torch_result.numpy().shape assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-4) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_flatten1() test_flatten2() msadapter/testing/ut/pytorch/functional/test_flip.py000066400000000000000000000064411436713554100233560ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_flip1(): np_1 = np.random.randn(4, 4) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.flip(ms_tensor_1, (1, )) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.flip(torch_tensor_1, (1, )) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_flip2(): np_1 = np.random.randn(4, 5, 6, 7) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.flip(ms_tensor_1, (0, 3)) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.flip(torch_tensor_1, (0, 3)) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_flip3(): np_1 = np.random.randn(3, 4, 5) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.flip(ms_tensor_1, (0, 1)) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.flip(torch_tensor_1, (0, 1)) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_flip4(): np_1 = np.reshape(np.arange(8), (2, 2, 2)) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.flip(ms_tensor_1, [0, 1]) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.flip(torch_tensor_1, [0, 1]) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_fliplr1(): np_1 = np.random.randn(4, 4) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.fliplr(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.fliplr(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_fliplr2(): np_1 = np.random.randn(3, 4, 5) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.fliplr(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.fliplr(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_fliplr3(): np_1 = np.random.randn(4, 5, 6, 7) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.fliplr(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.fliplr(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_fliplr4(): np_1 = np.reshape(np.arange(8), (2, 2, 2)) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.fliplr(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.fliplr(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_flip1() test_flip2() test_flip3() test_flip4() test_fliplr1() test_fliplr2() test_fliplr3() test_fliplr4() msadapter/testing/ut/pytorch/functional/test_function.py000066400000000000000000000546161436713554100242600ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_range1(): ms_result = ms_torch.range(1, 10) torch_result = torch.range(1, 10) assert np.allclose(ms_result.numpy(), torch_result.numpy()) def test_range2(): ms_result = ms_torch.range(1.0, 3, 0.5) torch_result = torch.range(1.0, 3, 0.5) assert np.allclose(ms_result.numpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_astensor(): a = np.array([1, 2, 3]) torch_r = torch.as_tensor(a) ms_r = ms_torch.as_tensor(a) assert np.allclose(torch_r.numpy(), ms_r.numpy()) def test_zeros_like(): input_t = torch.empty(2, 3) torch_r = torch.zeros_like(input_t) input_ms = ms_torch.empty(2, 3) ms_r = ms_torch.zeros_like(input_ms) assert np.allclose(torch_r.numpy(), ms_r.numpy()) def test_ones_like(): input_t = torch.empty(2, 3) torch_r = torch.ones_like(input_t) input_ms = ms_torch.empty(2, 3) ms_r = ms_torch.ones_like(input_ms) assert np.allclose(torch_r.numpy(), ms_r.numpy()) def test_empty_like(): input_t = torch.ones((2, 3)) torch_r = torch.empty_like(input_t) input_ms = ms_torch.ones((2, 3)) ms_r = ms_torch.empty_like(input_ms) assert np.allclose(torch_r.shape, ms_r.shape) def test_full(): torch_r = torch.full((2, 3), 3.141592) ms_r = ms_torch.full((2, 3), 3.141592) assert np.allclose(torch_r.numpy(), ms_r.numpy()) def test_full_like(): input_t = torch.full((2, 3), 3.141592) torch_r = torch.full_like(input_t, 0.5) input_ms = ms_torch.full((2, 3), 3.141592) ms_r = ms_torch.full_like(input_ms, 0.5) assert np.allclose(torch_r.numpy(), ms_r.numpy()) def test_where(): x = torch.randn(3, 2) y = torch.ones(3, 2) output1 = torch.where(x > 0, x, y) x1 = torch.randn(2, 2, dtype=torch.double) output2 = torch.where(x1 > 0, x1, 0.) x_m = ms.Tensor(x.numpy()) y_m = ms.Tensor(y.numpy()) output1_ms = ms_torch.where(x_m > 0, x_m, y_m) x1_m = ms.Tensor(x1.numpy()) output2_ms = ms_torch.where(x1_m > 0, x1_m, 0.) assert np.allclose(output1.numpy(), output1_ms.numpy()) assert np.allclose(output2.numpy(), output2_ms.numpy()) def test_seed(): torch.manual_seed(12) x = torch.empty((1, 4)) t_r = torch.nn.init.uniform(x) torch.seed() t_r1 = torch.nn.init.uniform(x) seed_v = torch.initial_seed() ms_torch.manual_seed(12) ms_r = ms.ops.uniform((1, 4), ms.Tensor(1.0, ms.float32), ms.Tensor(2.0, ms.float32)) ms_torch.seed() ms_r1 = ms.ops.uniform((1, 4), ms.Tensor(1.0, ms.float32), ms.Tensor(2.0, ms.float32)) assert np.allclose(t_r.shape, ms_r.shape) assert np.allclose(t_r1.shape, ms_r1.shape) def test_rand(): t_r1 = torch.rand(4) t_r2 = torch.rand(2, 3) ms_r1 = ms_torch.rand(4) ms_r2 = ms_torch.rand(2, 3) assert np.allclose(t_r1.shape, ms_r1.shape) assert np.allclose(t_r2.shape, ms_r2.shape) def test_linspace(): t_r1 = torch.linspace(3, 10, steps=5) t_r2 = torch.linspace(-10, 10, steps=5) t_r3 = torch.linspace(start=-10, end=10, steps=5) t_r4 = torch.linspace(start=-10, end=10, steps=1) ms_r1 = ms_torch.linspace(3, 10, steps=5) ms_r2 = ms_torch.linspace(-10, 10, steps=5) ms_r3 = ms_torch.linspace(start=-10, end=10, steps=5) ms_r4 = ms_torch.linspace(start=-10, end=10, steps=1) assert np.allclose(t_r1.numpy(), ms_r1.numpy()) assert np.allclose(t_r2.numpy(), ms_r2.numpy()) assert np.allclose(t_r3.numpy(), ms_r3.numpy()) assert np.allclose(t_r4.numpy(), ms_r4.numpy()) def test_take(): src = torch.tensor([[4, 3, 5], [6, 7, 8]]) output_t = torch.take(src, torch.tensor([0, 2, 5])) ms_src = ms_torch.tensor([[4, 3, 5], [6, 7, 8]]) output_ms = ms_torch.take(ms_src, ms_torch.tensor([0, 2, 5])) assert np.allclose(output_t.numpy(), output_ms.numpy()) def test_abs(): t_r = torch.abs(torch.tensor([-1, -2, 3])) ms_r = ms_torch.abs(ms_torch.tensor([-1, -2, 3])) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_atan2(): a = torch.tensor([0., 1.]) b = torch.tensor([1., 1.]) t_r = torch.atan2(a, b) ms_a = ms_torch.tensor([0., 1.]) ms_b = ms_torch.tensor([1., 1.]) ms_r = ms_torch.atan2(ms_a, ms_b) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_clamp(): a = torch.tensor([[1., 25., 5., 7.], [4., 11., 6., 21.]], dtype=torch.float32) t_r = torch.clamp(a, min=5, max=20) input = ms_torch.tensor(np.array([[1., 25., 5., 7.], [4., 11., 6., 21.]]), ms_torch.float32) ms_r = ms_torch.clamp(input, min=5, max=20) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_cos(): x = torch.tensor([0.24, 0.83, 0.31, 0.09], dtype=torch.float32) t_r = torch.cos(x) y = ms_torch.tensor([0.24, 0.83, 0.31, 0.09], ms_torch.float32) ms_r = ms_torch.cos(y) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_fft(): t = torch.arange(0, 4) t_r = torch.fft.fft(t) ms_t = ms_torch.arange(0, 4) ms_r = ms_torch.fft.fft(ms_t) t1 = torch.tensor([0. + 1.j, 2. + 3.j, 4. + 5.j, 6. + 7.j]) t_r1 = torch.fft.fft(t1) t1_ms = ms_torch.tensor([0. + 1.j, 2. + 3.j, 4. + 5.j, 6. + 7.j]) ms_r1 = ms_torch.fft.fft(t1_ms) assert np.allclose(t_r.numpy(), ms_r.numpy()) assert np.allclose(t_r1.numpy(), ms_r1.numpy()) def test_fmod(): t_r = torch.fmod(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) t_r1 = torch.fmod(torch.tensor([1, 2, 3, 4, 5]), -1.5) ms_r = ms_torch.fmod(ms_torch.tensor([-3., -2, -1, 1, 2, 3]), 2) ms_r1 = ms_torch.fmod(ms_torch.tensor([1, 2, 3, 4, 5]), -1.5) assert np.allclose(t_r.numpy(), ms_r.numpy()) assert np.allclose(t_r1.numpy(), ms_r1.numpy()) def test_frac(): t_r = torch.frac(torch.tensor([1, 2.5, -3.2])) ms_r = ms_torch.frac(ms_torch.tensor([1, 2.5, -3.2])) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_log(): x = torch.tensor([0.24, 0.83, 0.31, 0.09], dtype=torch.float32) t_r = torch.log10(x) t_r1 = torch.log1p(x) t_r2 = torch.log2(x) y = ms_torch.tensor([0.24, 0.83, 0.31, 0.09], ms_torch.float32) ms_r = ms_torch.log10(y) ms_r1 = ms_torch.log1p(y) ms_r2 = ms_torch.log2(y) assert np.allclose(t_r.numpy(), ms_r.numpy(), 1e-3) assert np.allclose(t_r1.numpy(), ms_r1.numpy(), 1e-3) assert np.allclose(t_r2.numpy(), ms_r2.numpy(), 1e-3) def test_sin(): x = torch.tensor([0.24, 0.83, 0.31, 0.09], dtype=torch.float32) t_r = torch.sin(x) y = ms_torch.tensor([0.24, 0.83, 0.31, 0.09], ms_torch.float32) ms_r = ms_torch.sin(y) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_norm(): x = torch.tensor([0.24, 0.83, 0.31, 0.09], dtype=torch.float32) t_r = torch.norm(x, 2) y = ms_torch.tensor([0.24, 0.83, 0.31, 0.09], ms_torch.float32) ms_r = ms_torch.norm(y, 2) x1 = torch.tensor([0.24, 0.83, 0.31, 0.09], dtype=torch.float32) t_r1 = torch.norm(x1, float('inf')) y1 = ms_torch.tensor([0.24, 0.83, 0.31, 0.09], ms_torch.float32) ms_r1 = ms_torch.norm(y1, float('inf')) assert np.allclose(t_r.numpy(), ms_r.numpy()) assert np.allclose(t_r1.numpy(), ms_r1.numpy()) def test_bartlett_window(): t_r = torch.bartlett_window(5) ms_r = ms_torch.bartlett_window(5) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_hamming_window(): t_r = torch.hamming_window(12) ms_r = ms_torch.hamming_window(12) assert np.allclose(t_r.shape, ms_r.shape) def test_hann_windoww(): t_r = torch.hann_window(12, periodic=False) ms_r = ms_torch.hann_window(12) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_cumsum(): x = torch.tensor([0.24, 0.83, 0.31, 0.09], dtype=torch.float32) t_r = torch.cumsum(x, dim=0) y = ms_torch.tensor([0.24, 0.83, 0.31, 0.09], ms_torch.float32) ms_r = ms_torch.cumsum(y, dim=0) assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_einsum(): # trace in_1 = torch.randn(4, 4) t_r = torch.einsum('ii', in_1) # ms_r = ms_torch.einsum('ii', ms_torch.tensor(in_1.numpy())) # diagonal torch.einsum('ii->i', torch.randn(4, 4)) # outer product x = torch.randn(5) y = torch.randn(4) torch.einsum('i,j->ij', x, y) # batch matrix multiplication As = torch.randn(3, 2, 5) Bs = torch.randn(3, 5, 4) torch.einsum('bij,bjk->bik', As, Bs) # with sublist format and ellipsis torch.einsum(As, [..., 0, 1], Bs, [..., 1, 2], [..., 0, 2]) # batch permute A = torch.randn(2, 3, 4, 5) torch.einsum('...ij->...ji', A) # equivalent to torch.nn.functional.bilinear A = torch.randn(3, 5, 4) l = torch.randn(2, 5) r = torch.randn(2, 4) torch.einsum('bn,anm,bm->ba', l, A, r) def test_histc(): data1 = np.array([1, 2, 1, 0, -1, -2, 2, 2, 3, 3, 4, 5, 6]).astype(np.float32) t_r1 = torch.histc(torch.tensor(data1), bins=4, min=3, max=3) ms_r1 = ms_torch.histc(ms_torch.tensor(data1), bins=4, min=3, max=3) t_r2 = torch.histc(torch.tensor([[1., 2, 1, 0, -1, -2, 2.1, 2.9, 3, 3.1, 4, 5, 6]]), bins=4, min=0, max=3) ms_r2 = ms_torch.histc(ms_torch.tensor([[1., 2, 1, 0, -1, -2, 2.1, 2.9, 3, 3.1, 4, 5, 6]]), bins=4, min=0, max=3) t_r3 = torch.histc(torch.tensor([1., 1, 1]), bins=4, min=3, max=3) ms_r3 = ms_torch.histc(ms_torch.tensor([1., 1, 1]), bins=4, min=3, max=3) assert np.allclose(t_r1.numpy(), ms_r1.numpy()) assert np.allclose(t_r2.numpy(), ms_r2.numpy()) assert np.allclose(t_r3.numpy(), ms_r3.numpy()) def test_triu(): a = torch.randn(3, 3) t_r = torch.triu(a) t_r1 = torch.triu(a, diagonal=1) t_r2 = torch.triu(a, diagonal=-1) ms_r = ms_torch.triu(ms_torch.tensor(a.numpy())) ms_r1 = ms_torch.triu(ms_torch.tensor(a.numpy()), diagonal=1) ms_r2 = ms_torch.triu(ms_torch.tensor(a.numpy()), diagonal=-1) assert np.allclose(t_r.numpy(), ms_r.numpy()) assert np.allclose(t_r1.numpy(), ms_r1.numpy()) assert np.allclose(t_r2.numpy(), ms_r2.numpy()) def test_index_select(): data = np.random.randn(3, 4 ,5) x_torch = torch.tensor(data) indices = torch.tensor([0, 2]) torch_out = torch.index_select(x_torch, 1, indices) x_ms = ms_torch.tensor(data) indices = ms_torch.tensor([0, 2]) ms_out = ms_torch.index_select(x_ms, 1, indices) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_bmm(): input = np.random.randn(2, 3, 4) mat2 = np.random.randn(2, 4, 5) torch_input = torch.tensor(input) torch_mat2 = torch.tensor(mat2) torch_out = torch.bmm(torch_input, torch_mat2) ms_input = ms_torch.tensor(input) ms_mat2 = ms_torch.tensor(mat2) ms_out = ms_torch.bmm(ms_input, ms_mat2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) def test_baddbmm(): x = np.random.randn(10, 3, 5) b1 = np.random.randn(10, 3, 4) b2 = np.random.randn(10, 4, 5) torch_input = torch.tensor(x) torch_batch1 = torch.tensor(b1) torch_batch2 = torch.tensor(b2) torch_out = torch.baddbmm(torch_input, torch_batch1, torch_batch2, beta=2, alpha=1.5) ms_torch_input = ms_torch.tensor(x) ms_torch_batch1 = ms_torch.tensor(b1) ms_torch_batch2 = ms_torch.tensor(b2) ms_torch_out = ms_torch.baddbmm(ms_torch_input, ms_torch_batch1, ms_torch_batch2, beta=2, alpha=1.5) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_argmin(): x = np.random.randn(2, 3, 2) torch_x = torch.tensor(x) torch_out = torch.argmin(torch_x, 1, False) ms_x = ms_torch.tensor(x) ms_out = ms_torch.argmin(ms_x, 1, False) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_argmax(): x = np.random.randn(2, 3, 2) torch_x = torch.tensor(x) torch_out = torch.argmax(torch_x, 1, False) ms_x = ms_torch.tensor(x) ms_out = ms_torch.argmax(ms_x, 1, False) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_broadcast_to(): torch_x = torch.tensor([1]) torch_out = torch.broadcast_to(torch_x, (2,2)) ms_x = ms_torch.tensor([1]) ms_out = ms_torch.broadcast_to(ms_x, (2,2)) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_ravel(): x = np.random.randn(2, 1, 2) torch_x = torch.tensor(x, dtype=torch.bool) torch_out = torch.ravel(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.bool) ms_out = ms_torch.ravel(ms_x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_numel(): tensor = np.zeros((1, 2, 3, 1, 2)).astype(np.int32) torch_tensor = torch.tensor(tensor) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.numel(torch_tensor) ms_output = ms_torch.numel(ms_tensor) assert np.allclose(ms_output, torch_output) def test_logsumexp(): x = np.random.randn(3, 4, 5).astype(np.float32) torch_x = torch.tensor(x) torch_out = torch.logsumexp(torch_x, dim=2, keepdim=True) ms_torch_x = ms_torch.tensor(x) ms_torch_out = ms_torch.logsumexp(ms_torch_x, dim=2, keepdim=True) assert np.allclose(torch_out.size(), ms_torch_out.size()) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_addmv(): M = np.random.randn(2).astype(np.float32) mat = np.random.randn(2, 3).astype(np.float32) vec = np.random.randn(3).astype(np.float32) torch_M = torch.tensor(M) torch_mat = torch.tensor(mat) torch_vec = torch.tensor(vec) torch_out = torch.addmv(torch_M, torch_mat, torch_vec, beta=False, alpha=0.5) ms_torch_M = ms_torch.tensor(M) ms_torch_mat = ms_torch.tensor(mat) ms_torch_vec = ms_torch.tensor(vec) ms_torch_out = ms_torch.addmv(ms_torch_M, ms_torch_mat, ms_torch_vec, beta=False, alpha=0.5) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_dot(): x = [2, 3, 1] y = [2, 1, 2] torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch.dot(torch_x, torch_y) ms_torch_x = ms_torch.tensor(x) ms_torch_y = ms_torch.tensor(y) ms_torch_out = ms_torch.dot(ms_torch_x, ms_torch_y) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_inverse(): A1 = np.random.randn(4, 4).astype(np.float32) torch_A1 = torch.tensor(A1) torch_out1 = torch.inverse(torch_A1) ms_torch_A1 = ms_torch.tensor(A1) ms_torch_out1 = ms_torch.inverse(ms_torch_A1) assert np.allclose(torch_out1.numpy(), ms_torch_out1.numpy()) def test_scatter_1(): ms_out = ms_torch.scatter(ms_torch.full((2, 4), 2.), 1, ms_torch.tensor([[2], [3]]), 1.23) torch_out = torch.scatter(torch.full((2, 4), 2.), 1, torch.tensor([[2], [3]]), 1.23) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_topk(): x = np.random.randn(10).astype(np.float32) torch_x = torch.tensor(x) torch_v, torch_i = torch_x.topk(3) ms_x = ms_torch.tensor(x) ms_v, ms_i = ms_x.topk(3) assert np.allclose(torch_v.numpy(), ms_v.numpy()) assert np.allclose(torch_i.numpy(), ms_i.numpy()) def test_addbmm(): M_ = np.random.randn(3, 5) batch1_ = np.random.randn(10, 3, 4) batch2_ = np.random.randn(10, 4, 5) M = torch.tensor(M_) batch1 = torch.tensor(batch1_) batch2 = torch.tensor(batch2_) torch_output = torch.addbmm(M, batch1, batch2, alpha=2, beta=3) M = ms_torch.tensor(M_) batch1 = ms_torch.tensor(batch1_) batch2 = ms_torch.tensor(batch2_) ms_output = ms_torch.addbmm(M, batch1, batch2, alpha=2, beta=3) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_addr(): vec1 = torch.arange(1., 4.) vec2 = torch.arange(1., 3.) M = torch.zeros(3, 2) torch_output = torch.addr(M, vec1, vec2) vec1 = ms_torch.arange(1., 4.) vec2 = ms_torch.arange(1., 3.) M = ms_torch.zeros(3, 2) ms_output = ms_torch.addr(M, vec1, vec2) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_allclose(): a1 = ms_torch.allclose(ms_torch.tensor([10000., 1e-07]), ms_torch.tensor([10000.1, 1e-08])) b1 = ms_torch.allclose(ms_torch.tensor([10000., 1e-08]), ms_torch.tensor([10000.1, 1e-09])) c1 = ms_torch.allclose(ms_torch.tensor([1.0, float('nan')]), ms_torch.tensor([1.0, float('nan')])) d1 = ms_torch.allclose(ms_torch.tensor([1.0, float('nan')]), ms_torch.tensor([1.0, float('nan')]), equal_nan=True) a2 = torch.allclose(torch.tensor([10000., 1e-07]), torch.tensor([10000.1, 1e-08])) b2 = torch.allclose(torch.tensor([10000., 1e-08]), torch.tensor([10000.1, 1e-09])) c2 = torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')])) d2 = torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')]), equal_nan=True) assert a1 == a2 assert b1 == b2 assert c1 == c2 assert d1 == d2 def test_isclose(): a1 = ms_torch.isclose(ms_torch.tensor((1., 2, 3)), ms_torch.tensor((1 + 1e-10, 3, 4))) b1 = ms_torch.isclose(ms_torch.tensor((float('inf'), 4)), ms_torch.tensor((float('inf'), 6)), rtol=.5) a2 = torch.isclose(torch.tensor((1., 2, 3)), torch.tensor((1 + 1e-10, 3, 4))) b2 = torch.isclose(torch.tensor((float('inf'), 4)), torch.tensor((float('inf'), 6)), rtol=.5) assert np.allclose(a1.numpy(), a2.numpy()) assert np.allclose(b1.numpy(), b2.numpy()) def test_addmm(): _M = np.random.randn(2, 3) _mat1 = np.random.randn(2, 3) _mat2 = np.random.randn(3, 3) M = ms_torch.tensor(_M) mat1 = ms_torch.tensor(_mat1) mat2 = ms_torch.tensor(_mat2) a1 = ms_torch.addmm(M, mat1, mat2, alpha=2, beta=3) M = torch.tensor(_M) mat1 = torch.tensor(_mat1) mat2 = torch.tensor(_mat2) a2 = torch.addmm(M, mat1, mat2, alpha=2, beta=3) assert np.allclose(a1.numpy(), a2.numpy()) def test_cholesky(): _data1 = np.random.randn(3, 3) _data2 = np.random.randn(3, 2, 2) a = torch.tensor(_data1) a = a @ a.mT + 1e-3 torch_out1 = torch.cholesky(a) a = torch.tensor(_data2) a = a @ a.mT + 1e-03 torch_out2 = torch.cholesky(a) a = ms_torch.tensor(_data1) a = a @ a.mT + 1e-3 ms_out1 = ms_torch.cholesky(a) a = ms_torch.tensor(_data2) a = a @ a.mT + 1e-03 ms_out2 = ms_torch.cholesky(a) assert np.allclose(torch_out1.numpy(), ms_out1.numpy()) assert np.allclose(torch_out2.numpy(), ms_out2.numpy()) def test_dist(): _data1 = np.random.randn(4) _data2 = np.random.randn(4) x = torch.tensor(_data1) y = torch.tensor(_data2) # TODO: Not support float p yet #a1 = torch.dist(x, y, 3.5) b1 = torch.dist(x, y, 3) c1 = torch.dist(x, y, 0) d1 = torch.dist(x, y, 1) x = ms_torch.tensor(_data1) y = ms_torch.tensor(_data2) #a2 = ms_torch.dist(x, y, 3.5) b2 = ms_torch.dist(x, y, 3) c2 = ms_torch.dist(x, y, 0) d2 = ms_torch.dist(x, y, 1) #assert np.allclose(a1.numpy(), a2.numpy()) assert np.allclose(b1.numpy(), b2.numpy()) assert np.allclose(c1.numpy(), c2.numpy()) assert np.allclose(d1.numpy(), d2.numpy()) def test_aminmax(): a1, b1 = torch.aminmax(torch.tensor([1, -3, 5])) a2, b2 = ms_torch.aminmax(ms_torch.tensor([1, -3, 5])) assert np.allclose(a1.numpy(), a2.numpy()) assert np.allclose(b1.numpy(), b2.numpy()) def test_any(): data1 = np.random.randn(1, 2) data2 = np.random.randn(4, 2) a = torch.tensor(data1).bool() a1 = torch.any(a) a = torch.arange(0, 3) b1 = torch.any(a) a = torch.tensor(data2) < 0 c1 = torch.any(a, 1) d1 = torch.any(a, 0) a = torch.tensor(data1).bool() a2 = torch.any(a) a = torch.arange(0, 3) b2 = torch.any(a) a = torch.tensor(data2) < 0 c2 = torch.any(a, 1) d2 = torch.any(a, 0) assert np.allclose(a1.numpy(), a2.numpy()) assert np.allclose(b1.numpy(), b2.numpy()) assert np.allclose(c1.numpy(), c2.numpy()) assert np.allclose(d1.numpy(), d2.numpy()) def test_cholesky_inverse(): _data = np.random.randn(3, 3) a = torch.tensor(_data) a = torch.mm(a, a.t()) + 1e-05 * torch.eye(3) u = torch.cholesky(a) torch_out = torch.cholesky_inverse(u) a = ms_torch.tensor(_data) a = ms_torch.mm(a, a.t()) + 1e-05 * ms_torch.eye(3) u = ms_torch.cholesky(a) ms_out = ms_torch.cholesky_inverse(u) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_iscomplex(): a = ms_torch.tensor([1+1j, 2, 3]) assert ms_torch.is_complex(a) == True a = ms_torch.tensor([1, 2, 3]) assert ms_torch.is_complex(a) == False def test_isinf(): data = [1, float('inf'), 2, float('-inf'), float('nan')] a = torch.isinf(torch.tensor(data)) b = ms_torch.isinf(ms_torch.tensor(data)) assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype def test_isneginf(): data = [-float('inf'), float('inf'), 1.2] a = torch.isneginf(torch.tensor(data)) b = ms_torch.isneginf(ms_torch.tensor(data)) assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype def test_isposinf(): data = [-float('inf'), float('inf'), 1.2] a = torch.isposinf(torch.tensor(data)) b = ms_torch.isposinf(ms_torch.tensor(data)) assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype def test_isreal(): data = [1, 1+1j, 2+0j] a = torch.isreal(torch.tensor(data)) b = ms_torch.isreal(ms_torch.tensor(data)) assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype if __name__ == '__main__': test_range1() test_range2() test_astensor() test_zeros_like() test_ones_like() test_empty_like() test_full() test_full_like() test_where() test_seed() test_rand() test_linspace() test_take() test_abs() test_atan2() test_clamp() test_cos() test_fft() test_fmod() test_frac() test_log() test_sin() test_norm() test_bartlett_window() test_hamming_window() test_hann_windoww() test_cumsum() test_einsum() test_histc() test_triu() test_index_select() test_bmm() test_baddbmm() test_argmin() test_argmax() test_broadcast_to() test_ravel() test_numel() test_logsumexp() test_addmv() test_dot() test_inverse() test_scatter_1() test_topk() test_addbmm() test_allclose() test_isclose() test_addmm() test_cholesky() test_dist() test_aminmax() test_any() test_cholesky_inverse() test_isinf() test_isneginf() test_isposinf() test_isreal()msadapter/testing/ut/pytorch/functional/test_gather.py000066400000000000000000000034141436713554100236730ustar00rootroot00000000000000 #!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_gather_2d(): input = [[1, 2], [3, 4]] index = [[0, 0], [1, 0]] tourch_input = torch.tensor(input) tourch_index = torch.tensor(index) torch_out = torch.gather(tourch_input, 1, tourch_index) ms_input = ms_torch.tensor(input) ms_index = ms_torch.tensor(index) ms_out = ms_torch.gather(ms_input, 1, ms_index) assert np.allclose(torch_out.size(), ms_out.size()) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_gather_2d_long(): x = [ [2, 3, 4, 5, 0, 0], [1, 4, 3, 0, 0, 0], [4, 2, 2, 5, 7, 0], [1, 0, 0, 0, 0, 0]] input = torch.tensor(x) length = torch.LongTensor([[4], [3], [5], [1]]) torch_out = torch.gather(input, 1, length - 1) ms_input = ms_torch.tensor(x) ms_length = ms_torch.LongTensor([[4], [3], [5], [1]]) ms_out = ms_torch.gather(ms_input, 1, ms_length - 1) assert np.allclose(torch_out.size(), ms_out.size()) assert np.allclose(torch_out.numpy(), ms_out.numpy()) return def test_gather_3d(): input = np.random.randn(2, 2, 2) index = [[[1, 0],[0, 0]],[[0, 1],[0, 1]]] tourch_input = torch.tensor(input) tourch_index = torch.tensor(index) torch_out = torch.gather(tourch_input, 2, tourch_index) ms_input = ms_torch.tensor(input) ms_index = ms_torch.tensor(index) ms_out = ms_torch.gather(ms_input, 2, ms_index) assert np.allclose(torch_out.size(), ms_out.size()) assert np.allclose(torch_out.numpy(), ms_out.numpy()) if __name__ == '__main__': test_gather_2d() test_gather_2d_long() test_gather_3d()msadapter/testing/ut/pytorch/functional/test_log.py000066400000000000000000000011721436713554100232010ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_log1(): np_1 = np.array([0.2, 1, 3]).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.log(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.log(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-4) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_log1() msadapter/testing/ut/pytorch/functional/test_mask.py000066400000000000000000000012351436713554100233530ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_masked_select(): x = np.random.random((3,4)).astype(np.float32) torch_tensor = torch.tensor(x) torch_mask = torch_tensor.ge(0.5) torch_out = torch.masked_select(torch_tensor, torch_mask) ms_tensor = ms_torch.tensor(x) ms_mask = ms_tensor.ge(0.5) ms_out = ms_torch.masked_select(ms_tensor, ms_mask) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) if __name__ == '__main__': test_masked_select() msadapter/testing/ut/pytorch/functional/test_math.py000066400000000000000000001714151436713554100233610ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) #context.set_context(mode=ms.GRAPH_MODE) def test_round1(): np_array = np.array([[1.501, 0.5, 1000.8],[-0.9, -2.5, -24.49]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.round(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.round(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_round2(): np_array = np.random.randn(20).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.round(torch_tensor, decimals=2) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.round(ms_tensor, decimals=2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_round3(): np_array = np.random.randn(20).astype(np.float32)*10**3 torch_tensor = torch.tensor(np_array) torch_out = torch.round(torch_tensor, decimals=-2) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.round(ms_tensor, decimals=-2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_floor1(): np_array = np.array([[1.501, 0.5, 1000.8],[-0.9, -2.5, -24.49]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.floor(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.floor(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_ceil1(): np_array = np.array([[1.501, 0.5, 1000.8],[-0.9, -2.5, -24.49]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.ceil(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.ceil(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_sign1(): np_array = np.array([[1.1, 0], [-0, -0.5]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.sign(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.sign(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_sign2(): np_array = np.array([0, 1, 5, -2, -9, -0, 5]).astype(np.int32) torch_tensor = torch.tensor(np_array) torch_out = torch.sign(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.sign(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pow1(): np_array = np.array([0, 1, 5, -2, -9, -0, 5]).astype(np.float64) pow_array = np.array([-2, 2, 1.2, 0, -2, 2, 2]).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_pow_tensor = torch.tensor(pow_array) torch_out = torch.pow(torch_tensor,torch_pow_tensor) ms_tensor = ms_torch.tensor(np_array) ms_pow_tensor = ms_torch.tensor(pow_array) ms_out = ms_torch.pow(ms_tensor,ms_pow_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pow2(): np_array = np.array([0, 1, 5, -2, -9, -0, 5]).astype(np.int64) torch_tensor = torch.tensor(np_array) torch_out = torch.pow(torch_tensor, False) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.pow(ms_tensor, False) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype # TODO: On GPU and Ascend not support input above 7-dimentions. ''' def test_pow3(): np_array = np.array([[[[[[[[[0, 1, 5, 2, 9, -0, 5]]]]]]]]]).astype(np.int32) torch_tensor = torch.tensor(np_array) torch_out = torch.pow(torch_tensor, -2.2) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.pow(ms_tensor, -2.2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype ''' def test_pow4(): np_array = np.array([0, 1.5, 5, 2, -9.5, -0, 5]).astype(np.float16) torch_tensor = torch.tensor(np_array) torch_out = torch.pow(2, torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.pow(2, ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_exp1(): np_array = np.array([0, 1, 5, 2, 9, -0, 5]).astype(np.double) torch_tensor = torch.tensor(np_array) torch_out = torch.exp(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.exp(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_exp2(): np_array = np.array([0, 1, 5, 2, 9, -0, 5]).astype(np.bool8) torch_tensor = torch.tensor(np_array) torch_out = torch.exp(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.exp(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_exp3(): np_array = np.array([[[[[[[[0, 1]]]]]]]]).astype(np.int16) torch_tensor = torch.tensor(np_array) torch_out = torch.exp(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.exp(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_mul1(): a = np.array([1, 2]) b = np.array([[3], [4]]) torch_tensor_a = torch.tensor(a) torch_tensor_b = torch.tensor(b) torch_out = torch.mul(torch_tensor_a, torch_tensor_b) ms_tensor_a = ms_torch.tensor(a) ms_tensor_b = ms_torch.tensor(b) ms_out = ms_torch.mul(ms_tensor_a, ms_tensor_b) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_mul2(): b = np.array([[3], [4]]) torch_tensor_b = torch.tensor(b) torch_out = torch.mul(2, torch_tensor_b) ms_tensor_b = ms_torch.tensor(b) ms_out = ms_torch.mul(2, ms_tensor_b) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) # assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_absolute(): np_array = np.array([[0, -1.5, -0.55, 2, 9, -0, 5]]).astype(np.float16) torch_tensor = torch.tensor(np_array) torch_out = torch.absolute(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.absolute(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_acos(): np_array = np.random.rand(1, 1, 1, 1, 1, 1, 2, 3, 2).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_out = torch.acos(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.acos(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arccos(): np_array = np.random.rand(1, 2, 3, 4).astype(np.float64)*5 torch_tensor = torch.tensor(np_array) torch_out = torch.arccos(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.arccos(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_acosh(): np_array = np.random.rand(1, 1, 1, 1, 1, 2, 3, 2).astype(np.float64)*50+2 torch_tensor = torch.tensor(np_array) torch_out = torch.acosh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.acosh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arccosh(): np_array = np.random.rand(1, 4, 5, 6).astype(np.float64)*5 np_array[0, 0, 0, 0] = 1 torch_tensor = torch.tensor(np_array) torch_out = torch.arccosh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.arccosh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_add(): np_x = np.random.rand(1, 2, 3, 2).astype(np.int32) np_y = np.random.rand(1, 2, 3, 2).astype(np.int32) for x_dtype in (np.int16, np.int32, np.uint8, np.half, np.double, np.single): for y_dtype in (np.int16, np.int32, np.uint8, np.half, np.double, np.single): np_x_ = np_x.astype(x_dtype) np_y_ = np_y.astype(y_dtype) torch_x = torch.tensor(np_x_) torch_y = torch.tensor(np_y_) torch_out1 = torch.add(torch_x, 2.1, alpha=2) torch_out2 = torch.add(torch_x, torch_y, alpha=2) #torch_out2 = torch.add(torch_x, 2, torch_y) ms_x = ms_torch.tensor(np_x_) ms_y = ms_torch.tensor(np_y_) ms_out1 = ms_torch.add(ms_x, 2.1, alpha=2) ms_out2 = ms_torch.add(ms_x, ms_y, alpha=2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), equal_nan=True) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert ms_out1.asnumpy().shape == torch_out1.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), equal_nan=True) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_addcdiv(): np_array = np.random.rand(1, 2, 3, 2).astype(np.float64) np_x = np.random.rand(1, 2, 3, 2).astype(np.float64) np_y = np.random.rand(1, 2, 3, 2).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_x = torch.tensor(np_x) torch_y = torch.tensor(np_y) torch_out = torch.addcdiv(torch_tensor, torch_x, torch_y, value=2.1) #torch_out = torch.addcdiv(torch_tensor, 2.1, torch_x, torch_y) ms_tensor = ms_torch.tensor(np_array) ms_x = ms_torch.tensor(np_x) ms_y = ms_torch.tensor(np_y) ms_out = ms_torch.addcdiv(ms_tensor, ms_x, ms_y, value=2.1) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_addcmul(): #np_array = np.random.rand(1, 2, 3, 2).astype(np.float64) np_array = np.array([1, 1, 1]).astype(np.float64) #np_x = np.random.rand(1, 2, 3, 2).astype(np.float64) #np_y = np.random.rand(1, 2, 3, 2).astype(np.float64) np_x = np.array([[1], [2], [3]]).astype(np.float64) np_y = np.array([[1, 2, 3]]).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_x = torch.tensor(np_x) torch_y = torch.tensor(np_y) torch_out = torch.addcmul(torch_tensor, torch_x, torch_y, value=2.1) #torch_out = torch.addcmul(torch_tensor, 2.1, torch_x, torch_y) ms_tensor = ms_torch.tensor(np_array) ms_x = ms_torch.tensor(np_x) ms_y = ms_torch.tensor(np_y) ms_out = ms_torch.addcmul(ms_tensor, ms_x, ms_y, value=2.1) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_angle(): np_array = np.array([[[1, -20, -1j, 0j, -0j, -1+1j], [0, -0, 3-3j, -0-0.6j, -1-0.01j, 1+1.732j]]]) torch_tensor = torch.tensor(np_array) torch_out = torch.angle(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.angle(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_asin(): np_array = np.random.rand(1, 1, 1, 1, 1, 1, 2, 3, 2).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_out = torch.asin(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.asin(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arcsin(): np_array = np.random.rand(1, 2, 3, 4).astype(np.float64)*5 torch_tensor = torch.tensor(np_array) torch_out = torch.arcsin(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.arcsin(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_asinh(): np_array = np.random.rand(1, 1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 100 torch_tensor = torch.tensor(np_array) torch_out = torch.asinh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.asinh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arcsinh(): np_array = np.random.rand(1, 2, 3, 2).astype(np.float64)*50 np_array[0, 0, 0, 0] = 0 torch_tensor = torch.tensor(np_array) torch_out = torch.arcsinh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.arcsinh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_atan(): np_array = np.random.rand(1, 1, 1, 1, 1, 1, 2, 3, 2).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_out = torch.atan(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.atan(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arctan(): np_array = np.random.rand(1, 2, 3, 4).astype(np.float64)*5-2 torch_tensor = torch.tensor(np_array) torch_out = torch.arctan(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.arctan(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_atanh(): np_array = np.random.rand(1, 1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 2 torch_tensor = torch.tensor(np_array) torch_out = torch.atanh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.atanh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arctanh(): np_array = np.random.rand(1, 4, 5, 6).astype(np.float64)*2 np_array[0, 0, 0, 0] = 0 torch_tensor = torch.tensor(np_array) torch_out = torch.arctanh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.arctanh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arctan2(): np_array = np.random.rand(1, 4, 5, 6).astype(np.float64)*5-2 np_other = np.random.rand(1, 4, 5, 6).astype(np.float64)*5-2 torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out = torch.arctan2(torch_tensor, torch_other) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out = ms_torch.arctan2(ms_tensor, ms_other) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_bitwise_not(): np_array1 = np.arange(-5, 12).astype(np.uint8) np_array2 = np.arange(-10, 10).astype(np.int16) np_array3 = np.array([[[False, True, True, False, False]]]) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_tensor3 = torch.tensor(np_array3) torch_out1 = torch.bitwise_not(torch_tensor1) torch_out2 = torch.bitwise_not(torch_tensor2) torch_out3 = torch.bitwise_not(torch_tensor3) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_tensor3 = ms_torch.tensor(np_array3) ms_out1 = ms_torch.bitwise_not(ms_tensor1) ms_out2 = ms_torch.bitwise_not(ms_tensor2) ms_out3 = ms_torch.bitwise_not(ms_tensor3) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_bitwise(): np_array1 = np.arange(-5, 11).reshape(1, 2, 2, 4).astype(np.bool8) np_array2 = np.arange(-10, 6).reshape(2, 2, 4).astype(np.bool8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.bitwise_and(torch_tensor1, torch_tensor2) torch_out2 = torch.bitwise_or(torch_tensor1, 3) torch_out3 = torch.bitwise_xor(True, torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.bitwise_and(ms_tensor1, ms_tensor2) ms_out2 = ms_torch.bitwise_or(ms_tensor1, 3) ms_out3 = ms_torch.bitwise_xor(True, ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_bitwise_shift(): np_array1 = np.arange(-5, 11).reshape(2, 2, 4).astype(np.int8) np_array2 = np.array([1,2,7,1,2,5,1,2,-1,-3,-4,0,2,1,0,0]).reshape(2, 2, 4).astype(np.int8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.bitwise_left_shift(torch_tensor1, torch_tensor2) torch_out2 = torch.bitwise_right_shift(torch_tensor1, 3) torch_out3 = torch.bitwise_right_shift(True, torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.bitwise_left_shift(ms_tensor1, ms_tensor2) ms_out2 = ms_torch.bitwise_right_shift(ms_tensor1, 3) ms_out3 = ms_torch.bitwise_right_shift(True, ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_clip(): np_array = np.random.rand(1, 2, 3, 4).astype(np.float32) min = 0.35 max = 0.65 torch_tensor = torch.tensor(np_array) torch_out1 = torch.clip(torch_tensor, torch.tensor(min), torch.tensor(max)) torch_out2 = torch.clip(torch_tensor, 0.85, max) torch_out3 = torch.clip(torch_tensor, None, max) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.clip(ms_tensor, ms_torch.tensor(min), ms_torch.tensor(max)) ms_out2 = ms_torch.clip(ms_tensor, 0.85, max) ms_out3 = ms_torch.clip(ms_tensor, None, max) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_conj_physical(): np_array1 = np.array([[1, 2, 3, 4]]).astype(np.int16) np_array2 = np.array([[1+5j, 2-1j, 3.0, -4j, -0j]]) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.conj_physical(torch_tensor1) torch_out2 = torch.conj_physical(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.conj_physical(ms_tensor1) ms_out2 = ms_torch.conj_physical(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_copysign(): np_other = np.array([[0.0, -0.0, 1, 1.5, -0.5, -20]]).astype(np.half) np_array = np.array([0, -0, 0, 1, 1, 1]).astype(np.int32) for type_input in (np.int32, np.half, np.float64, np.float32): torch_tensor = torch.tensor(np_array.astype(type_input)) ms_tensor = ms_torch.tensor(np_array.astype(type_input)) for type_other in (np.half, np.float64, np.int8, np.int32): torch_other = torch.tensor(np_other.astype(type_other)) torch_out1 = torch.copysign(torch_tensor, torch_other) ms_other = ms_torch.tensor(np_other.astype(type_other)) ms_out1 = ms_torch.copysign(ms_tensor, ms_other) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype torch_out2 = torch.copysign(torch_tensor, -0) ms_out2 = ms_torch.copysign(ms_tensor, -0) assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_cosh(): np_array = np.random.rand(1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 20 torch_tensor = torch.tensor(np_array) torch_out = torch.cosh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.cosh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_deg2rad(): np_array = np.array([[180, -180, 360, 30, 0, 57, 80]]) torch_tensor = torch.tensor(np_array) torch_out = torch.deg2rad(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.deg2rad(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_digamma(): np_array = np.random.rand(1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 20 torch_tensor = torch.tensor(np_array) torch_out = torch.digamma(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.digamma(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_lgamma(): np_array = np.random.rand(1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 20 torch_tensor = torch.tensor(np_array) torch_out = torch.lgamma(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.lgamma(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_erf(): np_array = np.random.rand(1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 4 torch_tensor = torch.tensor(np_array) torch_out = torch.erf(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.erf(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_erfc(): np_array = np.random.rand(1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 4 torch_tensor = torch.tensor(np_array) torch_out = torch.erfc(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.erfc(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_erfinv(): np_array = np.random.rand(2, 3, 4).astype(np.float64) - 0.5 np_array = np_array * 1.9 np_array[0, 0, 0:3] = [-1, 1, 0.98] torch_tensor = torch.tensor(np_array) torch_out = torch.erfinv(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.erfinv(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_exp2_(): np_array = np.random.rand(2, 3, 4).astype(np.float64) - 0.5 np_array = np_array * 10 np_array[0, 0, 0:3] = [0, 10, 25] torch_tensor = torch.tensor(np_array) torch_out = torch.exp2(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.exp2(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_expm1(): np_array = np.random.rand(2, 3, 4).astype(np.float32) - 0.5 np_array = np_array * 10 np_array[0, 0, 0:3] = [0, 10, 25] torch_tensor = torch.tensor(np_array) torch_out = torch.expm1(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.expm1(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_fake_quantize(): np_array = np.random.randn(2, 2, 12).astype(np.float32)*10 scales = (np.random.randn(2).astype(np.float32) + 1.1) * 0.1 zero_points = np.array([0, 1]).astype(np.int32) axis = 1 min = 0 max = 255 torch_tensor = torch.tensor(np_array) torch_scales = torch.tensor(scales) torch_zero_points = torch.tensor(zero_points) torch_out1 = torch.fake_quantize_per_channel_affine(torch_tensor, torch_scales, torch_zero_points, axis, min, max) torch_out2 = torch.fake_quantize_per_tensor_affine(torch_tensor, 0.04, 1, min, max) ms_tensor = ms_torch.tensor(np_array) ms_scales = ms_torch.tensor(scales) ms_zero_points = ms_torch.tensor(zero_points) ms_out1 = ms_torch.fake_quantize_per_channel_affine(ms_tensor, ms_scales, ms_zero_points, axis, min, max) ms_out2 = ms_torch.fake_quantize_per_tensor_affine(ms_tensor, 0.04, 1, min, max) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert ms_out1.asnumpy().shape == torch_out1.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_fix(): np_array = np.random.rand(3, 4, 5) * 10 - 5 torch_tensor = torch.tensor(np_array) torch_out = torch.fix(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.fix(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_float_power(): np_array = np.random.rand(3, 4, 5).astype(np.float32)*5 np_exponent = np.random.rand(4, 5)*5 np_array = np_array.astype(np.float32) np_exponent = np_exponent.astype(np.int32) torch_tensor = torch.tensor(np_array) torch_exponent = torch.tensor(np_exponent) torch_out1 = torch.float_power(torch_tensor, torch_exponent) torch_out2 = torch.float_power(torch_tensor, 2+3j) torch_out3 = torch.float_power(2, torch_exponent) ms_tensor = ms_torch.tensor(np_array) ms_exponent = ms_torch.tensor(np_exponent) ms_out1 = ms_torch.float_power(ms_tensor, ms_exponent) ms_out2 = ms_torch.float_power(ms_tensor, 2+3j) ms_out3 = ms_torch.float_power(2, ms_exponent) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_floor_divide(): np_array = np.random.rand(3, 4, 5).astype(np.float32)*5 np_other = np.random.rand(4, 5)*5 np_array = np_array.astype(np.float32) np_other = np_other.astype(np.float32) torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out1 = torch.floor_divide(torch_tensor, torch_other) torch_out2 = torch.floor_divide(torch_tensor, 2) torch_out3 = torch.floor_divide(2, torch_other) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out1 = ms_torch.floor_divide(ms_tensor, ms_other) ms_out2 = ms_torch.floor_divide(ms_tensor, 2) ms_out3 = ms_torch.floor_divide(2, ms_other) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_frexp(): np_array = np.arange(16).reshape(1, 4, 4).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_out1, torch_out2 = torch.frexp(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out1, ms_out2 = ms_torch.frexp(ms_tensor) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_gradient(): np_array = np.array([4., 1., 1., 16.]) np_coordinates = np.array([-2., -1., 1., 4.]) torch_tensor = torch.tensor(np_array) torch_coordinates = torch.tensor(np_coordinates) torch_out1 = torch.gradient(torch_tensor, spacing=torch.tensor(1)) ms_tensor = ms_torch.tensor(np_array) ms_coordinates = ms_torch.tensor(np_coordinates) ms_out1 = ms_torch.gradient(ms_tensor, spacing=ms_torch.tensor(1)) assert ms_out1[0].asnumpy().dtype == torch_out1[0].numpy().dtype for i in range(len(ms_out1)): assert np.allclose(ms_out1[i].asnumpy(), torch_out1[i].numpy()) def test_imag(): np_array = np.array([[1+5j, 2-1j, 3.0, -4j, -0j]]) torch_tensor = torch.tensor(np_array) torch_out = torch.imag(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.imag(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_ldexp(): np_array = np.random.rand(1, 4, 5, 6)*5-2 np_other = np.random.rand(1, 4, 5, 6)*5-2 torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out = torch.ldexp(torch_tensor, torch_other) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out = ms_torch.ldexp(ms_tensor, ms_other) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_lerp(): np_array = np.random.rand(1, 4, 5, 5)*5-2 np_other = np.random.rand(4, 5, 5)*5-2 np_weight = np.random.rand(1, 4, 5, 5)*5-2 torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_weight = torch.tensor(np_weight) torch_out1 = torch.lerp(torch_tensor, torch_other, torch_weight) torch_out2 = torch.lerp(torch_tensor, torch_other, 3) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_weight = ms_torch.tensor(np_weight) ms_out1 = ms_torch.lerp(ms_tensor, ms_other, ms_weight) ms_out2 = ms_torch.lerp(ms_tensor, ms_other, 3) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), atol=1e-6) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert ms_out1.asnumpy().shape == torch_out1.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), atol=1e-6) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_logaddexp(): np_array = np.random.rand(1, 4, 5, 6)*5-2 np_other = np.random.rand(1, 4, 5, 6)*5-2 torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out = torch.logaddexp(torch_tensor, torch_other) torch_out2 = torch.logaddexp2(torch_tensor, torch_other) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out = ms_torch.logaddexp(ms_tensor, ms_other) ms_out2 = ms_torch.logaddexp2(ms_tensor, ms_other) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_logical_and(): np_array = np.random.randn(120).reshape(2, 2, 6, 5) * 16 - 8 np_array = np_array.astype(np.int16) np_other = np.random.randn(60).reshape(2, 6, 5) * 2 np_other = np_other.astype(np.bool8) torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out = torch.logical_and(torch_tensor, torch_other) torch_out2 = torch.logical_or(torch_tensor, torch.tensor(3)) torch_out3 = torch.logical_xor(torch.tensor(True), torch_other) torch_out4 = torch.logical_not(torch_tensor) torch_out5 = torch.logical_not(torch_other) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out = ms_torch.logical_and(ms_tensor, ms_other) ms_out2 = ms_torch.logical_or(ms_tensor, ms_torch.tensor(3)) ms_out3 = ms_torch.logical_xor(ms_torch.tensor(True), ms_other) ms_out4 = ms_torch.logical_not(ms_tensor) ms_out5 = ms_torch.logical_not(ms_other) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype assert np.allclose(ms_out4.asnumpy(), torch_out4.numpy()) assert ms_out4.asnumpy().dtype == torch_out4.numpy().dtype assert np.allclose(ms_out5.asnumpy(), torch_out5.numpy()) assert ms_out5.asnumpy().dtype == torch_out5.numpy().dtype def test_logit(): np_array = np.random.randn(120).reshape(4, 6, 5) * 3 - 1 np_array = np_array.astype(np.float32) eps = 0.001 torch_tensor = torch.tensor(np_array) torch_out1 = torch.logit(torch_tensor, 0.001) torch_out2 = torch.logit(torch_tensor, 0.4) torch_out3 = torch.logit(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.logit(ms_tensor, 0.001) ms_out2 = ms_torch.logit(ms_tensor, 0.4) ms_out3 = ms_torch.logit(ms_tensor) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), atol=1e-4) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), atol=1e-4, equal_nan=True) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy(), atol=1e-4, equal_nan=True) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_tanh(): np_array = np.random.rand(1, 1, 1, 1, 3, 4, 5).astype(np.float64) - 0.5 np_array = np_array * 20 torch_tensor = torch.tensor(np_array) torch_out = torch.tanh(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.tanh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_sigmoid(): np_array = np.random.rand(1, 1, 1, 1, 3, 4, 5).astype(np.float64) - 0.5 np_array = np_array * 20 torch_tensor = torch.tensor(np_array) torch_out = torch.sigmoid(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.sigmoid(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_hypot(): np_array = np.random.rand(2, 3, 4, 5) np_other = -np.random.rand(4, 5) torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out = torch.hypot(torch_tensor, torch_other) torch_out2 = torch.hypot(torch_tensor, torch.tensor(3)) ms_tensor = ms_torch.tensor(np_array) ms_other = ms_torch.tensor(np_other) ms_out = ms_torch.hypot(ms_tensor, ms_other) ms_out2 = ms_torch.hypot(ms_tensor, ms_torch.tensor(3)) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_i0(): np_array = np.random.rand(2, 3, 4, 5).astype(np.double) np_array2 = np.arange(12).reshape(3, 4) - 6 torch_tensor = torch.tensor(np_array) torch_tensor2 = torch.tensor(np_array2) torch_out = torch.i0(torch_tensor) torch_out2 = torch.i0(torch_tensor2) ms_tensor = ms_torch.tensor(np_array) ms_tensor2 = ms_torch.tensor(np_array2) ms_out = ms_torch.i0(ms_tensor) ms_out2 = ms_torch.i0(ms_tensor2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_igamma(): np_array = np.random.rand(2, 3, 4, 5) * 3 + 1 np_other = np.random.rand(4, 5) * 3 + 1 #np.int16, np.int64, np.bool8, np.half, np.single, np.double for type1 in (np.int64, np.bool8, np.half, np.double): for type2 in (np.int64, np.bool8, np.half, np.double): np_array_ = np_array.astype(type1) np_other_ = np_other.astype(type2) try: torch_tensor = torch.tensor(np_array_) torch_other = torch.tensor(np_other_) torch_out = torch.igamma(torch_tensor, torch_other) torch_out2 = torch.igammac(torch_tensor, torch_other) #print(torch_tensor.dtype, torch_other.dtype, torch_out.dtype) except: continue ms_tensor = ms_torch.tensor(np_array_) ms_other = ms_torch.tensor(np_other_) ms_out = ms_torch.igamma(ms_tensor, ms_other) ms_out2 = ms_torch.igammac(ms_tensor, ms_other) if ms_out.dtype == ms.float16: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True, atol=1e-3) assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), equal_nan=True, atol=1e-3) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_mvlgamma(): np_array = np.random.rand(2, 3, 4, 5).astype(np.double) for p in (1, 2, 3): torch_tensor = torch.tensor(np_array) + (p-1)/2 ms_tensor = ms_torch.tensor(np_array) + (p-1)/2 torch_out = torch.mvlgamma(torch_tensor, p) ms_out = ms_torch.mvlgamma(ms_tensor, p) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_nan_to_num(): np_array = np.array([np.nan, np.inf, -np.inf, 3]) typetest = [np.half, np.float32, np.double] for i in typetest: np_array_ = np_array.astype(i) torch_tensor = torch.tensor(np_array_) ms_tensor = ms_torch.tensor(np_array_) torch_out = torch.nan_to_num(torch_tensor) ms_out = ms_torch.nan_to_num(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype torch_tensor = torch.tensor(np_array) ms_tensor = ms_torch.tensor(np_array) torch_out = torch.nan_to_num(torch_tensor, 1.1, 2, -1) ms_out = ms_torch.nan_to_num(ms_tensor, 1.1, 2, -1) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_neg(): np_array1 = np.array([np.nan, np.inf, -np.inf, 3, 0, -0]) np_array2 = np.array([-0j, 0j, -0.0, -2+3j, 5+2j, 1, 1j]) np_array3 = np.array([0, 1, 2, 3, 4, 5]).astype(np.uint8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_tensor3 = torch.tensor(np_array3) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_tensor3 = ms_torch.tensor(np_array3) torch_out = torch.neg(torch_tensor1) ms_out = ms_torch.neg(ms_tensor1) torch_out2 = torch.negative(torch_tensor2) ms_out2 = ms_torch.negative(ms_tensor2) torch_out3 = torch.negative(torch_tensor3) ms_out3 = ms_torch.negative(ms_tensor3) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_nextafter(): np_array1 = np.array([np.nan, np.inf, -np.inf, 3, 0, -0]).astype(np.double) np_other = np.array([1, 2, 3, 3, -5, 6]).astype(np.double) torch_tensor1 = torch.tensor(np_array1) torch_other = torch.tensor(np_other) ms_tensor1 = ms_torch.tensor(np_array1) ms_other = ms_torch.tensor(np_other) torch_out = torch.nextafter(torch_tensor1, torch_other) ms_out = ms_torch.nextafter(ms_tensor1, ms_other) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_positive(): np_array1 = np.array([np.nan, np.inf, -np.inf, 3, 0, -0]) np_array2 = np.array([-0j, 0j, -0.0, -2+3j, 5+2j, 1, 1j]) np_array3 = np.array([0, 1, 2, 3, 4, 5]).astype(np.uint8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_tensor3 = torch.tensor(np_array3) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_tensor3 = ms_torch.tensor(np_array3) torch_out = torch.positive(torch_tensor1) ms_out = ms_torch.positive(ms_tensor1) torch_out2 = torch.positive(torch_tensor2) ms_out2 = ms_torch.positive(ms_tensor2) torch_out3 = torch.positive(torch_tensor3) ms_out3 = ms_torch.positive(ms_tensor3) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype ms_out[0] = 7 torch_out[0] = 7 assert np.allclose(ms_tensor1.asnumpy(), torch_tensor1.numpy()) assert ms_tensor1.asnumpy().dtype == torch_tensor1.numpy().dtype def test_rad2deg(): np_array = np.array([[3.14, 1.57, 0, -1.57, -1, 20]]) torch_tensor = torch.tensor(np_array) torch_out = torch.rad2deg(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.rad2deg(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_real(): np_array = np.array([[1+5j, 2-1j, 3.0, -4j, -0j]]) torch_tensor = torch.tensor(np_array) torch_out = torch.real(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.real(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_reciprocal(): np_array1 = np.array([[1+5j, 2-1j, 3.0, -4j, -0.1j]]) #np_array1 = np.array([[0j, -0j, 0, -0, -0-0j, 0+0j]]) #np_array2 = np.array([[np.inf, -0.0, 1.0, np.nan]]) np_array2 = np.array([[np.inf, 0.0, 1.0, np.nan]]) np_array3 = np.array([[1, 1, 2]]).astype(np.int16) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_tensor3 = torch.tensor(np_array3) torch_out1 = torch.reciprocal(torch_tensor1) torch_out2 = torch.reciprocal(torch_tensor2) torch_out3 = torch.reciprocal(torch_tensor3) ms_tensor1 = ms_torch.tensor(np_array1) ms_out1 = ms_torch.reciprocal(ms_tensor1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out2 = ms_torch.reciprocal(ms_tensor2) ms_tensor3 = ms_torch.tensor(np_array3) ms_out3 = ms_torch.reciprocal(ms_tensor3) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), equal_nan=True) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_remainder(): np_array1 = np.array([[-3., -2, -1, 1, 2, 3]]).astype(np.float16) np_array2 = np.array([[1, 2, 3, 4, 5]]).astype(np.int16) np_other = np.array([[6, 7, 3, 4, 5]]).astype(np.float32) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_other = torch.tensor(np_other) torch_out1 = torch.remainder(torch_tensor1, 2) torch_out2 = torch.remainder(torch_tensor2, 3) torch_out3 = torch.remainder(torch_tensor2, torch_other) ms_tensor1 = ms_torch.tensor(np_array1) ms_out1 = ms_torch.remainder(ms_tensor1, 2) ms_tensor2 = ms_torch.tensor(np_array2) ms_out2 = ms_torch.remainder(ms_tensor2, 3) ms_other = ms_torch.tensor(np_other) ms_out3 = ms_torch.remainder(ms_tensor2, ms_other) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_rsqrt(): np_array1 = np.array([[-3, -2, -1, 0, 2, 3]]).astype(np.float32) np_array2 = np.array([[1, 2, 3, 4, 5]]).astype(np.bool8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.rsqrt(torch_tensor1) torch_out2 = torch.rsqrt(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.rsqrt(ms_tensor1) ms_out2 = ms_torch.rsqrt(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), equal_nan=True) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_sgn(): np_array1 = np.array([[-3, -2, -0.0, 0.0, 2, 3]]).astype(np.float16) np_array2 = np.array([[-3, -2, -0, 0, 2, 3]]).astype(np.int16) np_array3 = np.array([[1+1j, 2, -0j, 0j, -0, 3-4j, -2+2j, -5j]]) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_tensor3 = torch.tensor(np_array3) torch_out1 = torch.sgn(torch_tensor1) torch_out2 = torch.sgn(torch_tensor2) torch_out3 = torch.sgn(torch_tensor3) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_tensor3 = ms_torch.tensor(np_array3) ms_out1 = ms_torch.sgn(ms_tensor1) ms_out2 = ms_torch.sgn(ms_tensor2) ms_out3 = ms_torch.sgn(ms_tensor3) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), equal_nan=True) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_signbit(): np_array1 = np.array([[-3, -2, -0.0, 0.0, 2, 3]]).astype(np.float16) np_array2 = np.array([[-3, -2, -0, 0, 2, 3]]).astype(np.int16) np_array3 = np.array([[-3, -2, -0, 0, 2, 3]]).astype(np.bool8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_tensor3 = torch.tensor(np_array3) torch_out1 = torch.signbit(torch_tensor1) torch_out2 = torch.signbit(torch_tensor2) torch_out3 = torch.signbit(torch_tensor3) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_tensor3 = ms_torch.tensor(np_array3) ms_out1 = ms_torch.signbit(ms_tensor1) ms_out2 = ms_torch.signbit(ms_tensor2) ms_out3 = ms_torch.signbit(ms_tensor3) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), equal_nan=True) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_sinc(): np_array1 = np.random.randn(10).astype(np.float32) np_array2 = np.random.randn(10).astype(np.float64) np_array1[0] = 0 torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.sinc(torch_tensor1) torch_out2 = torch.sinc(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.sinc(ms_tensor1) ms_out2 = ms_torch.sinc(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), atol=1e-6) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_sinh(): np_array1 = np.random.randn(10).astype(np.float32) np_array2 = np.random.randn(10).astype(np.float64)*10 np_array1[0] = 0 torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.sinh(torch_tensor1) torch_out2 = torch.sinh(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.sinh(ms_tensor1) ms_out2 = ms_torch.sinh(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_square(): np_array1 = np.random.randn(10).astype(np.complex64) np_array2 = np.arange(10).astype(np.int32) np_array1[0] = 1j torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.square(torch_tensor1) torch_out2 = torch.square(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.square(ms_tensor1) ms_out2 = ms_torch.square(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_sub(): np_x = np.random.rand(1, 2, 3, 2).astype(np.float16) np_y = np.random.rand(1, 2, 3, 2).astype(np.float64) for x_dtype in (np.int16, np.int32, np.uint8, np.half, np.double, np.single): for y_dtype in (np.int16, np.int32, np.uint8, np.half, np.double, np.single): np_x_ = np_x.astype(x_dtype) np_y_ = np_y.astype(y_dtype) torch_x = torch.tensor(np_x_) torch_y = torch.tensor(np_y_) torch_out1 = torch.sub(torch_x, 2.1, alpha=2) torch_out2 = torch.sub(torch_x, torch_y, alpha=2) #torch_out2 = torch.sub(torch_x, 2, torch_y) ms_x = ms_torch.tensor(np_x_) ms_y = ms_torch.tensor(np_y_) ms_out1 = ms_torch.sub(ms_x, 2.1, alpha=2) ms_out2 = ms_torch.sub(ms_x, ms_y, alpha=2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), equal_nan=True) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert ms_out1.asnumpy().shape == torch_out1.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), equal_nan=True) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape def test_tan(): np_array1 = np.random.randn(10).astype(np.float64) * 5 np_array2 = np.random.randn(10).astype(np.float32) * 5 np_array1[0] = 1.57 torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.tan(torch_tensor1) torch_out2 = torch.tan(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.tan(ms_tensor1) ms_out2 = ms_torch.tan(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_true_divide(): np_array1 = np.random.randn(10).astype(np.float64) * 5 np_array2 = np.random.randn(10).astype(np.float32) * 5 for x_dtype in (np.int16, np.int32, np.uint8, np.half, np.double, np.single): np_array1_ = np_array1.astype(x_dtype) torch_tensor1 = torch.tensor(np_array1_) ms_tensor1 = ms_torch.tensor(np_array1_) for y_dtype in (np.int16, np.int32, np.uint8, np.half, np.double, np.single): np_array2_ = np_array2.astype(y_dtype) torch_tensor2 = torch.tensor(np_array2_) ms_tensor2 = ms_torch.tensor(np_array2_) torch_out1 = torch.true_divide(torch_tensor1, torch_tensor2) ms_out1 = ms_torch.true_divide(ms_tensor1, ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), equal_nan=True) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype torch_out2 = torch.true_divide(torch_tensor1, 2) torch_out3 = torch.true_divide(torch_tensor1, 2.0) torch_out4 = torch.true_divide(2.0, torch_tensor1) ms_out2 = ms_torch.true_divide(ms_tensor1, 2) ms_out3 = ms_torch.true_divide(ms_tensor1, 2.0) ms_out4 = ms_torch.true_divide(2.0, ms_tensor1) assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), equal_nan=True) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy(), equal_nan=True) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype assert np.allclose(ms_out4.asnumpy(), torch_out4.numpy(), equal_nan=True) assert ms_out4.asnumpy().dtype == torch_out4.numpy().dtype torch_out5 = torch.true_divide(2, 2) ms_out5 = ms_torch.true_divide(2, 2) assert np.allclose(ms_out5.asnumpy(), torch_out5.numpy()) assert ms_out5.asnumpy().dtype == torch_out5.numpy().dtype def test_trunc(): np_array1 = np.random.randn(10).astype(np.float64) * 5 np_array2 = np.random.randn(10).astype(np.float32) * 5 torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.trunc(torch_tensor1) torch_out2 = torch.trunc(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.trunc(ms_tensor1) ms_out2 = ms_torch.trunc(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_xlogy(): np_array1 = np.random.randn(10).astype(np.float64) * 5 np_array2 = np.random.randn(10).astype(np.float32) * 5 np_array1[0:2] = 0 np_array2[1:3] = np.nan np_array2[5] = np.inf np_array2[0] = 2 torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch.xlogy(torch_tensor1, torch_tensor2) torch_out2 = torch.xlogy(2, torch_tensor2) torch_out3 = torch.xlogy(torch_tensor1, 3.0) ms_tensor1 = ms_torch.tensor(np_array1) ms_tensor2 = ms_torch.tensor(np_array2) ms_out1 = ms_torch.xlogy(ms_tensor1, ms_tensor2) ms_out2 = ms_torch.xlogy(2, ms_tensor2) ms_out3 = ms_torch.xlogy(ms_tensor1, 3.0) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), equal_nan=True) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), equal_nan=True) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy(), equal_nan=True) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype if __name__ == '__main__': test_round1() test_round2() test_round3() test_floor1() test_ceil1() test_sign1() test_sign2() test_pow1() test_pow2() #test_pow3() test_pow4() test_exp1() test_exp2() test_exp3() test_mul1() test_mul2() test_absolute() test_acos() test_arccos() test_acosh() test_arccosh() test_add() test_addcdiv() test_addcmul() test_angle() test_asin() test_arcsin() test_asinh() test_arcsinh() test_atan() test_arctan() test_atanh() test_arctanh() test_arctan2() test_bitwise_not() test_bitwise() test_bitwise_shift() test_clip() test_conj_physical() test_copysign() test_cosh() test_deg2rad() test_digamma() test_lgamma() test_erf() test_erfc() test_erfinv() test_exp2_() test_expm1() test_fake_quantize() test_fix() test_float_power() test_floor_divide() test_frexp() test_gradient() test_imag() test_ldexp() test_lerp() test_logaddexp() test_logical_and() test_logit() test_tanh() test_sigmoid() test_hypot() test_i0() test_igamma() test_mvlgamma() test_nan_to_num() test_neg() test_nextafter() test_positive() test_rad2deg() test_real() test_reciprocal() test_remainder() test_rsqrt() test_sgn() test_signbit() test_sinc() test_square() test_sub() test_tan() test_true_divide() test_trunc() test_xlogy() msadapter/testing/ut/pytorch/functional/test_meshgrid.py000066400000000000000000000057301436713554100242260ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_meshgrid1(): np_1 = np.array([1, 2, 3]) np_2 = np.array([5, 6, 7, 8]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.meshgrid(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.meshgrid(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_result[0].asnumpy(), torch_result[0].numpy()) assert ms_result[0].asnumpy().dtype == torch_result[0].numpy().dtype assert np.allclose(ms_result[1].asnumpy(), torch_result[1].numpy()) assert ms_result[1].asnumpy().dtype == torch_result[1].numpy().dtype def test_meshgrid2(): np_1 = np.array([1, 2, 3]) np_2 = np.array([5, 6, 7, 8]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.meshgrid([ms_tensor_1, ms_tensor_2]) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.meshgrid([torch_tensor_1, torch_tensor_2]) assert np.allclose(ms_result[0].asnumpy(), torch_result[0].numpy()) assert ms_result[0].asnumpy().dtype == torch_result[0].numpy().dtype assert np.allclose(ms_result[1].asnumpy(), torch_result[1].numpy()) assert ms_result[1].asnumpy().dtype == torch_result[1].numpy().dtype def test_meshgrid3(): np_1 = np.array([1, 2, 3]) np_2 = np.array([5, 6, 7, 8]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.meshgrid((ms_tensor_1, ms_tensor_2)) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.meshgrid((torch_tensor_1, torch_tensor_2)) assert np.allclose(ms_result[0].asnumpy(), torch_result[0].numpy()) assert ms_result[0].asnumpy().dtype == torch_result[0].numpy().dtype assert np.allclose(ms_result[1].asnumpy(), torch_result[1].numpy()) assert ms_result[1].asnumpy().dtype == torch_result[1].numpy().dtype def test_meshgrid4(): np_1 = np.array([1, 2, 3]) np_2 = np.array([5, 6, 7, 8]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.meshgrid(ms_tensor_1, ms_tensor_2, indexing='xy') torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.meshgrid(torch_tensor_1, torch_tensor_2, indexing='xy') assert np.allclose(ms_result[0].asnumpy(), torch_result[0].numpy()) assert ms_result[0].asnumpy().dtype == torch_result[0].numpy().dtype assert np.allclose(ms_result[1].asnumpy(), torch_result[1].numpy()) assert ms_result[1].asnumpy().dtype == torch_result[1].numpy().dtype if __name__ == '__main__': test_meshgrid1() test_meshgrid2() test_meshgrid3() test_meshgrid4() msadapter/testing/ut/pytorch/functional/test_mm.py000066400000000000000000000023721436713554100230340ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_mm1(): np_1 = np.array([[1, 2],[3, 4]]) np_2 = np.array([[1, 2],[3, 4]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.mm(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.mm(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_mm2(): np_1 = np.array([[1., 2],[3, 4]]) np_2 = np.array([[1, 2.],[3, 4]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.mm(ms_tensor_1, ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.mm(torch_tensor_1, torch_tensor_2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_mm1() test_mm2() msadapter/testing/ut/pytorch/functional/test_other.py000066400000000000000000000143041436713554100235420ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) #context.set_context(mode=ms.GRAPH_MODE) def test_view_as_real1(): np_array = np.array([[[[1+2j, 3+4j, 5+6j]],[[10+20j, 30+40j, 50+60j]]]]) torch_tensor = torch.tensor(np_array) torch_out = torch.view_as_real(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.view_as_real(ms_tensor) torch_out[0, 0, 0, 0, 0] = 0 ms_out[0, 0, 0, 0, 0] = 0 # assert np.allclose(torch_tensor.asnumpy(), ms_tensor.numpy()) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_view_as_real2(): np_array = np.array([-2.1, 0, 1.0j]) torch_tensor = torch.tensor(np_array) torch_out = torch.view_as_real(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.view_as_real(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_bincount1(): np_array = np.random.randint(0, 8, (7,)).astype(np.uint8) weight_array = np.linspace(0, 2, len(np_array)) torch_tensor = torch.tensor(np_array) torch_weight = torch.tensor(weight_array) torch_out = torch.bincount(torch_tensor, torch_weight) ms_tensor = ms_torch.tensor(np_array) ms_weight = ms_torch.tensor(weight_array) ms_out = ms_torch.bincount(ms_tensor, ms_weight) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_bincount2(): np_array = np.random.randint(0, 4, (8,)).astype(np.int64) torch_tensor = torch.tensor(np_array) torch_out = torch.bincount(torch_tensor, minlength=5) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.bincount(ms_tensor, minlength=5) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_nonzero1(): np_array = np.random.random((2,3,4)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.nonzero(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.nonzero(ms_tensor) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_nonzero2(): np_array = np.random.random((2,3,4)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.nonzero(torch_tensor, as_tuple=True) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.nonzero(ms_tensor, as_tuple=True) for i in range(len(ms_out)): assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) assert ms_out[i].numpy().dtype == torch_out[i].numpy().dtype def test_nonzero3(): np_array = np.array([1, 1, 1, 0, 1]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.nonzero(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.nonzero(ms_tensor) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_nonzero4(): np_array = np.array([1, 1, 1, 0, 1]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.nonzero(torch_tensor, as_tuple=True) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.nonzero(ms_tensor, as_tuple=True) for i in range(len(ms_out)): assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) assert ms_out[i].numpy().dtype == torch_out[i].numpy().dtype def test_frombuffer(): import array a = array.array('i', [1, 2, 3]) torch_tensor = torch.frombuffer(a, dtype=torch.int32) ms_tensor = ms_torch.frombuffer(a, dtype=ms_torch.int32) assert np.allclose(torch_tensor.numpy(), ms_tensor.numpy()) assert torch_tensor.numpy().dtype == ms_tensor.numpy().dtype def as_strided1(): np_array = np.random.random((3,3)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.as_strided(torch_tensor, (2, 2), (1, 2)) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.as_strided(ms_tensor, (2, 2), (1, 2)) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_unique1(): np_array = np.array([1, 3, 2, 3]) torch_tensor = torch.tensor(np_array) torch_out = torch.unique(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.unique(ms_tensor) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype # def test_unique2(): # np_array = np.array([1, 3, 2, 3]) # torch_tensor = torch.tensor(np_array) # torch_out = torch.unique(torch_tensor, sorted=False) # # ms_tensor = ms_torch.tensor(np_array) # ms_out = ms_torch.unique(ms_tensor, sorted=False) # # assert np.allclose(ms_out.numpy(), torch_out.numpy()) # assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_unique3(): np_array = np.array([[1, 3], [2, 3]]) torch_tensor = torch.tensor(np_array) torch_out = torch.unique(torch_tensor, sorted=True) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.unique(ms_tensor, sorted=True) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype # def test_unique4(): # np_array = np.array([[1, 3], [2, 3]]) # torch_tensor = torch.tensor(np_array) # torch_out = torch.unique(torch_tensor, sorted=False, return_inverse=True) # # ms_tensor = ms_torch.tensor(np_array) # ms_out = ms_torch.unique(ms_tensor, sorted=False, return_inverse=True) # # for i in range(len(torch_out)): # assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) if __name__ == '__main__': test_view_as_real1() test_view_as_real2() test_bincount1() test_bincount2() test_nonzero1() test_nonzero2() test_nonzero3() test_nonzero4() test_frombuffer() test_unique1() # test_unique2() test_unique3() # test_unique4() msadapter/testing/ut/pytorch/functional/test_permute.py000066400000000000000000000011301436713554100240730ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_permute(): tensor = np.random.random((3, 4, 4)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.permute(torch_tensor, (1, 0, 2)) ms_output = ms_torch.permute(ms_tensor, (1, 0, 2)) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) if __name__ == '__main__': test_permute()msadapter/testing/ut/pytorch/functional/test_reduction.py000066400000000000000000000251761436713554100244260ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context from ms_adapter.pytorch.nn import Module context.set_context(mode=ms.PYNATIVE_MODE) #context.set_context(mode=ms.GRAPH_MODE) def test_max(): np_array = np.array([[1, 2],[3, 4]]).astype(np.int32) ms_tensor = ms_torch.tensor(np_array) class MaxMs(Module): def __init__(self, dim=None, keepdim=False, out=None): super(MaxMs, self).__init__() self.max = ms_torch.max self.dim = dim self.keepdim = keepdim self.out =out def forward(self, input): x = self.max(input, dim=self.dim, keepdim=self.keepdim, out=self.out) return x ms_max1 = MaxMs() ms_max2 = MaxMs(dim=1, keepdim=True) ms_max3 = MaxMs(dim=0) ms_out1 = ms_max1(ms_tensor) ms_out2 = ms_max2(ms_tensor) ms_out3 = ms_max3(ms_tensor) #print(ms_out1,"\n",ms_out2,"\n",ms_out3) def test_max1(): np_array = np.array([[1, 2],[3, 4]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.max(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.max(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_max2(): np_array = np.array([[1, 2],[3, 4]]).astype(np.int32) torch_tensor = torch.tensor(np_array) torch_out = torch.max(torch_tensor, dim=1, keepdim=True) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.max(ms_tensor, dim=1, keepdim=True) assert np.allclose(ms_out.values.asnumpy(), torch_out.values.numpy()) assert np.allclose(ms_out.indices.asnumpy(), torch_out.indices.numpy()) assert ms_out.values.asnumpy().dtype == torch_out.values.numpy().dtype def test_max3(): np_array = np.array([1, 2, 3, 4]).astype(np.int32) torch_tensor = torch.tensor(np_array) torch_out = torch.max(torch_tensor, dim=0) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.max(ms_tensor, dim=0) assert np.allclose(ms_out.values.asnumpy(), torch_out.values.numpy()) assert np.allclose(ms_out.indices.asnumpy(), torch_out.indices.numpy()) assert ms_out.values.asnumpy().dtype == torch_out.values.numpy().dtype def test_mean_sum1(): np_array = np.array([[1, 2],[9, 8]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out1 = torch.mean(torch_tensor, dim=0, keepdim=True, dtype=torch.float16) torch_out2 = torch.sum(torch_tensor, dim=0, keepdim=True, dtype=torch.float16) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.mean(ms_tensor, dim=0, keepdim=True, dtype=ms_torch.float16) ms_out2 = ms_torch.sum(ms_tensor, dim=0, keepdim=True, dtype=ms_torch.float16) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_mean_sum2(): np_array = np.array([[1, 2],[9, 8]]).astype(np.int32) torch_tensor = torch.tensor(np_array) torch_out1 = torch.mean(torch_tensor, dim=0, keepdim=True, dtype=float) torch_out2 = torch.sum(torch_tensor, dim=0, keepdim=True, dtype=float) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.mean(ms_tensor, dim=0, keepdim=True, dtype=float) ms_out2 = ms_torch.sum(ms_tensor, dim=0, keepdim=True, dtype=float) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_mean_sum3(): np_array = np.array([1, 2, 9, 8]).astype(np.double) torch_tensor = torch.tensor(np_array) torch_out1 = torch.mean(torch_tensor) torch_out2 = torch.sum(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.mean(ms_tensor) ms_out2 = ms_torch.sum(ms_tensor) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_mean_sum4(): np_array = np.arange(0, 144).reshape(3, 3, 4, 4).astype(np.float32).astype(np.half) dim = (0, 2, 3) torch_tensor = torch.tensor(np_array) torch_out1 = torch.mean(torch_tensor, dim) torch_out2 = torch.sum(torch_tensor, dim) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.mean(ms_tensor, dim) ms_out2 = ms_torch.sum(ms_tensor, dim) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_mean_sum5(): np_array = np.array([2.0045, 2.0045, 2.0045]).astype(np.double) torch_tensor = torch.tensor(np_array) torch_out1 = torch.mean(torch_tensor, dtype=torch.float16) torch_out2 = torch.sum(torch_tensor, dtype=torch.float16) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.mean(ms_tensor, dtype=ms_torch.float16) ms_out2 = ms_torch.sum(ms_tensor, dtype=ms_torch.float16) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_median1(): np_array = np.arange(0, 36).astype(np.double) torch_tensor = torch.tensor(np_array) torch_out = torch.median(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.median(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_median2(): np_array = np.arange(0, 36).astype(np.double) torch_tensor = torch.tensor(np_array) torch_out = torch.median(torch_tensor, dim=0) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.median(ms_tensor, dim=0) assert np.allclose(ms_out.values.asnumpy(), torch_out.values.numpy()) assert np.allclose(ms_out.indices.asnumpy(), torch_out.indices.numpy()) assert ms_out.values.asnumpy().dtype == torch_out.values.numpy().dtype def test_median3(): np_array = np.arange(0, 36).reshape(2, 3, 6).astype(np.double) torch_tensor = torch.tensor(np_array) torch_out = torch.median(torch_tensor, dim=1, keepdim=True) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.median(ms_tensor, dim=1, keepdim=True) assert np.allclose(ms_out.values.asnumpy(), torch_out.values.numpy()) assert np.allclose(ms_out.indices.asnumpy(), torch_out.indices.numpy()) assert ms_out.values.asnumpy().dtype == torch_out.values.numpy().dtype def test_median4(): np_array = np.arange(0, 36).reshape(2, 3, 6).astype(np.int16) torch_tensor = torch.tensor(np_array) torch_out = torch.median(torch_tensor, dim=-1) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.median(ms_tensor, dim=-1) assert np.allclose(ms_out.values.asnumpy(), torch_out.values.numpy()) assert np.allclose(ms_out.indices.asnumpy(), torch_out.indices.numpy()) assert ms_out.values.asnumpy().dtype == torch_out.values.numpy().dtype def test_min1(): np_array = np.array([[1, 2],[3, 4]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.min(torch_tensor) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.min(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_min2(): np_array = np.array([[1, 2],[3, 4]]).astype(np.int32) torch_tensor = torch.tensor(np_array) torch_out, torch_indices = torch.min(torch_tensor, dim=1) ms_tensor = ms_torch.tensor(np_array) ms_out, ms_indices = ms_torch.min(ms_tensor,dim=1) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_min3(): np_array = np.array([[1, 2],[3, 4]]).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out, torch_indices = torch.min(torch_tensor, dim=1, keepdim=True) ms_tensor = ms_torch.tensor(np_array) ms_out, ms_indices = ms_torch.min(ms_tensor, dim=1, keepdim=True) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_min4(): np_array = np.array([[1, 2],[4, 3]]).astype(np.float32) torch_out1 = torch.Tensor() torch_out2 = torch.LongTensor() torch_tensor = torch.tensor(np_array) torch_out, torch_indices = torch.min(torch_tensor, dim=1, keepdim=True, out=(torch_out1, torch_out2)) ms_out1 = ms_torch.FloatTensor(2, 1) ms_out2 = ms_torch.IntTensor(2, 1) ms_tensor = ms_torch.tensor(np_array) ms_torch.min(ms_tensor, dim=1, keepdim=True, out=(ms_out1, ms_out2)) assert np.allclose(ms_out1.asnumpy(), torch_out.numpy()) assert np.allclose(ms_out2.asnumpy(), torch_indices.numpy()) assert ms_out1.asnumpy().dtype == torch_out.numpy().dtype def test_prod(): np_array = np.random.randn(2, 3, 4) * 2 torch_tensor = torch.tensor(np_array) torch_out1 = torch.prod(torch_tensor) torch_out2 = torch.prod(torch_tensor, dtype=int) torch_out3 = torch.prod(torch_tensor, dtype=torch.float32, dim=1, keepdim=False) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.prod(ms_tensor) ms_out2 = ms_torch.prod(ms_tensor, dtype=int) ms_out3 = ms_torch.prod(ms_tensor, dtype=ms_torch.float32, dim=1, keepdim=False) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype if __name__ == '__main__': test_max() test_max1() test_max2() test_max3() test_mean_sum1() test_mean_sum2() test_mean_sum3() test_mean_sum4() test_mean_sum5() test_median1() test_median2() test_median3() test_median4() test_min1() test_min2() test_min3() test_min4() test_prod() msadapter/testing/ut/pytorch/functional/test_reshape.py000066400000000000000000000022061436713554100240460ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_reshape1(): input = np.array([1,2,3,4]) ms_tensor = ms_torch.tensor(input) ms_result = ms_torch.reshape(ms_tensor, [2,2]) torch_tensor = torch.tensor(input) torch_result = torch.reshape(torch_tensor, [2,2]) assert ms_result.asnumpy().shape == torch_result.numpy().shape assert np.allclose(ms_result.numpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_reshape2(): input = np.array([1,2,3,4]) ms_tensor = ms_torch.tensor(input) ms_result = ms_torch.reshape(ms_tensor, [-1,]) torch_tensor = torch.tensor(input) torch_result = torch.reshape(torch_tensor, [-1,]) assert ms_result.asnumpy().shape == torch_result.numpy().shape assert np.allclose(ms_result.numpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_reshape1() test_reshape2()msadapter/testing/ut/pytorch/functional/test_softmax.py000066400000000000000000000021341436713554100241000ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_softmax(): np_array = np.random.randn(4, 3, 3, 20) torch_tensor = torch.tensor(np_array) torch_out1 = torch.softmax(torch_tensor, -1, torch.float32) torch_out2 = torch.softmax(torch_tensor, 3, float) torch_out3 = torch.softmax(torch_tensor, 2) ms_tensor = ms_torch.tensor(np_array) ms_out1 = ms_torch.softmax(ms_tensor, -1, ms_torch.float32) ms_out2 = ms_torch.softmax(ms_tensor, 3, float) ms_out3 = ms_torch.softmax(ms_tensor, 2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype if __name__ == '__main__': test_softmax() msadapter/testing/ut/pytorch/functional/test_sort.py000066400000000000000000000030571436713554100234130ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.PYNATIVE_MODE) #context.set_context(mode=ms.GRAPH_MODE) def test_sort(): x = [0, 1] * 9 torch_x = torch.tensor(x, dtype=torch.float32) torch_out1, torch_out2 = torch.sort(torch_x, dim=-1, descending=False, stable=True) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out1, ms_out2 = ms_torch.sort(ms_x) assert np.allclose(torch_out1.numpy(), ms_out1.numpy()) assert np.allclose(torch_out2.numpy(), ms_out2.numpy()) torch_out1, torch_out2 = torch.sort(torch_x, dim=0, descending=True, stable=True) ms_out1, ms_out2 = ms_torch.sort(ms_x, 0, True) assert np.allclose(torch_out1.numpy(), ms_out1.numpy()) assert np.allclose(torch_out2.numpy(), ms_out2.numpy()) def test_msort(): x = [0, 1] * 9 torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.msort(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.msort(ms_x) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_argsort(): x = [0, 1] * 9 torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.argsort(torch_x, -1, True) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.argsort(ms_x, -1, True) #assert np.allclose(torch_out.numpy(), ms_out.numpy()) # 'stable' is different if __name__ == '__main__': test_sort() test_msort() test_argsort()msadapter/testing/ut/pytorch/functional/test_split.py000066400000000000000000000053041436713554100235540ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_split1(): np_array = np.random.random((5,2)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.split(torch_tensor, split_size_or_sections=2) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.split(ms_tensor, split_size_or_sections=2) for i in range(len(ms_out)): assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) assert ms_out[i].numpy().dtype == torch_out[i].numpy().dtype def test_split2(): np_array = np.random.random((5,2)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.split(torch_tensor, split_size_or_sections=[1,4]) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.split(ms_tensor, split_size_or_sections=[1,4]) for i in range(len(ms_out)): assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) assert ms_out[i].numpy().dtype == torch_out[i].numpy().dtype def test_split3(): np_array = np.random.random((5,4)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.split(torch_tensor, split_size_or_sections=[1,3], dim=1) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.split(ms_tensor, split_size_or_sections=[1,3], dim=1) for i in range(len(ms_out)): assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) assert ms_out[i].numpy().dtype == torch_out[i].numpy().dtype def test_split4(): np_array = np.random.random((3,16,5,4)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.split(torch_tensor, split_size_or_sections=4, dim=1) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.split(ms_tensor, split_size_or_sections=4, dim=1) for i in range(len(ms_out)): assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) assert ms_out[i].numpy().dtype == torch_out[i].numpy().dtype def test_split5(): np_array = np.random.random((3,16,5,4)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch.split(torch_tensor, split_size_or_sections=[1,2,3,4,6], dim=1) ms_tensor = ms_torch.tensor(np_array) ms_out = ms_torch.split(ms_tensor, split_size_or_sections=[1,2,3,4,6], dim=1) for i in range(len(ms_out)): assert np.allclose(ms_out[i].numpy(), torch_out[i].numpy()) assert ms_out[i].numpy().dtype == torch_out[i].numpy().dtype if __name__ == '__main__': test_split1() test_split2() test_split3() test_split4() test_split5() msadapter/testing/ut/pytorch/functional/test_sqrt.py000066400000000000000000000025251436713554100234140ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_sqrt1(): np_1 = np.array([1, 2, 3]) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.sqrt(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.sqrt(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_sqrt2(): np_1 = np.array([1.0, 2.0, 3.0]) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.sqrt(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.sqrt(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_sqrt3(): np_1 = np.array([4, 9, 16]) ms_tensor_1 = ms_torch.tensor(np_1) ms_result = ms_torch.sqrt(ms_tensor_1) torch_tensor_1 = torch.tensor(np_1) torch_result = torch.sqrt(torch_tensor_1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_sqrt1() test_sqrt2() test_sqrt3() msadapter/testing/ut/pytorch/functional/test_squeeze.py000066400000000000000000000022561436713554100241050ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_squeeze1(): x = np.ones(shape=[3, 2, 1]) torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.squeeze(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.squeeze(ms_x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_squeeze2(): x = np.ones(shape=[3, 1, 2]) torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.squeeze(torch_x, dim=0) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.squeeze(ms_x, dim=0) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_squeeze3(): x = np.ones(shape=[3, 1, 2, 1]) torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.squeeze(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.squeeze(ms_x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) if __name__ == '__main__': test_squeeze1() test_squeeze2() test_squeeze3()msadapter/testing/ut/pytorch/functional/test_stack.py000066400000000000000000000024541436713554100235310ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_stack1(): np_1 = np.array([[1, 2], [3, 4]]) np_2 = np.array([[5, 6], [7, 8]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.stack((ms_tensor_1, ms_tensor_2), dim=1) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.stack((torch_tensor_1, torch_tensor_2), dim=1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_stack2(): np_1 = np.array([[1, 2], [3, 4]]) np_2 = np.array([[5, 6], [7, 8]]) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_result = ms_torch.stack((ms_tensor_1, ms_tensor_2)) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_result = torch.stack((torch_tensor_1, torch_tensor_2)) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_stack1() test_stack2() msadapter/testing/ut/pytorch/functional/test_t.py000066400000000000000000000014541436713554100226660ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_t1(): x = np.random.randn(3, 4) torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.t(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.t(ms_x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_t2(): x = np.random.randn(4) torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.t(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.t(ms_x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) if __name__ == '__main__': test_t1() test_t2()msadapter/testing/ut/pytorch/functional/test_transpose.py000066400000000000000000000034711436713554100244420ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_transpose1(): np_1 = np.random.randn(1, 2, 3) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.transpose(ms_tensor, 0, 1) torch_tensor = torch.tensor(np_1) torch_result = torch.transpose(torch_tensor, 0, 1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_transpose2(): np_1 = np.random.randn(1, 2, 3) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.transpose(ms_tensor, 0, -1) torch_tensor = torch.tensor(np_1) torch_result = torch.transpose(torch_tensor, 0, -1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_transpose3(): np_1 = np.random.randn(1, 2, 3) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.transpose(ms_tensor, 0, -3) torch_tensor = torch.tensor(np_1) torch_result = torch.transpose(torch_tensor, 0, -3) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_transpose4(): np_1 = np.random.randn(1, 2, 3) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.transpose(ms_tensor, -1, -2) torch_tensor = torch.tensor(np_1) torch_result = torch.transpose(torch_tensor, -1, -2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_transpose1() test_transpose2() test_transpose3() test_transpose4()msadapter/testing/ut/pytorch/functional/test_unbind.py000066400000000000000000000011601436713554100236740ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_unbind(): x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.unbind(torch_x, 0) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.unbind(ms_x, 0) for ms_out_, torch_out_ in zip(ms_out, torch_out): assert np.allclose(ms_out_.numpy(), torch_out_.numpy()) if __name__ == '__main__': test_unbind()msadapter/testing/ut/pytorch/functional/test_unsqueeze.py000066400000000000000000000011001436713554100244330ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_unsqueeze(): x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.unsqueeze(torch_x, 0) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.unsqueeze(ms_x, 0) assert np.allclose(ms_out.numpy(), torch_out.numpy()) if __name__ == '__main__': test_unsqueeze()msadapter/testing/ut/pytorch/functional/test_zeros.py000066400000000000000000000013601436713554100235610ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_zeros1(): ms_result = ms_torch.zeros(1, 2, 3) torch_result = torch.zeros(1, 2, 3) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_zeros2(): ms_result = ms_torch.zeros((1, 2, 3)) torch_result = torch.zeros((1, 2, 3)) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_zeros1() test_zeros2() msadapter/testing/ut/pytorch/nn/000077500000000000000000000000001436713554100172575ustar00rootroot00000000000000msadapter/testing/ut/pytorch/nn/__init__.py000066400000000000000000000000001436713554100213560ustar00rootroot00000000000000msadapter/testing/ut/pytorch/nn/functional/000077500000000000000000000000001436713554100214215ustar00rootroot00000000000000msadapter/testing/ut/pytorch/nn/functional/test_avg_pooling.py000066400000000000000000000171721436713554100253460ustar00rootroot00000000000000import numpy as np import torch import mindspore as ms import ms_adapter.pytorch as ms_torch ms.context.set_context(mode=ms.PYNATIVE_MODE) def test_avg_pool1d_1(): N = np.random.randint(1, 33) C = np.random.randint(2, 65) tensor = np.random.randn(N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool1d(torch_tensor, kernel_size, stride, padding) ms_output = ms_torch.nn.functional.avg_pool1d(ms_tensor, kernel_size, stride, padding) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # TODO: Not support ceil_mode yet ''' def test_avg_pool1d_2(): B = np.random.randint(1, 33) N = np.random.randint(1, 33) C = np.random.randint(2, 65) tensor = np.random.randn(B, N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool1d(torch_tensor, kernel_size, stride, padding, ceil_mode=True) ms_output = ms_torch.nn.functional.avg_pool1d(ms_tensor, kernel_size, stride, padding, ceil_mode=True) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) ''' def test_avg_pool1d_3(): B = np.random.randint(1, 33) N = np.random.randint(1, 33) C = np.random.randint(2, 65) tensor = np.random.randn(B, N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool1d(torch_tensor, kernel_size, stride, padding, count_include_pad=True) ms_output = ms_torch.nn.functional.avg_pool1d(ms_tensor, kernel_size, stride, padding, count_include_pad=True) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avg_pool3d_1(): N = np.random.randint(1, 17) C = np.random.randint(2, 33) D = np.random.randint(1, 17) H = np.random.randint(1, 33) W = np.random.randint(1, 33) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool3d(torch_tensor, kernel_size, stride, padding) ms_output = ms_torch.nn.functional.avg_pool3d(ms_tensor, kernel_size, stride, padding) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avg_pool3d_2(): N = np.random.randint(1, 17) C = np.random.randint(2, 33) D = np.random.randint(5, 17) H = np.random.randint(5, 33) W = np.random.randint(5, 33) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, (D - 1)/2), np.random.randint(1, (H - 1)/2), np.random.randint(1, (W - 1)/2)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool3d(torch_tensor, kernel_size, stride, padding, ceil_mode=True) ms_output = ms_torch.nn.functional.avg_pool3d(ms_tensor, kernel_size, stride, padding, ceil_mode=True) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avg_pool3d_3(): C = np.random.randint(2, 33) D = np.random.randint(5, 17) H = np.random.randint(5, 33) W = np.random.randint(5, 33) tensor = np.random.randn(C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, (D - 1)/2), np.random.randint(1, (H - 1)/2), np.random.randint(1, (W - 1)/2)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool3d(torch_tensor, kernel_size, stride, padding, count_include_pad=True) ms_output = ms_torch.nn.functional.avg_pool3d(ms_tensor, kernel_size, stride, padding, count_include_pad=True) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avg_pool3d_4(): C = np.random.randint(2, 33) D = np.random.randint(5, 17) H = np.random.randint(5, 33) W = np.random.randint(5, 33) tensor = np.random.randn(C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, (D - 1)/2), np.random.randint(1, (H - 1)/2), np.random.randint(1, (W - 1)/2)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool3d(torch_tensor, kernel_size, stride, padding, divisor_override=2) ms_output = ms_torch.nn.functional.avg_pool3d(ms_tensor, kernel_size, stride, padding, divisor_override=2) if ms.context.get_context("device_target") == "GPU": assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) elif ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) if __name__ == '__main__': test_avg_pool1d_1() test_avg_pool1d_2() test_avg_pool1d_3() test_avg_pool3d_1() test_avg_pool3d_2() test_avg_pool3d_3() test_avg_pool3d_4() msadapter/testing/ut/pytorch/nn/functional/test_batch_norm.py000066400000000000000000000161201436713554100251460ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_batch_norm1(): np_input = np.random.randn(8, 8, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.batch_norm( ms_input, ms_running_mean, ms_running_var, ms_weight, ms_bias, training=True) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.batch_norm( torch_input, torch_running_mean, torch_running_var, torch_weight, torch_bias, training=True) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_batch_norm2(): np_input = np.random.randn(4, 8, 6, 5).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_result = ms_torch.nn.functional.batch_norm( ms_input, ms_running_mean, ms_running_var, training=True) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_result = torch.nn.functional.batch_norm( torch_input, torch_running_mean, torch_running_var, training=True) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_batch_norm3(): np_input = np.random.randn(4, 8, 6, 5, 4).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_weight = ms_torch.tensor(weight) ms_result = ms_torch.nn.functional.batch_norm( ms_input, ms_running_mean, ms_running_var, weight=ms_weight, training=False) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_weight = torch.tensor(weight) torch_result = torch.nn.functional.batch_norm( torch_input, torch_running_mean, torch_running_var, weight=torch_weight, training=False) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_batch_norm4(): np_input = np.random.randn(10, 8, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.batch_norm( ms_input, ms_running_mean, ms_running_var, bias=ms_bias, training=True) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.batch_norm( torch_input, torch_running_mean, torch_running_var, bias=torch_bias, training=True) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_batch_norm5(): np_input = np.random.randn(10, 8, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_result = ms_torch.nn.functional.batch_norm( ms_input, ms_running_mean, ms_running_var, training=True) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_result = torch.nn.functional.batch_norm( torch_input, torch_running_mean, torch_running_var, training=True) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_batch_norm6(): np_input = np.random.randn(4, 8, 6, 5).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_result = ms_torch.nn.functional.batch_norm( ms_input, ms_running_mean, ms_running_var, training=True, momentum=0.01, eps=1e-4) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_result = torch.nn.functional.batch_norm( torch_input, torch_running_mean, torch_running_var, training=True, momentum=0.01, eps=1e-4) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape if __name__ == "__main__": test_batch_norm1() test_batch_norm2() test_batch_norm3() test_batch_norm4() test_batch_norm5() test_batch_norm6() msadapter/testing/ut/pytorch/nn/functional/test_conv.py000066400000000000000000000247311436713554100240060ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_conv1d1(): np_input = np.random.randn(33, 16, 30).astype(np.float32) np_weight = np.random.randn(20, 16, 5).astype(np.float32) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_out = torch.nn.functional.conv1d(torch_tensor, torch_weight) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_out = ms_torch.nn.functional.conv1d(ms_tensor, ms_weight) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv1d2(): out_channel = 1 np_input = np.random.randn(3, 8, 30).astype(np.float32) np_weight = np.random.randn(out_channel, 8, 5).astype(np.float32) np_bias = np.ones(out_channel).astype(np.float32)*0.5 torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_out = torch.nn.functional.conv1d(torch_tensor, torch_weight, torch_bias, 7, (3,), (4,), 1) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_out = ms_torch.nn.functional.conv1d(ms_tensor, ms_weight, ms_bias, 7, (3,), (4,), 1) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv1d3(): out_channel = 8 np_input = np.random.randn(15, 32, 30).astype(np.float64) np_weight = np.random.randn(out_channel, 8, 5).astype(np.float64) np_bias = np.ones(out_channel).astype(np.float32)*0.5 torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_out = torch.nn.functional.conv1d(torch_tensor, torch_weight, torch_bias, (1,), 'same', (2,), 4) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_out = ms_torch.nn.functional.conv1d(ms_tensor, ms_weight, ms_bias, (1,), 'same', (2,), 4) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv2d1(): np_input = np.random.randn(1, 4, 5, 5).astype(np.float32) np_weight = np.random.randn(8, 4, 3, 3).astype(np.float32) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_out = torch.nn.functional.conv2d(torch_tensor, torch_weight) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_out = ms_torch.nn.functional.conv2d(ms_tensor, ms_weight) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv2d2(): out_channel = 1 np_input = np.random.randn(3, 8, 9, 9).astype(np.float32) np_weight = np.random.randn(out_channel, 8, 3, 5).astype(np.float32) np_bias = np.ones(out_channel).astype(np.float32)*0.5 torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_out = torch.nn.functional.conv2d(torch_tensor, torch_weight, torch_bias, 7, 3, (2,), 1) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_out = ms_torch.nn.functional.conv2d(ms_tensor, ms_weight, ms_bias, 7, 3, (2,), 1) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv2d3(): out_channel = 8 np_input = np.random.randn(4, 16, 8, 8).astype(np.float64) np_weight = np.random.randn(out_channel, 4, 3, 3).astype(np.float64) np_bias = np.ones(out_channel).astype(np.float32)*0.5 torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_out = torch.nn.functional.conv2d(torch_tensor, torch_weight, torch_bias, (2, 2), (3, 2), (1, 3), 4) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_out = ms_torch.nn.functional.conv2d(ms_tensor, ms_weight, ms_bias, (2, 2), (3, 2), (1, 3), 4) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv3d1(): np_input = np.random.randn(1, 4, 5, 5, 5).astype(np.float32) np_weight = np.random.randn(8, 4, 3, 3, 3).astype(np.float32) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_out = torch.nn.functional.conv3d(torch_tensor, torch_weight) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_out = ms_torch.nn.functional.conv3d(ms_tensor, ms_weight) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv3d2(): out_channel = 1 np_input = np.random.randn(3, 8, 9, 9, 9).astype(np.float32) np_weight = np.random.randn(out_channel, 8, 3, 5, 4).astype(np.float32) np_bias = np.ones(out_channel).astype(np.float32)*0.5 torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_out = torch.nn.functional.conv3d(torch_tensor, torch_weight, torch_bias, 7, 3, (2,), 1) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_out = ms_torch.nn.functional.conv3d(ms_tensor, ms_weight, ms_bias, 7, 3, (2,), 1) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_conv3d3(): out_channel = 8 np_input = np.random.randn(4, 16, 18, 18, 18).astype(np.float64) np_weight = np.random.randn(out_channel, 4, 5, 5, 5).astype(np.float64) np_bias = np.ones(out_channel).astype(np.float32)*0.5 torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_out = torch.nn.functional.conv3d(torch_tensor, torch_weight, torch_bias, (2, 2, 1), (4, 6, 2), (1, 3, 3), 4) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_out = ms_torch.nn.functional.conv3d(ms_tensor, ms_weight, ms_bias, (2, 2, 1), (4, 6, 2), (1, 3, 3), 4) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-2, atol=1e-2) else: assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_unfold(): np_input = np.random.randn(7, 8, 9, 10) torch_tensor = torch.tensor(np_input) torch_out = torch.nn.functional.unfold(torch_tensor, (2, 3), 1, (1, 2), (2, 1)) ms_tensor = ms_torch.tensor(np_input) ms_out = ms_torch.nn.functional.unfold(ms_tensor, (2, 3), 1, (1, 2), (2, 1)) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_fold(): np_input1 = np.random.randn(7, 8, 24) np_input2 = np.random.randn(18, 6) torch_tensor1 = torch.tensor(np_input1) torch_tensor2 = torch.tensor(np_input2) torch_out1 = torch.nn.functional.fold(torch_tensor1, (4, 5), (2, 2), 1, (1, 2), (2, 1)) torch_out2 = torch.nn.functional.fold(torch_tensor2, (7, 4), 3, 2, 2, 3) ms_tensor1 = ms_torch.tensor(np_input1) ms_tensor2 = ms_torch.tensor(np_input2) ms_out1 = ms_torch.nn.functional.fold(ms_tensor1, (4, 5), (2, 2), 1, (1, 2), (2, 1)) ms_out2 = ms_torch.nn.functional.fold(ms_tensor2, (7, 4), 3, 2, 2, 3) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), rtol=1e-3, atol=1e-5) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert ms_out1.asnumpy().shape == torch_out1.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), rtol=1e-3, atol=1e-5) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape if __name__ == '__main__': test_conv1d1() test_conv1d2() test_conv1d3() test_conv2d1() test_conv2d2() test_conv2d3() test_conv3d1() test_conv3d2() test_conv3d3() test_unfold() test_fold() msadapter/testing/ut/pytorch/nn/functional/test_conv_transpose1d.py000066400000000000000000000122351436713554100263250ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import pytest import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_conv_transpose1d1(): np_input = np.random.randn(20, 16, 50).astype(np.float32) np_weight = np.random.randn(16, 33, 5).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose1d(ms_tensor, ms_weight) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose1d(torch_tensor, torch_weight) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose1d2(): np_input = np.random.randn(20, 15, 50).astype(np.float32) np_weight = np.random.randn(15, 33, 5).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose1d( ms_tensor, ms_weight, stride=2, padding=1, groups=3, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose1d( torch_tensor, torch_weight, stride=2, padding=1, groups=3, dilation=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose1d3(): np_input = np.random.randn(20, 15, 50).astype(np.float32) np_weight = np.random.randn(15, 33, 5).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose1d( ms_tensor, ms_weight, stride=2, padding=0, output_padding=0, groups=3, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose1d( torch_tensor, torch_weight, stride=2, padding=0, output_padding=0, groups=3, dilation=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape # def test_conv_transpose1d4(): # np_input = np.random.randn(1, 3, 50).astype(np.float32) # np_weight = np.random.randn(3, 64, 5).astype(np.float32) # # ms_tensor = ms_torch.tensor(np_input) # ms_weight = ms_torch.tensor(np_weight) # ms_result = ms_torch.nn.functional.conv_transpose1d( # ms_tensor, ms_weight, stride=3, padding=2, outputpadding=2, bias=None) # # torch_tensor = torch.tensor(np_input) # torch_weight = torch.tensor(np_weight) # torch_result = torch.nn.functional.conv_transpose1d( # torch_tensor, torch_weight, stride=3, padding=2, output_padding=2, bias=None) # assert np.allclose(ms_result.asnumpy().round(decimals=1e-4), torch_result.numpy(), rtol=1e-3, atol=1e-5) # assert ms_result.asnumpy().dtype == torch_result.numpy().dtype # assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose1d5(): np_input = np.random.randn(20, 16, 50).astype(np.float32) np_weight = np.random.randn(16, 33, 5).astype(np.float32) np_bias = np.random.randn(np_weight.shape[1]).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_result = ms_torch.nn.functional.conv_transpose1d(ms_tensor, ms_weight, bias=ms_bias) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_result = torch.nn.functional.conv_transpose1d(torch_tensor, torch_weight, bias=torch_bias) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose1d6(): np_input = np.random.randn(20, 15, 50, 10).astype(np.float32) np_weight = np.random.randn(15, 11, 5, 10).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) with pytest.raises(ValueError): ms_result = ms_torch.nn.functional.conv_transpose1d( ms_tensor, ms_weight, stride=2, padding=1, groups=3, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) with pytest.raises(RuntimeError): torch_result = torch.nn.functional.conv_transpose1d( torch_tensor, torch_weight, stride=2, padding=1, groups=3, dilation=2) if __name__ == "__main__": test_conv_transpose1d1() test_conv_transpose1d2() # reversed for pad_same # test_conv_transpose1d4() test_conv_transpose1d5() test_conv_transpose1d6() msadapter/testing/ut/pytorch/nn/functional/test_conv_transpose2d.py000066400000000000000000000121461436713554100263270ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import pytest import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_conv_transpose2d1(): np_input = np.random.randn(1, 4, 5, 5).astype(np.float32) np_weight = np.random.randn(4, 8, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose2d(ms_tensor, ms_weight) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose2d(torch_tensor, torch_weight) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose2d2(): np_input = np.random.randn(1, 4, 5, 5).astype(np.float32) np_weight = np.random.randn(4, 8, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose2d( ms_tensor, ms_weight, stride=2, padding=1, groups=2, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose2d( torch_tensor, torch_weight, stride=2, padding=1, groups=2, dilation=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose2d3(): np_input = np.random.randn(1, 4, 5, 5).astype(np.float32) np_weight = np.random.randn(4, 8, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose2d( ms_tensor, ms_weight, stride=2, padding=0, output_padding=0, groups=2, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose2d( torch_tensor, torch_weight, stride=2, padding=0, output_padding=0, groups=2, dilation=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose2d4(): np_input = np.random.randn(1, 4, 5, 5).astype(np.float32) np_weight = np.random.randn(4, 8, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose2d( ms_tensor, ms_weight, stride=2, padding=(2, 4), groups=2, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose2d( torch_tensor, torch_weight, stride=2, padding=(2, 4), groups=2, dilation=2) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose2d5(): np_input = np.random.randn(1, 4, 5, 5).astype(np.float32) np_weight = np.random.randn(4, 8, 3, 3).astype(np.float32) np_bias = np.random.randn(np_weight.shape[1]).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_bias = ms_torch.tensor(np_bias) ms_result = ms_torch.nn.functional.conv_transpose2d(ms_tensor, ms_weight, bias=ms_bias) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_bias = torch.tensor(np_bias) torch_result = torch.nn.functional.conv_transpose2d(torch_tensor, torch_weight, bias=torch_bias) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose2d6(): np_input = np.random.randn(20, 15, 50).astype(np.float32) np_weight = np.random.randn(15, 11, 5).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) with pytest.raises(ValueError): ms_result = ms_torch.nn.functional.conv_transpose2d( ms_tensor, ms_weight, stride=2, padding=1, groups=3, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) with pytest.raises(RuntimeError): torch_result = torch.nn.functional.conv_transpose2d( torch_tensor, torch_weight, stride=2, padding=1, groups=3, dilation=2) if __name__ == "__main__": test_conv_transpose2d1() test_conv_transpose2d2() test_conv_transpose2d3() test_conv_transpose2d4() test_conv_transpose2d5() test_conv_transpose2d6() msadapter/testing/ut/pytorch/nn/functional/test_conv_transpose3d.py000066400000000000000000000106521436713554100263300ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import pytest import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_conv_transpose3d1(): np_input = np.random.randn(20, 16, 50, 10, 20).astype(np.float32) np_weight = np.random.randn(16, 34, 3, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose3d(ms_tensor, ms_weight) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose3d(torch_tensor, torch_weight) if ms.context.get_context("device_target") == "GPU": assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-4) else: assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose3d2(): np_input = np.random.randn(20, 16, 50, 10, 20).astype(np.float32) np_weight = np.random.randn(16, 34, 3, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose3d( ms_tensor, ms_weight, stride=2, padding=1, groups=1, dilation=1) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose3d( torch_tensor, torch_weight, stride=2, padding=1, groups=1, dilation=1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose3d3(): np_input = np.random.randn(20, 16, 50, 10, 20).astype(np.float32) np_weight = np.random.randn(16, 34, 3, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose3d( ms_tensor, ms_weight, stride=2, padding=0, output_padding=0, groups=1, dilation=1) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose3d( torch_tensor, torch_weight, stride=2, padding=0, output_padding=0, groups=1, dilation=1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose3d4(): np_input = np.random.randn(20, 16, 50, 10, 20).astype(np.float32) np_weight = np.random.randn(16, 34, 3, 3, 3).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) ms_result = ms_torch.nn.functional.conv_transpose3d( ms_tensor, ms_weight, stride=2, padding=(1, 1, 2), groups=1, dilation=1) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) torch_result = torch.nn.functional.conv_transpose3d( torch_tensor, torch_weight, stride=2, padding=(1, 1, 2), groups=1, dilation=1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_conv_transpose3d5(): np_input = np.random.randn(20, 15, 50).astype(np.float32) np_weight = np.random.randn(15, 11, 5).astype(np.float32) ms_tensor = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(np_weight) with pytest.raises(ValueError): ms_result = ms_torch.nn.functional.conv_transpose3d( ms_tensor, ms_weight, stride=2, padding=1, groups=3, dilation=2) torch_tensor = torch.tensor(np_input) torch_weight = torch.tensor(np_weight) with pytest.raises(RuntimeError): torch_result = torch.nn.functional.conv_transpose3d( torch_tensor, torch_weight, stride=2, padding=1, groups=3, dilation=2) if __name__ == "__main__": test_conv_transpose3d1() test_conv_transpose3d2() test_conv_transpose3d3() test_conv_transpose3d4() test_conv_transpose3d5() msadapter/testing/ut/pytorch/nn/functional/test_dropout.py000066400000000000000000000114511436713554100245300ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_dropout1d1(): data = np.array([[-1, 0, 1, -2], [2, 2, 3, 4] ]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.dropout1d(ms_input, 0.8, False) torch_input = torch.tensor(data) torch_out = torch.nn.functional.dropout1d(torch_input, 0.8, False) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_dropout1d2(): data = np.array([[[-1, 0, 1, -2], [2, 2, 3, 4]], [[1, 1, 1, 1], [1, 1, 1, 1]]]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.dropout1d(ms_input, 0.7, True, True) torch_input = torch.tensor(data) torch_out = torch.nn.functional.dropout1d(torch_input, 0.7, True, True) # assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) # assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_dropout2d1(): data = np.ones([2, 3, 2, 4]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.dropout2d(ms_input, 0.6) torch_input = torch.tensor(data) torch_out = torch.nn.functional.dropout2d(torch_input, 0.6) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_dropout3d1(): data = np.ones([2, 3, 2, 4]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.dropout3d(ms_input, 0.6) torch_input = torch.tensor(data) torch_out = torch.nn.functional.dropout3d(torch_input, 0.6) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_dropout3d2(): data = np.ones([2, 2, 3, 2, 4]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.dropout3d(ms_input, 0.4) torch_input = torch.tensor(data) torch_out = torch.nn.functional.dropout3d(torch_input, 0.4) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_dropout(): data = np.ones([2, 3, 2, 4]).astype(np.float32)*2 ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.dropout(ms_input, 0.6) torch_input = torch.tensor(data) torch_out = torch.nn.functional.dropout(torch_input, 0.6) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_alphadropout1(): data = np.ones([2, 3, 2]).astype(np.float32) data[:,:,0]=-1*data[:,:,0] ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.alpha_dropout(ms_input, 0.1, True) torch_input = torch.tensor(data) torch_out = torch.nn.functional.alpha_dropout(torch_input, 0.1, True) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_alphadropout2(): data = np.array([[1.0, -3.0, 4.0],[-2.0, 8.5, 0]]) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.alpha_dropout(ms_input, 0.6, True) torch_input = torch.tensor(data) torch_out = torch.nn.functional.alpha_dropout(torch_input, 0.6, True) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_featurealphadropout1(): data = np.ones([2, 4, 3, 2]).astype(np.float32) data[:, :, :, 0] = -1 * data[:, :, :, 0] ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.feature_alpha_dropout(ms_input, 0.5, True) torch_input = torch.tensor(data) torch_out = torch.nn.functional.feature_alpha_dropout(torch_input, 0.5, True) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_featurealphadropout2(): data = np.ones([5, 6]).astype(np.float32) data[2:3, 0:2] = -1 ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.feature_alpha_dropout(ms_input, 0.35, True) torch_input = torch.tensor(data) torch_out = torch.nn.functional.feature_alpha_dropout(torch_input, 0.35, True) assert ms_out.shape == torch_out.shape assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_dropout1d1() test_dropout1d2() test_dropout2d1() test_dropout3d1() test_dropout3d2() test_dropout() test_alphadropout1() test_alphadropout2() test_featurealphadropout1() test_featurealphadropout2() msadapter/testing/ut/pytorch/nn/functional/test_functional.py000066400000000000000000000265141436713554100252040ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import ms_adapter.pytorch as ms_torch import numpy as np import mindspore as ms import torch from ms_adapter.pytorch.nn.functional import interpolate, adaptive_avg_pool2d, avg_pool2d ms.context.set_context(mode=ms.PYNATIVE_MODE) def test_interpolate1(): tensor = np.arange(1, 5).reshape((1, 1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.interpolate(torch_tensor, size=3, mode="bilinear", align_corners=True) ms_tensor = ms_torch.tensor(tensor) ms_output = interpolate(ms_tensor, size=3, mode="bilinear", align_corners=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_interpolate2(): tensor = np.arange(1, 5).reshape((1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.interpolate(torch_tensor, scale_factor=4, mode='linear', align_corners=True) ms_tensor = ms_torch.tensor(tensor) ms_output = interpolate(ms_tensor, scale_factor=4, mode='linear', align_corners=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_interpolate3(): tensor = np.arange(1, 5).reshape((1, 1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.interpolate(torch_tensor, size=3, mode="bilinear", align_corners=False) ms_tensor = ms_torch.tensor(tensor) ms_output = interpolate(ms_tensor, size=3, mode="bilinear", align_corners=False) def test_interpolate4(): tensor = np.arange(1, 5).reshape((1, 1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.interpolate(torch_tensor, size=3, mode="nearest") ms_tensor = ms_torch.tensor(tensor) ms_output = interpolate(ms_tensor, size=3, mode="nearest") assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_interpolate5(): tensor = np.arange(1, 5).reshape((1, 1, 1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.interpolate(torch_tensor, size=3, mode="nearest") ms_tensor = ms_torch.tensor(tensor) ms_output = interpolate(ms_tensor, size=3, mode="nearest") assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_adaptive_avg_pool2d(): tensor = np.random.randn(1, 32, 9, 9).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.adaptive_avg_pool2d(torch_tensor, (3, 5)) ms_tensor = ms_torch.tensor(tensor) ms_output = adaptive_avg_pool2d(ms_tensor, (3, 5)) assert (torch_output.shape == ms_output.shape) #TODO assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avg_pool2d(): tensor = np.random.randn(1, 32, 9, 9).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.avg_pool2d(torch_tensor, (3, 5), stride=None) ms_tensor = ms_torch.tensor(tensor) ms_output = avg_pool2d(ms_tensor, (3, 5), stride=None) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_upsample_nearest1(): data = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_nearest(torch_tensor, (3, 5)) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_nearest(ms_tensor, (3, 5)) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_nearest2(): data = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_nearest(torch_tensor, 3) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_nearest(ms_tensor, 3) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_nearest3(): data = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_nearest(torch_tensor, scale_factor=2) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_nearest(ms_tensor, scale_factor=2) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_nearest4(): data = np.random.randn(2, 3, 4, 5, 6).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_nearest(torch_tensor, (3, 4, 5)) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_nearest(ms_tensor, (3, 4, 5)) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_nearest5(): data = np.random.randn(2, 3, 4, 5, 6).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_nearest(torch_tensor, scale_factor=2) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_nearest(ms_tensor, scale_factor=2) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_nearest6(): data = np.random.randn(2, 3, 4, 5, 6).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_nearest(torch_tensor, 3) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_nearest(ms_tensor, 3) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_bilinear1(): data = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_bilinear(torch_tensor, (3, 5)) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_bilinear(ms_tensor, (3, 5)) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_bilinear2(): data = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_bilinear(torch_tensor, 3) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_bilinear(ms_tensor, 3) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_upsample_bilinear3(): data = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.upsample_bilinear(torch_tensor, scale_factor=2) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.upsample_bilinear(ms_tensor, scale_factor=2) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-4) def test_adaptive_avg_pool1d(): data = np.random.randint(0, 10, [1, 3, 6]).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.adaptive_avg_pool1d(torch_tensor, output_size=2) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.adaptive_avg_pool1d(ms_tensor, output_size=2) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_adaptive_max_pool1d(): data = np.random.randint(0, 10, [1, 3, 6]).astype(np.float32) torch_tensor = torch.tensor(data) torch_output = torch.nn.functional.adaptive_max_pool1d(torch_tensor, output_size=2) ms_tensor = ms_torch.tensor(data) ms_output = ms_torch.nn.functional.adaptive_max_pool1d(ms_tensor, output_size=2) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_adaptive_max_pool2d(): tensor = np.random.randn(1, 32, 9, 9).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_output = torch.nn.functional.adaptive_max_pool2d(torch_tensor, (3, 5)) ms_tensor = ms_torch.tensor(tensor) ms_output = ms_torch.nn.functional.adaptive_max_pool2d(ms_tensor, (3, 5)) assert (torch_output.shape == ms_output.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # def test_adaptive_avg_pool3d(): # tensor = np.random.randn(1, 32, 9, 9, 9).astype(np.float32) # # torch_tensor = torch.tensor(tensor) # torch_output = torch.nn.functional.adaptive_avg_pool3d(torch_tensor, (3, 5, 3)) # # ms_tensor = ms_torch.tensor(tensor) # ms_output = ms_torch.nn.functional.adaptive_avg_pool3d(ms_tensor, (3, 5, 3)) # # assert (torch_output.shape == ms_output.shape) # assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # # # def test_adaptive_max_pool3d(): # tensor = np.random.randn(1, 32, 9, 9, 9).astype(np.float32) # # torch_tensor = torch.tensor(tensor) # torch_output = torch.nn.functional.adaptive_max_pool3d(torch_tensor, (3, 5, 3)) # # ms_tensor = ms_torch.tensor(tensor) # ms_output = ms_torch.nn.functional.adaptive_max_pool3d(ms_tensor, (3, 5, 3)) # # assert (torch_output.shape == ms_output.shape) # assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_affine_grid1(): size = (1, 1, 3, 2) theta_np = np.array([[0.5, 0, 0.75], [0, 0.75, 0]]).reshape(1, 2, 3) torch_size = torch.Size(size) torch_theta = torch.from_numpy(theta_np) torch_out = torch.nn.functional.affine_grid(torch_theta, torch_size) ms_torch_size = ms_torch.Size(size) ms_torch_theta = ms_torch.from_numpy(theta_np) ms_torch_out = ms_torch.nn.functional.affine_grid(ms_torch_theta, ms_torch_size) assert (torch_out.shape == ms_torch_out.shape) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_affine_grid2(): theta = [[[0.8, 0.5, 0],[-0.5, 0.8, 0]]] size_ones = np.ones((1, 3, 2, 3)) torch_ones = torch.tensor(size_ones) torch_size = torch_ones.size() torch_theta = torch.tensor(theta, dtype=torch.float32) torch_out = torch.nn.functional.affine_grid(torch_theta, torch_size, align_corners=True) ms_torch_ones = ms_torch.tensor(size_ones) ms_torch_size = ms_torch_ones.size() ms_torch_theta = ms_torch.tensor(theta, dtype=ms_torch.float32) ms_torch_out = ms_torch.nn.functional.affine_grid(ms_torch_theta, ms_torch_size, align_corners=True) assert (torch_out.shape == ms_torch_out.shape) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) if __name__ == '__main__': test_interpolate1() test_interpolate2() test_interpolate3() test_interpolate4() test_interpolate5() test_upsample_nearest1() test_upsample_nearest2() test_upsample_nearest3() test_upsample_nearest4() test_upsample_nearest5() test_upsample_nearest6() test_upsample_bilinear1() test_upsample_bilinear2() test_upsample_bilinear3() test_adaptive_avg_pool1d() test_adaptive_max_pool1d() test_adaptive_avg_pool2d() test_adaptive_max_pool2d() # test_adaptive_avg_pool3d() # test_adaptive_max_pool3d() test_affine_grid1() test_affine_grid2()msadapter/testing/ut/pytorch/nn/functional/test_grid_sample.py000066400000000000000000000014161436713554100253220ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch from torch.nn.functional import grid_sample import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_grid_sample(): data = np.array(np.ones(shape=(2, 2, 2, 2))).astype(np.float32) grid_data = np.arange(0.2, 1, 0.1).reshape((2, 2, 1, 2)).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.grid_sample(ms_input, grid=ms_torch.tensor(grid_data)) torch_input = torch.tensor(data) torch_out = grid_sample(torch_input, grid=torch.tensor(grid_data)) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) if __name__ == '__main__': test_grid_sample() msadapter/testing/ut/pytorch/nn/functional/test_group_norm.py000066400000000000000000000066271436713554100252340ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_group_norm1(): np_input = np.random.randn(10, 8, 6).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.group_norm( ms_input, num_groups=2, weight=ms_weight, bias=ms_bias, eps=1e-5) torch_input = torch.tensor(np_input) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.group_norm( torch_input, num_groups=2, weight=torch_weight, bias=torch_bias, eps=1e-5) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_group_norm2(): np_input = np.random.randn(10, 8, 6).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_result = ms_torch.nn.functional.group_norm( ms_input, num_groups=2, eps=1e-5) torch_input = torch.tensor(np_input) torch_result = torch.nn.functional.group_norm( torch_input, num_groups=2, eps=1e-5) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_group_norm3(): np_input = np.random.randn(10, 8, 6, 6).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.group_norm( ms_input, num_groups=2, bias=ms_bias, eps=1e-5) torch_input = torch.tensor(np_input) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.group_norm( torch_input, num_groups=2, bias=torch_bias, eps=1e-5) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_group_norm4(): np_input = np.random.randn(10, 8, 6, 6, 6).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(weight) ms_result = ms_torch.nn.functional.group_norm( ms_input, num_groups=2, weight=ms_weight, eps=1e-5) torch_input = torch.tensor(np_input) torch_weight = torch.tensor(weight) torch_result = torch.nn.functional.group_norm( torch_input, num_groups=2, weight=torch_weight, eps=1e-5) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape if __name__ == "__main__": test_group_norm1() test_group_norm2() test_group_norm3() test_group_norm4() msadapter/testing/ut/pytorch/nn/functional/test_gumbelsoftmax.py000066400000000000000000000013751436713554100257150ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context import random context.set_context(mode=ms.PYNATIVE_MODE) def test_gumbel_softmax(): data = np.array([[-1, 0, 2, 3, 4, 7]],).astype(np.float32) # seed = 10 ms_input = ms_torch.tensor(data) # random.seed(seed) ms_out = ms_torch.nn.functional.gumbel_softmax(ms_input) torch_input = torch.tensor(data) # torch.manual_seed(seed) torch_out = torch.nn.functional.gumbel_softmax(torch_input) # assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_gumbel_softmax() msadapter/testing/ut/pytorch/nn/functional/test_hardsigmoid.py000066400000000000000000000022171436713554100253260ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_hardsigmoid1(): data = np.array([[-4, -3, 1, 0, -1.5, 2.8, 3.7, 100]]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.hardsigmoid(ms_input) torch_input = torch.tensor(data) torch_out = torch.nn.functional.hardsigmoid(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_hardsigmoid2(): data = np.array([[-4, -3, 1, 0, -1.5, 2.8, 3.7, 100]]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.hardsigmoid(ms_input, True) torch_input = torch.tensor(data) torch_out = torch.nn.functional.hardsigmoid(torch_input, True) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_hardsigmoid1() test_hardsigmoid2()msadapter/testing/ut/pytorch/nn/functional/test_instance_norm.py000066400000000000000000000176241436713554100257030ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_instance_norm1(): np_input = np.random.randn(10, 8, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.instance_norm( ms_input, ms_running_mean, ms_running_var, ms_weight, ms_bias, use_input_stats=True) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.instance_norm( torch_input, torch_running_mean, torch_running_var, torch_weight, torch_bias, use_input_stats=True) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_instance_norm2(): np_input = np.random.randn(10, 8, 6, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.instance_norm( ms_input, ms_running_mean, ms_running_var, ms_weight, ms_bias, use_input_stats=False) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.instance_norm( torch_input, torch_running_mean, torch_running_var, torch_weight, torch_bias, use_input_stats=False) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_instance_norm3(): np_input = np.random.randn(10, 8, 6, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_result = ms_torch.nn.functional.instance_norm( ms_input, ms_running_mean, ms_running_var, use_input_stats=False) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_result = torch.nn.functional.instance_norm( torch_input, torch_running_mean, torch_running_var, use_input_stats=False) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_instance_norm4(): np_input = np.random.randn(10, 8, 6, 6, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.instance_norm( ms_input, ms_running_mean, ms_running_var, ms_bias, use_input_stats=False) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.instance_norm( torch_input, torch_running_mean, torch_running_var, torch_bias, use_input_stats=False) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_instance_norm5(): np_input = np.random.randn(10, 8, 6, 6, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_weight = ms_torch.tensor(weight) ms_result = ms_torch.nn.functional.instance_norm( ms_input, ms_running_mean, ms_running_var, ms_weight, use_input_stats=False) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_weight = torch.tensor(weight) torch_result = torch.nn.functional.instance_norm( torch_input, torch_running_mean, torch_running_var, torch_weight, use_input_stats=False) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_instance_norm6(): np_input = np.random.randn(10, 8, 6).astype(np.float32) running_mean = np.random.randn(np_input.shape[1]).astype(np.float32) running_var = np.abs(np.random.randn(np_input.shape[1])).astype(np.float32) weight = np.random.randn(np_input.shape[1]).astype(np.float32) bias = np.random.randn(np_input.shape[1]).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_running_mean = ms_torch.tensor(running_mean) ms_running_var = ms_torch.tensor(running_var) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.instance_norm( ms_input, ms_running_mean, ms_running_var, ms_weight, ms_bias, use_input_stats=True, momentum=0.2, eps=1e-7) torch_input = torch.tensor(np_input) torch_running_mean = torch.tensor(running_mean) torch_running_var = torch.tensor(running_var) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.instance_norm( torch_input, torch_running_mean, torch_running_var, torch_weight, torch_bias, use_input_stats=True, momentum=0.2, eps=1e-7) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape if __name__ == "__main__": test_instance_norm1() test_instance_norm2() test_instance_norm3() test_instance_norm4() test_instance_norm5() test_instance_norm6() msadapter/testing/ut/pytorch/nn/functional/test_layer_norm.py000066400000000000000000000133021436713554100252000ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_layer_norm1(): np_input = np.random.randn(10, 8, 6).astype(np.float32) normalized_shape = (8, 6) weight = np.random.randn(*normalized_shape).astype(np.float32) bias = np.random.randn(*normalized_shape).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.layer_norm( ms_input, normalized_shape, ms_weight, ms_bias) torch_input = torch.tensor(np_input) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.layer_norm( torch_input, normalized_shape, torch_weight, torch_bias) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_layer_norm2(): np_input = np.random.randn(10, 8, 6, 6).astype(np.float32) normalized_shape = (8, 6, 6) weight = np.random.randn(*normalized_shape).astype(np.float32) bias = np.random.randn(*normalized_shape).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.layer_norm( ms_input, normalized_shape, ms_weight, ms_bias) torch_input = torch.tensor(np_input) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.layer_norm( torch_input, normalized_shape, torch_weight, torch_bias) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_layer_norm3(): np_input = np.random.randn(10, 8, 6, 6, 6).astype(np.float32) normalized_shape = (8, 6, 6, 6) weight = np.random.randn(*normalized_shape).astype(np.float32) bias = np.random.randn(*normalized_shape).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.layer_norm( ms_input, normalized_shape, ms_weight, ms_bias) torch_input = torch.tensor(np_input) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.layer_norm( torch_input, normalized_shape, torch_weight, torch_bias) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-4) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_layer_norm4(): np_input = np.random.randn(10, 8, 6, 6, 6).astype(np.float32) normalized_shape = (8, 6, 6, 6) bias = np.random.randn(*normalized_shape).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.layer_norm( ms_input, normalized_shape, ms_bias) torch_input = torch.tensor(np_input) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.layer_norm( torch_input, normalized_shape, torch_bias) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_layer_norm5(): np_input = np.random.randn(10, 8, 6, 6, 6).astype(np.float32) normalized_shape = (8, 6, 6, 6) weight = np.random.randn(*normalized_shape).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(weight) ms_result = ms_torch.nn.functional.layer_norm(ms_input, normalized_shape, ms_weight) torch_input = torch.tensor(np_input) torch_weight = torch.tensor(weight) torch_result = torch.nn.functional.layer_norm(torch_input, normalized_shape, torch_weight) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape def test_layer_norm6(): np_input = np.random.randn(10, 8, 6).astype(np.float32) normalized_shape = (8, 6) weight = np.random.randn(*normalized_shape).astype(np.float32) bias = np.random.randn(*normalized_shape).astype(np.float32) ms_input = ms_torch.tensor(np_input) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_result = ms_torch.nn.functional.layer_norm( ms_input, normalized_shape, ms_weight, ms_bias, eps=1e-7) torch_input = torch.tensor(np_input) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_result = torch.nn.functional.layer_norm( torch_input, normalized_shape, torch_weight, torch_bias, eps=1e-7) assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), rtol=1e-3, atol=1e-5) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype assert ms_result.asnumpy().shape == torch_result.numpy().shape if __name__ == "__main__": test_layer_norm1() test_layer_norm2() test_layer_norm3() test_layer_norm4() test_layer_norm5() test_layer_norm6() msadapter/testing/ut/pytorch/nn/functional/test_linear.py000066400000000000000000000053711436713554100243120ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_linear1(): data = np.ones((1, 2, 5)).astype(np.float32) weight = np.ones((3, 5)).astype(np.float32) ms_data = ms_torch.tensor(data) ms_weight = ms_torch.tensor(weight) ms_out = ms_torch.nn.functional.linear(ms_data, ms_weight) torch_data = torch.tensor(data) torch_weight = torch.tensor(weight) torch_out = torch.nn.functional.linear(torch_data, torch_weight) assert ms_out.shape == torch_out.shape assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_linear2(): data = np.ones((1, 2, 5)).astype(np.float32) weight = np.ones((3, 5)).astype(np.float32) bias = np.ones((3)).astype(np.float32) ms_data = ms_torch.tensor(data) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_out = ms_torch.nn.functional.linear(ms_data, ms_weight, bias=ms_bias) torch_data = torch.tensor(data) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_out = torch.nn.functional.linear(torch_data, torch_weight, bias=torch_bias) assert ms_out.shape == torch_out.shape assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_linear3(): data = np.ones((1, 2, 5)).astype(np.float32) weight = np.ones((5)).astype(np.float32) ms_data = ms_torch.tensor(data) ms_weight = ms_torch.tensor(weight) ms_out = ms_torch.nn.functional.linear(ms_data, ms_weight) torch_data = torch.tensor(data) torch_weight = torch.tensor(weight) torch_out = torch.nn.functional.linear(torch_data, torch_weight) assert ms_out.shape == torch_out.shape assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_bilinear(): data1 = np.random.randn(2, 2, 3) data2 = np.random.randn(2, 2, 5) weight = np.random.randn(4, 3, 5) bias = np.random.randn(4) torch_input1 = torch.tensor(data1) torch_input2 = torch.tensor(data2) torch_weight = torch.tensor(weight) torch_bias = torch.tensor(bias) torch_out = torch.nn.functional.bilinear(torch_input1, torch_input2, torch_weight, torch_bias) ms_input1 = ms_torch.tensor(data1) ms_input2 = ms_torch.tensor(data2) ms_weight = ms_torch.tensor(weight) ms_bias = ms_torch.tensor(bias) ms_out = ms_torch.nn.functional.bilinear(ms_input1, ms_input2, ms_weight, ms_bias) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.shape == torch_out.shape assert ms_out.numpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_linear1() test_linear2() test_linear3() test_bilinear() msadapter/testing/ut/pytorch/nn/functional/test_log_softmax.py000066400000000000000000000033261436713554100253600ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_log_softmax1(): ms_result = ms_torch.nn.functional.log_softmax(ms_torch.tensor([1, 2, 3, 4, 5]), \ dim=0, dtype=ms_torch.float32) torch_result = torch.nn.functional.log_softmax(torch.tensor([1, 2, 3, 4, 5]), \ dim=0, dtype=torch.float32) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_log_softmax2(): x = np.array([[-0.08082921, -0.13706027, -0.4711177, -0.05606057], [-0.46082982, 1.1761844, -1.016654, -1.743829], [-1.5062045, 0.6910976, 0.4839723, 1.1502692]]).astype(np.float32) expect = np.array([[-1.2939762, -1.3502073, -1.6842647, -1.2692076], [-1.9445671, -0.3075528, -2.5003912, -3.2275662], [-3.452001, -1.2546989, -1.4618242, -0.79552734]]).astype(np.float32) output = ms_torch.nn.functional.log_softmax(ms_torch.tensor(x), 1) assert np.allclose(output.asnumpy(), expect) def test_log_softmax_dim(): input = np.array([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]]) torch_input = torch.tensor(input, dtype=torch.float32) torch_output = torch.nn.functional.log_softmax(torch_input) print(torch_output) ms_torch_input = ms_torch.tensor(input, dtype=ms_torch.float32) ms_torch_output = ms_torch.nn.functional.log_softmax(ms_torch_input) print(ms_torch_output) assert np.allclose(torch_output.numpy(), ms_torch_output.asnumpy()) if __name__ == '__main__': test_log_softmax1() test_log_softmax2() test_log_softmax_dim() msadapter/testing/ut/pytorch/nn/functional/test_loss.py000066400000000000000000000151431436713554100240160ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.PYNATIVE_MODE) def test_ctc_loss(): np_log_probs = np.random.randn(24, 2, 10) np_targets = np.random.rand(2, 10)*10 np_input_length = np.array([8, 10]) np_target_length = np.array([5, 6]) torch_log_probs = torch.tensor(np_log_probs) torch_targets = torch.tensor(np_targets) torch_input_length = torch.tensor(np_input_length) torch_target_length = torch.tensor(np_target_length) result_torch = torch.nn.functional.ctc_loss(torch_log_probs, torch_targets, torch_input_length, torch_target_length) ms_log_probs = ms_torch.tensor(np_log_probs) ms_targets = ms_torch.tensor(np_targets) ms_input_length = ms_torch.tensor(np_input_length) ms_target_length = ms_torch.tensor(np_target_length) result_ms = ms_torch.nn.functional.ctc_loss(ms_log_probs, ms_targets, ms_input_length, ms_target_length) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_gaussian_nll_loss(): np_input = np.random.randn(5, 6, 7, 8) np_targets = np.random.rand(5, 6, 7, 8) np_var = np.ones([5, 6, 7, 8]) torch_log_probs = torch.tensor(np_input) torch_targets = torch.tensor(np_targets) torch_var = torch.tensor(np_var) result_torch = torch.nn.functional.gaussian_nll_loss(torch_log_probs, torch_targets, torch_var) ms_log_probs = ms_torch.tensor(np_input) ms_targets = ms_torch.tensor(np_targets) ms_var = ms_torch.tensor(np_var) result_ms = ms_torch.nn.functional.gaussian_nll_loss(ms_log_probs, ms_targets, ms_var) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_hinge_embedding_loss(): np_input = np.random.randn(5, 6, 7, 8) np_targets = np.sign(np.random.rand(5, 6, 7, 8)) torch_log_probs = torch.tensor(np_input) torch_targets = torch.tensor(np_targets) result_torch = torch.nn.functional.hinge_embedding_loss(torch_log_probs, torch_targets) ms_log_probs = ms_torch.tensor(np_input) ms_targets = ms_torch.tensor(np_targets) result_ms = ms_torch.nn.functional.hinge_embedding_loss(ms_log_probs, ms_targets) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_margin_ranking_loss(): np_input1 = np.random.randn(5, 6, 7, 8) np_input2 = np.random.randn(5, 6, 7, 8) np_targets = np.sign(np.random.rand(5, 6, 7, 8)) torch_input1 = torch.tensor(np_input1) torch_input2 = torch.tensor(np_input2) torch_targets = torch.tensor(np_targets) result_torch = torch.nn.functional.margin_ranking_loss(torch_input1, torch_input2, torch_targets) ms_input1 = ms_torch.tensor(np_input1) ms_input2 = ms_torch.tensor(np_input2) ms_targets = ms_torch.tensor(np_targets) result_ms = ms_torch.nn.functional.margin_ranking_loss(ms_input1, ms_input2, ms_targets) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape # def test_multilabel_margin_loss(): # np_input = np.random.randn(5, 10) # np_targets = np.random.rand(5, 10)*10-1 # # torch_input = torch.tensor(np_input) # torch_targets = torch.tensor(np_targets, dtype=torch.int64) # result_torch = torch.nn.functional.multilabel_margin_loss(torch_input, torch_targets) # # ms_input = ms_torch.tensor(np_input) # ms_targets = ms_torch.tensor(np_targets, dtype=ms.int64) # result_ms = ms_torch.nn.functional.multilabel_margin_loss(ms_input, ms_targets) # # assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) # assert result_ms.shape == result_torch.shape def test_multilabel_soft_margin_loss(): np_input = np.array([[0.3, 0.6, 0.6], [0.9, 0.4, 0.2]]) np_targets = np.array([[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]]) torch_input = torch.tensor(np_input) torch_targets = torch.tensor(np_targets) result_torch = torch.nn.functional.multilabel_soft_margin_loss(torch_input, torch_targets) ms_input = ms_torch.tensor(np_input) ms_targets = ms_torch.tensor(np_targets) result_ms = ms_torch.nn.functional.multilabel_soft_margin_loss(ms_input, ms_targets) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_multi_margin_loss(): np_input = np.array([[0.1, 0.2, 0.4, 0.8]]) np_targets = np.array([3]) torch_input = torch.tensor(np_input) torch_targets = torch.tensor(np_targets, dtype=torch.int64) result_torch = torch.nn.functional.multi_margin_loss(torch_input, torch_targets) ms_input = ms_torch.tensor(np_input) ms_targets = ms_torch.tensor(np_targets, dtype=ms.int64) result_ms = ms_torch.nn.functional.multi_margin_loss(ms_input, ms_targets) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_huber_loss(): np_input = np.array([1.0, 2, 10, 2]) np_targets = np.array([1.0, 5, 1, 20]) torch_input = torch.tensor(np_input) torch_targets = torch.tensor(np_targets) result_torch = torch.nn.functional.huber_loss(torch_input, torch_targets) ms_input = ms_torch.tensor(np_input) ms_targets = ms_torch.tensor(np_targets) result_ms = ms_torch.nn.functional.huber_loss(ms_input, ms_targets) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_triplet_margin_loss(): np_anc = np.random.randn(100, 128) np_pos = np.random.randn(100, 128) np_neg = np.random.randn(100, 128) torch_anc = torch.tensor(np_anc, requires_grad=True) torch_pos = torch.tensor(np_pos, requires_grad=True) torch_neg = torch.tensor(np_neg, requires_grad=True) result_torch = torch.nn.functional.triplet_margin_loss(torch_anc, torch_pos, torch_neg) ms_anc = ms_torch.tensor(np_anc, requires_grad=True) ms_pos = ms_torch.tensor(np_pos, requires_grad=True) ms_neg = ms_torch.tensor(np_neg, requires_grad=True) result_ms = ms_torch.nn.functional.triplet_margin_loss(ms_anc, ms_pos, ms_neg) assert np.allclose(result_ms.asnumpy(), result_torch.detach().numpy()) assert result_ms.shape == result_torch.shape if __name__ == '__main__': test_ctc_loss() test_gaussian_nll_loss() test_hinge_embedding_loss() test_margin_ranking_loss() # test_multilabel_margin_loss() test_multilabel_soft_margin_loss() test_multi_margin_loss() test_huber_loss() test_triplet_margin_loss() msadapter/testing/ut/pytorch/nn/functional/test_max_pooling.py000066400000000000000000000315411436713554100253520ustar00rootroot00000000000000import numpy as np import torch import mindspore as ms import ms_adapter.pytorch as ms_torch ms.context.set_context(mode=ms.PYNATIVE_MODE) def test_max_pool1d_1(): N = np.random.randint(1, 33) C = np.random.randint(2, 65) tensor = np.random.randn(N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output, torch_indices = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, return_indices=True) ms_output, ms_indices = ms_torch.nn.functional.max_pool1d(ms_tensor, kernel_size, stride, padding, return_indices=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool1d_2(): N = np.random.randint(1, 33) C = np.random.randint(2, 65) tensor = np.random.randn(N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding) ms_output = ms_torch.nn.functional.max_pool1d(ms_tensor, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool1d_3(): B = np.random.randint(1, 33) N = np.random.randint(1, 33) C = np.random.randint(2, 65) tensor = np.random.randn(B, N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output, torch_indices = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, return_indices=True) ms_output, ms_indices = ms_torch.nn.functional.max_pool1d(ms_tensor, kernel_size, stride, padding, return_indices=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool1d_4(): B = np.random.randint(1, 33) N = np.random.randint(1, 33) C = np.random.randint(2, 65) # tensor = np.random.randn(B, N, C).astype(np.float32) # now ms.ops.pad is not supported on Ascend, max_pool1d use ms.ops.Pad instead, the pad value is 0. # therefore some result at boarder may be wrong when tensor value are smaller than zero. # use np.random.random instead of np.random.randn to generate value that are bigger than zero. tensor = np.random.random((B, N, C)).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding) ms_output = ms_torch.nn.functional.max_pool1d(ms_tensor, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool1d_5(): B = np.random.randint(1, 33) N = np.random.randint(1, 33) C = np.random.randint(2, 65) tensor = np.random.randn(B, N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, ceil_mode=True) ms_output = ms_torch.nn.functional.max_pool1d(ms_tensor, kernel_size, stride, padding, ceil_mode=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool1d_6(): B = np.random.randint(1, 33) N = np.random.randint(1, 33) C = np.random.randint(5, 65) tensor = np.random.randn(B, N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, (C-1)/2) padding = np.random.randint(0, kernel_size/2 + 1) stride = np.random.randint(1, C + 1) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, dilation=2, ceil_mode=True) ms_output = ms_torch.nn.functional.max_pool1d(ms_tensor, kernel_size, stride, padding, dilation=2, ceil_mode=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool3d_1(): N = np.random.randint(1, 17) C = np.random.randint(2, 33) D = np.random.randint(1, 17) H = np.random.randint(1, 33) W = np.random.randint(1, 33) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding) ms_output = ms_torch.nn.functional.max_pool3d(ms_tensor, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool3d_2(): N = np.random.randint(1, 17) C = np.random.randint(2, 33) D = np.random.randint(1, 17) H = np.random.randint(1, 33) W = np.random.randint(1, 33) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, return_indices=True) ms_output, ms_indices = ms_torch.nn.functional.max_pool3d(ms_tensor, kernel_size, stride, padding, return_indices=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool3d_3(): N = np.random.randint(1, 17) C = np.random.randint(2, 33) D = np.random.randint(5, 17) H = np.random.randint(5, 33) W = np.random.randint(5, 33) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, (D - 1)/2), np.random.randint(1, (H - 1)/2), np.random.randint(1, (W - 1)/2)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, dilation=2, return_indices=True) ms_output, ms_indices = ms_torch.nn.functional.max_pool3d(ms_tensor, kernel_size, stride, padding, dilation=2, return_indices=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool3d_4(): N = np.random.randint(1, 17) C = np.random.randint(2, 33) D = np.random.randint(5, 17) H = np.random.randint(5, 33) W = np.random.randint(5, 33) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, (D - 1)/2), np.random.randint(1, (H - 1)/2), np.random.randint(1, (W - 1)/2)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, dilation=2, return_indices=True, ceil_mode=True) ms_output, ms_indices = ms_torch.nn.functional.max_pool3d(ms_tensor, kernel_size, stride, padding, dilation=2, return_indices=True, ceil_mode=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool3d_5(): C = np.random.randint(2, 33) D = np.random.randint(5, 17) H = np.random.randint(5, 33) W = np.random.randint(5, 33) tensor = np.random.randn(C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, (D - 1)/2), np.random.randint(1, (H - 1)/2), np.random.randint(1, (W - 1)/2)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, dilation=2, return_indices=True, ceil_mode=True) ms_output, ms_indices = ms_torch.nn.functional.max_pool3d(ms_tensor, kernel_size, stride, padding, dilation=2, return_indices=True, ceil_mode=True) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape def test_max_pool3d_6(): C = np.random.randint(2, 33) D = np.random.randint(5, 17) H = np.random.randint(5, 33) W = np.random.randint(5, 33) tensor = np.random.randn(C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = (np.random.randint(1, (D - 1)/2), np.random.randint(1, (H - 1)/2), np.random.randint(1, (W - 1)/2)) padding = (np.random.randint(0, kernel_size[0]/2 + 1), np.random.randint(0, kernel_size[1]/2 + 1), np.random.randint(0, kernel_size[2]/2 + 1)) stride = (np.random.randint(1, D + 1), np.random.randint(1, H + 1), np.random.randint(1, W + 1)) ms_tensor = ms_torch.tensor(tensor) torch_output = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, dilation=2) ms_output = ms_torch.nn.functional.max_pool3d(ms_tensor, kernel_size, stride, padding, dilation=2) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().shape == torch_output.numpy().shape if __name__ == '__main__': test_max_pool3d_1() test_max_pool3d_2() test_max_pool3d_3() test_max_pool3d_4() test_max_pool3d_5() test_max_pool3d_6() test_max_pool1d_1() test_max_pool1d_2() test_max_pool1d_3() test_max_pool1d_4() test_max_pool1d_5() test_max_pool1d_6() msadapter/testing/ut/pytorch/nn/functional/test_maxunpool.py000066400000000000000000000262161436713554100250630ustar00rootroot00000000000000import numpy as np import torch import mindspore as ms import ms_adapter.pytorch as ms_torch ms.context.set_context(mode=ms.PYNATIVE_MODE) def test_maxunpool1d_with_2dim(): N = np.random.randint(1, 33) C = np.random.randint(1, 257) tensor = np.random.randn(N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool1d(torch_pooling, torch_indices, kernel_size, stride, padding) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool1d(ms_pooling, ms_indices, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool1d_with_3dim(): B = np.random.randint(1, 33) N = np.random.randint(1, 33) C = np.random.randint(1, 257) tensor = np.random.randn(B, N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool1d(torch_pooling, torch_indices, kernel_size, stride, padding) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool1d(ms_pooling, ms_indices, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # def test_maxunpool1d_with_3dim_shape(): # B = np.random.randint(1, 33) # N = np.random.randint(1, 33) # C = np.random.randint(1, 257) # tensor = np.random.randn(B, N, C).astype(np.float32) # torch_tensor = torch.tensor(tensor) # kernel_size = np.random.randint(1, C + 1) # padding = np.random.randint(0, kernel_size/2 + 1) # stride = 1 # torch_pooling, torch_indices = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, return_indices=True) # torch_output = torch.nn.functional.max_unpool1d(torch_pooling, torch_indices, kernel_size, stride, padding, output_size=torch_tensor.size()) # ms_pooling = ms_torch.tensor(torch_pooling.numpy()) # ms_indices = ms_torch.tensor(torch_indices.numpy()) # ms_output = ms_torch.nn.functional.max_unpool1d(ms_pooling, ms_indices, kernel_size, stride, padding, tensor.shape) # assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool1d_with_2dim_shape(): N = np.random.randint(1, 33) C = np.random.randint(1, 257) tensor = np.random.randn(N, C).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, C + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool1d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool1d(torch_pooling, torch_indices, kernel_size, stride, padding, output_size=[C]) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool1d(ms_pooling, ms_indices, kernel_size, stride, padding, tensor.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool2d_with_3dim(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 257) tensor = np.random.randn(C, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, H + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool2d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool2d(torch_pooling, torch_indices, kernel_size, stride, padding) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool2d(ms_pooling, ms_indices, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool2d_with_3dim_shape(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 257) tensor = np.random.randn(C, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, H + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool2d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool2d(torch_pooling, torch_indices, kernel_size, stride, padding, [H, W]) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool2d(ms_pooling, ms_indices, kernel_size, stride, padding,tensor.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool2d_with_4dim(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 129) N = np.random.randint(1, 33) tensor = np.random.randn(N, C, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, H + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool2d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool2d(torch_pooling, torch_indices, kernel_size, stride, padding) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool2d(ms_pooling, ms_indices, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool2d_with_4dim_shape(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 129) N = np.random.randint(1, 33) tensor = np.random.randn(N, C, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, H + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool2d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool2d(torch_pooling, torch_indices, kernel_size, stride, padding, torch_tensor.size()) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool2d(ms_pooling, ms_indices, kernel_size, stride, padding, tensor.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool3d_with_4dim(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 129) D = np.random.randint(1, H + 1) tensor = np.random.randn(C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, D + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool3d(torch_pooling, torch_indices, kernel_size, stride, padding) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool3d(ms_pooling, ms_indices, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool3d_with_4dim_shape(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 129) D = np.random.randint(1, H + 1) tensor = np.random.randn(C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, D + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool3d(torch_pooling, torch_indices, kernel_size, stride, padding, [D, H, W]) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool3d(ms_pooling, ms_indices, kernel_size, stride, padding, tensor.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool3d_with_5dim(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 129) D = np.random.randint(1, H + 1) N = np.random.randint(1, 17) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, D + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool3d(torch_pooling, torch_indices, kernel_size, stride, padding) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool3d(ms_pooling, ms_indices, kernel_size, stride, padding) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool3d_with_5dim_shape(): H = W = np.random.randint(1, 33) C = np.random.randint(1, 129) D = np.random.randint(1, H + 1) N = np.random.randint(1, 17) tensor = np.random.randn(N, C, D, H, W).astype(np.float32) torch_tensor = torch.tensor(tensor) kernel_size = np.random.randint(1, D + 1) padding = np.random.randint(0, kernel_size/2 + 1) stride = 1 torch_pooling, torch_indices = torch.nn.functional.max_pool3d(torch_tensor, kernel_size, stride, padding, return_indices=True) torch_output = torch.nn.functional.max_unpool3d(torch_pooling, torch_indices, kernel_size, stride, padding, torch_tensor.size()) ms_pooling = ms_torch.tensor(torch_pooling.numpy()) ms_indices = ms_torch.tensor(torch_indices.numpy()) ms_output = ms_torch.nn.functional.max_unpool3d(ms_pooling, ms_indices, kernel_size, stride, padding, tensor.shape) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) if __name__ == '__main__': test_maxunpool1d_with_2dim() test_maxunpool1d_with_3dim() test_maxunpool2d_with_3dim() test_maxunpool2d_with_4dim() test_maxunpool3d_with_4dim() test_maxunpool3d_with_5dim() test_maxunpool1d_with_2dim_shape() # test_maxunpool1d_with_3dim_shape() test_maxunpool2d_with_3dim_shape() test_maxunpool2d_with_4dim_shape() test_maxunpool3d_with_4dim_shape() test_maxunpool3d_with_5dim_shape() msadapter/testing/ut/pytorch/nn/functional/test_normalize.py000066400000000000000000000041601436713554100250330ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_normalize1(): data = np.array([[-1, 0, 1, -2], [2, 2, 3, 4] ]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.normalize(ms_input) torch_input = torch.tensor(data) torch_out = torch.nn.functional.normalize(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_normalize2(): data = np.array([[[-1, 0, 1, -2], [2, 2, 3, 4] ]]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.normalize(ms_input, 2.2, 0, 10) torch_input = torch.tensor(data) torch_out = torch.nn.functional.normalize(torch_input, 2.2, 0, 10) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def local_response_norm1(): data = np.random.random((4, 4, 4)).astype(np.float32)*10 ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.local_response_norm(ms_input, 3, 0.01) torch_input = torch.tensor(data) torch_out = torch.nn.functional.local_response_norm(torch_input, 3, 0.01) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def local_response_norm2(): data = np.random.random((2, 4, 4, 3, 2)).astype(np.float32)*10 ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.local_response_norm(ms_input, 3, 0.01) torch_input = torch.tensor(data) torch_out = torch.nn.functional.local_response_norm(torch_input, 3, 0.01) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_normalize1() test_normalize2() local_response_norm1() local_response_norm2() msadapter/testing/ut/pytorch/nn/functional/test_one_hot.py000066400000000000000000000021211436713554100244610ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_one_hot1(): data = np.array([0, 2, 3, 4, 7]).astype(np.int64) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.one_hot(ms_input) torch_input = torch.tensor(data) torch_out = torch.nn.functional.one_hot(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_one_hot2(): data = np.array([[[0, 2, 3], [1, 1, 2]]]).astype(np.int64) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.one_hot(input=ms_input, num_classes=5) torch_input = torch.tensor(data) torch_out = torch.nn.functional.one_hot(input=torch_input, num_classes=5) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_one_hot1() test_one_hot2() msadapter/testing/ut/pytorch/nn/functional/test_pad.py000066400000000000000000000021011436713554100235700ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch from torch.nn import functional as F import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_torch from ms_adapter.pytorch.nn.functional import pad context.set_context(mode=ms.PYNATIVE_MODE) def test_pad_shape(): p1d = (1, 1) py_t4d = torch.ones(3, 3, 4, 2) py_out_1 = F.pad(py_t4d, p1d, "constant", 0) ms_t4d = ms_torch.ones(3, 3, 4, 2) ms_out_1 = pad(ms_t4d, p1d, "constant", 0) assert np.allclose(py_out_1.numpy(), ms_out_1.asnumpy()) p3d = (1, 1, 0, 2, 1, 1) py_out_2 = F.pad(py_t4d, p3d, "constant", 2) ms_out_2 = pad(ms_t4d, p3d, "constant", 2) assert np.allclose(py_out_2.numpy(), ms_out_2.asnumpy()) def test_pad_mode(): p1d = (1, 1) py_t4d = torch.ones(4, 3, 2) py_out_1 = F.pad(py_t4d, p1d, "reflect") ms_t4d = ms_torch.ones(4, 3, 2) ms_out_1 = pad(ms_t4d, p1d, "reflect") assert np.allclose(py_out_1.numpy(), ms_out_1.asnumpy()) if __name__ == '__main__': test_pad_shape() test_pad_mode()msadapter/testing/ut/pytorch/nn/functional/test_pdist.py000066400000000000000000000020231436713554100241520ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.GRAPH_MODE) def test_pdist1(): data = np.arange(0, 12).reshape(4, 3).astype(np.float32) torch_input = torch.tensor(data) torch_out = torch.nn.functional.pdist(torch_input, 3.1) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.pdist(ms_input, 3.1) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pdist2(): data = np.arange(0, 12).reshape(4, 3).astype(np.float32) torch_input = torch.tensor(data) torch_out = torch.nn.functional.pdist(torch_input, 0) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.pdist(ms_input, 0) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_pdist1() test_pdist2() msadapter/testing/ut/pytorch/nn/functional/test_relu.py000066400000000000000000000012021436713554100237740ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_relu1(): np_1 = np.random.randn(2, 3, 4, 5) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.nn.functional.relu(ms_tensor) torch_tensor = torch.tensor(np_1) torch_result = torch.nn.functional.relu(torch_tensor) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype if __name__ == '__main__': test_relu1() msadapter/testing/ut/pytorch/nn/functional/test_sigmoid.py000066400000000000000000000013021436713554100244610ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_sigmoid(): data = np.array([[-1, 0, 1, -2], [2, 2, 3, 4] ]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.sigmoid(ms_input) torch_input = torch.tensor(data) torch_out = torch.nn.functional.sigmoid(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_sigmoid() msadapter/testing/ut/pytorch/nn/functional/test_silu.py000066400000000000000000000022011436713554100240010ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_silu1(): data = np.array([[-4, -3, 1, 0, -1.5, 2.8, 3.7, 100]]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.silu(ms_input) torch_input = torch.tensor(data) torch_out = torch.nn.functional.silu(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_silu2(): data = np.array([[-4, -3, 1, 0, -1.5, 2.8, 3.7, 100]]).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.silu(ms_input, True) torch_input = torch.tensor(data) torch_out = torch.nn.functional.silu(torch_input, True) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_silu1() test_silu2()msadapter/testing/ut/pytorch/nn/functional/test_softplus.py000066400000000000000000000030741436713554100247150ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_softplus1(): data = np.array([[-1, 0, 2, 3, 4, 7]],).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.softplus(ms_input) torch_input = torch.tensor(data) torch_out = torch.nn.functional.softplus(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_softplus2(): data = np.array([[-1, 0, 2, 3, 4, 7]],).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.softplus(ms_input, beta=5.2, threshold=16.5) torch_input = torch.tensor(data) torch_out = torch.nn.functional.softplus(torch_input, beta=5.2, threshold=16.5) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_softplus3(): data = np.array([[-1, 0, 2, 19.9, 20, 21]],).astype(np.float32) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.softplus(ms_input, beta=1, threshold=20) torch_input = torch.tensor(data) torch_out = torch.nn.functional.softplus(torch_input, beta=1, threshold=20) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_softplus1() test_softplus2() test_softplus3()msadapter/testing/ut/pytorch/nn/functional/test_upsample.py000066400000000000000000000033271436713554100246650ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_upsample1(): np_1 = np.random.randn(2, 3, 4, 5) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.nn.functional.upsample(ms_tensor, (3, 3)) torch_tensor = torch.tensor(np_1) torch_result = torch.nn.functional.upsample(torch_tensor, (3, 3)) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_upsample2(): np_1 = np.random.randn(2, 3, 4) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.nn.functional.upsample(ms_tensor, 6, align_corners=True, mode='linear') torch_tensor = torch.tensor(np_1) torch_result = torch.nn.functional.upsample(torch_tensor, 6, align_corners=True, mode='linear') assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype # TODO: default is 'align_corners=False', which is not supported on CPU right now. ''' def test_upsample3(): np_1 = np.random.randn(2, 3, 4, 5).astype(np.float32) ms_tensor = ms_torch.tensor(np_1) ms_result = ms_torch.nn.functional.upsample(ms_tensor, (5, 5), mode='bilinear') torch_tensor = torch.tensor(np_1) torch_result = torch.nn.functional.upsample(torch_tensor, (5, 5), mode='bilinear') assert np.allclose(ms_result.asnumpy(), torch_result.numpy(), atol=1e-1) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype ''' if __name__ == '__main__': test_upsample1() test_upsample2() #test_upsample3() msadapter/testing/ut/pytorch/nn/functional/test_vision.py000066400000000000000000000037441436713554100243510ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_pixel_shuffle1(): data = np.arange(0, 240).reshape(1, 2, 8, 3, 5) torch_input = torch.tensor(data) torch_out = torch.nn.functional.pixel_shuffle(torch_input, 2) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.pixel_shuffle(ms_input, 2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pixel_shuffle2(): data = np.arange(0, 36).reshape(4, 3, 3).astype(np.float32) torch_input = torch.tensor(data) torch_out = torch.nn.functional.pixel_shuffle(torch_input, 2) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.pixel_shuffle(ms_input, 2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pixel_unshuffle1(): data = np.arange(0, 240).reshape(1, 2, 2, 6, 10) torch_input = torch.tensor(data) torch_out = torch.nn.functional.pixel_unshuffle(torch_input, 2) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.pixel_unshuffle(ms_input, 2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pixel_unshuffle2(): data = np.arange(0, 36).reshape(1, 6, 6).astype(np.float32) torch_input = torch.tensor(data) torch_out = torch.nn.functional.pixel_unshuffle(torch_input, 2) ms_input = ms_torch.tensor(data) ms_out = ms_torch.nn.functional.pixel_unshuffle(ms_input, 2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_pixel_shuffle1() test_pixel_shuffle2() test_pixel_unshuffle1() test_pixel_unshuffle2() msadapter/testing/ut/pytorch/nn/test_activation.py000066400000000000000000000450271436713554100230410ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import ms_adapter.pytorch as ms_torch from ms_adapter.pytorch.nn import Module, ReLU, ReLU6 from ms_adapter.pytorch import Tensor import ms_adapter.pytorch.nn as nn import numpy as np from mindspore import context import mindspore as ms import torch context.set_context(mode=ms.GRAPH_MODE) def test_relu1(): ms_net = ReLU() torch_net = torch.nn.ReLU() data = np.array([-1, 0, 1, -2, 4, 5, 6, 7]) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data.astype(np.float32)) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_relu2(): context.set_context(mode=ms.PYNATIVE_MODE) ms_net = ReLU(inplace=True) torch_net = torch.nn.ReLU(inplace=True) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data.astype(np.float32)) ms_output = ms_net(ms_input) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_hardtanh1(): ms_net = nn.Hardtanh() torch_net = torch.nn.Hardtanh() data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_hardtanh2(): context.set_context(mode=ms.PYNATIVE_MODE) ms_net = nn.Hardtanh(inplace=True) torch_net = torch.nn.Hardtanh(inplace=True) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_relu6(): context.set_context(mode=ms.PYNATIVE_MODE) ms_net = ReLU6(inplace=True) torch_net = torch.nn.ReLU6(inplace=True) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data.astype(np.float32)) ms_output = ms_net(ms_input) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_silu1(): ms_net = nn.SiLU() torch_net = torch.nn.SiLU() data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_silu2(): context.set_context(mode=ms.PYNATIVE_MODE) ms_net = nn.SiLU(inplace=True) torch_net = torch.nn.SiLU(inplace=True) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) ms_input = Tensor(data) torch_output = torch_net(torch_input) ms_output = ms_net(ms_input) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_hardswish1(): ms_net = nn.Hardswish() torch_net = torch.nn.Hardswish() data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_hardswish2(): context.set_context(mode=ms.PYNATIVE_MODE) ms_net = nn.Hardswish(inplace=True) torch_net = torch.nn.Hardswish(inplace=True) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) ms_input = Tensor(data) torch_output = torch_net(torch_input) ms_output = ms_net(ms_input) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_leakyReLU1(): ms_net = nn.LeakyReLU(negative_slope=1.2) torch_net = torch.nn.LeakyReLU(negative_slope=1.2) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_leakyReLU2(): context.set_context(mode=ms.PYNATIVE_MODE) ms_net = nn.LeakyReLU(negative_slope=0.1, inplace=True) torch_net = torch.nn.LeakyReLU(negative_slope=0.1, inplace=True) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) ms_input = Tensor(data) torch_output = torch_net(torch_input) ms_output = ms_net(ms_input) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_sigmoid(): ms_net = nn.Sigmoid() torch_net = torch.nn.Sigmoid() data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_logsigmoid(): ms_net = nn.LogSigmoid() torch_net = torch.nn.LogSigmoid() data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_elu(): ms_net = nn.ELU() torch_net = torch.nn.ELU() data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_rrelu(): ms_net = nn.RReLU() torch_net = torch.nn.RReLU() data = np.random.rand(2, 3).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_selu(): data = np.random.rand(2, 3).astype(np.float32) ms_net = nn.SELU(inplace=False) torch_net = torch.nn.SELU(inplace=False) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_celu(): data = np.random.rand(2, 3).astype(np.float32) ms_net = nn.CELU(inplace=False) torch_net = torch.nn.CELU(inplace=False) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_gelu(): data = np.random.rand(5, 6).astype(np.float32) ms_net = nn.GELU() ms_net2 = nn.GELU(approximate='tanh') torch_net = torch.nn.GELU() torch_net2 = torch.nn.GELU(approximate='tanh') torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) torch_output2 = torch_net2(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) ms_output2 = ms_net2(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_output2.asnumpy(), torch_output2.numpy()) def test_mish(): data = np.random.rand(2, 3).astype(np.float32) ms_net = nn.Mish(inplace=False) torch_net = torch.nn.Mish(inplace=False) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) if context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_softshrink(): data = np.random.rand(2, 3).astype(np.float32) ms_net = nn.Softshrink() torch_net = torch.nn.Softshrink() torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_tanh(): ms_net = nn.Tanh() torch_net = torch.nn.Tanh() data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_tanhshrink(): ms_net = nn.Tanhshrink() torch_net = torch.nn.Tanhshrink() data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_input.asnumpy(), torch_input.numpy()) def test_threshold1(): ms_net = nn.Threshold(threshold=3.5, value=20) torch_net = torch.nn.Threshold(threshold=3.5, value=20) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_threshold2(): #ms_net = nn.Threshold(threshold=3.5, value=3, inplace=True) #torch_net = torch.nn.Threshold(threshold=3.5, value=3, inplace=True) #data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) ms_net = nn.Threshold(threshold=1, value=9, inplace=True) torch_net = torch.nn.Threshold(threshold=1, value=9, inplace=True) data = np.ones((5, 6, 8)).astype(np.uint8) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().dtype == torch_output.numpy().dtype def test_softmax1(): ms_net = nn.Softmax(dim=-1) torch_net = torch.nn.Softmax(dim=-1) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_softmax2(): ms_net = nn.Softmax(dim=1) torch_net = torch.nn.Softmax(dim=1) data = np.array([[[-1, 0, 1, -2, 4, 5, 6, 7]], [[-1, 0, 1, -2, 4, 5, 6, 7]], [[5, 4, 1, 2, 3, 4, -5, -6]] ]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_logsoftmax1(): ms_net = nn.LogSoftmax(dim=-1) torch_net = torch.nn.LogSoftmax(dim=-1) data = np.array([[-1, 0, 1, -2, 4, 5, 6, 7]]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_logsoftmax2(): ms_net = nn.LogSoftmax(dim=0) torch_net = torch.nn.LogSoftmax(dim=0) data = np.array([[[-1, 0, 1, -2, 4, 5, 6, 7]], [[-1, 0, 1, -2, 4, 5, 6, 7]], [[5, 4, 1, 2, 3, 4, -5, -6]] ]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_softmin(): ms_net = nn.Softmin(dim=1) torch_net = torch.nn.Softmin(dim=1) data = np.array([[[-1, 0, 1, -2, 4, 5, 6, 7]], [[-1, 0, 1, -2, 4, 5, 6, 7]], [[5, 4, 1, 2, 3, 4, -5, -6]] ]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_softsign(): ms_net = nn.Softsign() torch_net = torch.nn.Softsign() data = np.array([[-1, 0, 1, -2, 4, -5, 6, -7]],).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_glu(): ms_net = nn.GLU(dim=0) torch_net = torch.nn.GLU(dim=0) data = np.array([[[-1, 0, 1, -2]], [[2, 2, 3, 4]] ]).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_hardshrink(): ms_net = nn.Hardshrink(lambd=0.8) torch_net = torch.nn.Hardshrink(lambd=0.8) data = np.random.randn(3, 4, 5).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.numpy(), torch_output.numpy()) assert ms_output.numpy().dtype == torch_output.numpy().dtype def test_hardsigmoid(): ms_net = nn.Hardsigmoid() torch_net = torch.nn.Hardsigmoid() data = np.array([[-4, -3, 1, 0, -1.5, 2.8, 3.7, 100]]).astype(np.float32) ms_input = Tensor(data) ms_out = ms_net(ms_input) torch_input = torch.tensor(data) torch_out = torch_net(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_multi_head_attention1(): _embed_dim = 20 _target_seq_length = 6 _batch_size = 5 _source_seq_length = 8 _num_heads = 2 ms_net = nn.MultiheadAttention(embed_dim=_embed_dim, num_heads=_num_heads, dropout=0.5) torch_net = torch.nn.MultiheadAttention(embed_dim=_embed_dim, num_heads=_num_heads, dropout=0.5) query = np.random.randn(_target_seq_length, _batch_size, _embed_dim).astype(np.float32) key = np.random.randn(_source_seq_length, _batch_size, _embed_dim).astype(np.float32) value = np.random.randn(_source_seq_length, _batch_size, _embed_dim).astype(np.float32) torch_query = torch.Tensor(query) torch_key = torch.Tensor(key) torch_value = torch.Tensor(value) torch_output = torch_net(torch_query, torch_key, torch_value, need_weights=False) ms_query = Tensor(query) ms_key = Tensor(key) ms_val = Tensor(value) ms_output = ms_net(ms_query, ms_key, ms_val, need_weights=False) assert ms_output[0].shape == torch_output[0].shape def test_multi_head_attention2(): _embed_dim = 20 _target_seq_length = 6 _batch_size = 5 _source_seq_length = 8 _num_heads = 2 ms_net = nn.MultiheadAttention(embed_dim=_embed_dim, num_heads=_num_heads, dropout=0.5, batch_first=True) torch_net = torch.nn.MultiheadAttention(embed_dim=_embed_dim, num_heads=_num_heads, dropout=0.5, batch_first=True) query = np.random.randn(_batch_size, _target_seq_length, _embed_dim).astype(np.float32) key = np.random.randn(_batch_size, _source_seq_length, _embed_dim).astype(np.float32) value = np.random.randn(_batch_size, _source_seq_length, _embed_dim).astype(np.float32) torch_query = torch.Tensor(query) torch_key = torch.Tensor(key) torch_value = torch.Tensor(value) torch_output = torch_net(torch_query, torch_key, torch_value, need_weights=False) ms_query = Tensor(query) ms_key = Tensor(key) ms_val = Tensor(value) ms_output = ms_net(ms_query, ms_key, ms_val, need_weights=False) assert ms_output[0].shape == torch_output[0].shape def test_prelu(): input = np.array([[[[0.1, 0.6], [0.9, 0.9]]]]).astype(np.float32) weight_init = 0.25 torch_input = torch.tensor(input) torch_out = torch.nn.PReLU(num_parameters=1, init=weight_init)(torch_input) ms_torch_input = ms_torch.tensor(input) ms_torch_out = ms_torch.nn.PReLU(num_parameters=1, init=weight_init)(ms_torch_input) assert np.allclose(torch_out.detach().numpy(), ms_torch_out.detach().numpy()) def test_softplus(): ms_net = nn.Softplus(beta=2, threshold=15) torch_net = torch.nn.Softplus(beta=2, threshold=15) data = np.random.randn(2, 3, 4, 5).astype(np.float32)*50 torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data) ms_output = ms_net(ms_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-5) def test_softmax2d(): ms_net = nn.Softmax2d() torch_net = torch.nn.Softmax2d() data1 = np.random.randn(2, 3, 4, 5).astype(np.float32)*50 data2 = np.random.randn(3, 4, 5).astype(np.float32)*50 torch_input1 = torch.Tensor(data1) torch_input2 = torch.Tensor(data2) torch_output1 = torch_net(torch_input1) torch_output2 = torch_net(torch_input2) ms_input1 = Tensor(data1) ms_input2 = Tensor(data2) ms_output1 = ms_net(ms_input1) ms_output2 = ms_net(ms_input2) assert np.allclose(ms_output1.asnumpy(), torch_output1.numpy()) assert np.allclose(ms_output2.asnumpy(), torch_output2.numpy()) if __name__ == '__main__': test_relu1() test_relu2() test_hardtanh1() test_hardtanh2() test_relu6() test_silu1() test_silu2() test_hardswish1() test_hardswish2() test_leakyReLU1() test_leakyReLU2() test_sigmoid() test_logsigmoid() test_elu() test_rrelu() test_selu() test_celu() test_gelu() test_mish() test_softshrink() test_tanh() test_tanhshrink() test_threshold1() test_threshold2() test_softmax1() test_softmax2() test_logsoftmax1() test_logsoftmax2() test_softmin() test_softsign() test_glu() test_hardshrink() test_multi_head_attention1() test_multi_head_attention2() test_prelu() test_softplus() test_softmax2d() msadapter/testing/ut/pytorch/nn/test_adaptive.py000066400000000000000000000054531436713554100224740ustar00rootroot00000000000000import numpy as np import torch import ms_adapter.pytorch as ms_torch from ms_adapter.pytorch.nn import AdaptiveLogSoftmaxWithLoss from mindspore import context import mindspore as ms from ms_adapter.pytorch.nn import Parameter #context.set_context(mode=ms.GRAPH_MODE) context.set_context(mode=ms.PYNATIVE_MODE) def test_adaptive_logsoftmax_withloss(): n = 50 in_fea = 200 n_class = 100 cutoffs = [10, 30, 60] seed = 1000 np.random.seed(seed) data = np.random.rand(n, in_fea).astype(np.float32)*10 target = np.random.rand(n)*n_class target = target.astype(np.int64) headweight = np.random.rand(13, 200) weight00 = np.random.rand(50, 200) weight01 = np.random.rand(20, 50) weight10 = np.random.rand(12, 200) weight11 = np.random.rand(30, 12) weight20 = np.random.rand(3, 200) weight21 = np.random.rand(40, 3) torch_net = torch.nn.AdaptiveLogSoftmaxWithLoss(in_fea, n_class, cutoffs) ms_net = AdaptiveLogSoftmaxWithLoss(in_fea, n_class, cutoffs) torch_net.head.weight = torch.nn.Parameter(torch.tensor(headweight, dtype=torch.float32)) torch_net.tail[0][0].weight = torch.nn.Parameter(torch.tensor(weight00, dtype=torch.float32)) torch_net.tail[0][1].weight = torch.nn.Parameter(torch.tensor(weight01, dtype=torch.float32)) torch_net.tail[1][0].weight = torch.nn.Parameter(torch.tensor(weight10, dtype=torch.float32)) torch_net.tail[1][1].weight = torch.nn.Parameter(torch.tensor(weight11, dtype=torch.float32)) torch_net.tail[2][0].weight = torch.nn.Parameter(torch.tensor(weight20, dtype=torch.float32)) torch_net.tail[2][1].weight = torch.nn.Parameter(torch.tensor(weight21, dtype=torch.float32)) ms_net.head.weight = Parameter(ms_torch.tensor(headweight, dtype=torch.float32)) ms_net.tail[0][0].weight = Parameter(ms_torch.tensor(weight00, dtype=torch.float32)) ms_net.tail[0][1].weight = Parameter(ms_torch.tensor(weight01, dtype=torch.float32)) ms_net.tail[1][0].weight = Parameter(ms_torch.tensor(weight10, dtype=torch.float32)) ms_net.tail[1][1].weight = Parameter(ms_torch.tensor(weight11, dtype=torch.float32)) ms_net.tail[2][0].weight = Parameter(ms_torch.tensor(weight20, dtype=torch.float32)) ms_net.tail[2][1].weight = Parameter(ms_torch.tensor(weight21, dtype=torch.float32)) torch_input = torch.tensor(data) torch_target = torch.tensor(target).long() ms_input = ms_torch.tensor(data) ms_target = ms_torch.tensor(target) torch_out, torch_loss = torch_net(torch_input, torch_target) ms_out, ms_loss = ms_net(ms_input, ms_target) # assert torch_out.shape == ms_out.shape # assert np.allclose(torch_loss.detach().numpy(), ms_loss.asnumpy()) # assert np.allclose(torch_out.detach().numpy(), ms_out.asnumpy()) if __name__ == '__main__': test_adaptive_logsoftmax_withloss() msadapter/testing/ut/pytorch/nn/test_adaptivepool.py000066400000000000000000000300011436713554100233510ustar00rootroot00000000000000import numpy as np import torch from mindspore import Tensor from ms_adapter.pytorch.nn import AdaptiveAvgPool1d, AdaptiveAvgPool2d, AdaptiveAvgPool3d from ms_adapter.pytorch.nn import AdaptiveMaxPool1d, AdaptiveMaxPool2d, AdaptiveMaxPool3d from mindspore import context import mindspore as ms context.set_context(mode=ms.GRAPH_MODE) # context.set_context(mode=ms.PYNATIVE_MODE) def test_adaptiveavgpool2d_compare1(): ms_net = AdaptiveAvgPool2d((3, 7)) torch_net = torch.nn.AdaptiveAvgPool2d((3, 7)) data = np.random.random((1, 64, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool2d_compare2(): ms_net = AdaptiveAvgPool2d(4) torch_net = torch.nn.AdaptiveAvgPool2d(4) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool2d_compare3(): ms_net = AdaptiveAvgPool2d((4, None)) torch_net = torch.nn.AdaptiveAvgPool2d((4, None)) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool2d_compare4(): ms_net = AdaptiveAvgPool2d((None, 4)) torch_net = torch.nn.AdaptiveAvgPool2d((None, 4)) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool2d_compare5(): ms_net = AdaptiveAvgPool2d((None, None)) torch_net = torch.nn.AdaptiveAvgPool2d((None, None)) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool1d_compare1(): ms_net = AdaptiveAvgPool1d(3) torch_net = torch.nn.AdaptiveAvgPool1d(3) data = np.random.random((1, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool3d_compare1(): ms_net = AdaptiveAvgPool3d((3, 4, 5)) torch_net = torch.nn.AdaptiveAvgPool3d((3, 4, 5)) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool3d_compare2(): ms_net = AdaptiveAvgPool3d(3) torch_net = torch.nn.AdaptiveAvgPool3d(3) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool3d_compare3(): ms_net = AdaptiveAvgPool3d(3) torch_net = torch.nn.AdaptiveAvgPool3d(3) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool3d_compare4(): ms_net = AdaptiveAvgPool3d((3, None, 5)) torch_net = torch.nn.AdaptiveAvgPool3d((3, None, 5)) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptiveavgpool3d_compare5(): ms_net = AdaptiveAvgPool3d((None, None, 5)) torch_net = torch.nn.AdaptiveAvgPool3d((None, None, 5)) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool2d_compare1(): ms_net = AdaptiveMaxPool2d((3, 7)) torch_net = torch.nn.AdaptiveMaxPool2d((3, 7)) data = np.random.random((1, 64, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool2d_compare2(): ms_net = AdaptiveMaxPool2d(4) torch_net = torch.nn.AdaptiveMaxPool2d(4) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool2d_compare3(): ms_net = AdaptiveMaxPool2d((4, None)) torch_net = torch.nn.AdaptiveMaxPool2d((4, None)) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool2d_compare4(): ms_net = AdaptiveMaxPool2d((None, 4)) torch_net = torch.nn.AdaptiveMaxPool2d((None, 4)) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool2d_compare5(): ms_net = AdaptiveMaxPool2d((None, None)) torch_net = torch.nn.AdaptiveMaxPool2d((None, None)) data = np.random.random((1, 3, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool2d_compare6(): ms_net = AdaptiveMaxPool2d((3, 7), return_indices=True) torch_net = torch.nn.AdaptiveMaxPool2d((3, 7), return_indices=True) data = np.random.random((1, 64, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output[0].shape, torch_output[0].shape) print(ms_output[1].shape, torch_output[1].shape) assert np.allclose(ms_output[0].shape, torch_output[0].shape) assert np.allclose(ms_output[1].shape, torch_output[1].shape) def test_adaptivemaxpool1d_compare1(): ms_net = AdaptiveMaxPool1d(3) torch_net = torch.nn.AdaptiveMaxPool1d(3) data = np.random.random((1, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool1d_compare2(): ms_net = AdaptiveMaxPool1d(3, True) torch_net = torch.nn.AdaptiveMaxPool1d(3, True) data = np.random.random((1, 10, 9)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output[0].shape, torch_output[0].shape) print(ms_output[1].shape, torch_output[1].shape) assert np.allclose(ms_output[0].shape, torch_output[0].shape) assert np.allclose(ms_output[1].shape, torch_output[1].shape) def test_adaptivemaxpool3d_compare1(): ms_net = AdaptiveMaxPool3d((3, 4, 5)) torch_net = torch.nn.AdaptiveMaxPool3d((3, 4, 5)) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool3d_compare2(): ms_net = AdaptiveMaxPool3d(3) torch_net = torch.nn.AdaptiveMaxPool3d(3) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool3d_compare3(): ms_net = AdaptiveMaxPool3d(3) torch_net = torch.nn.AdaptiveMaxPool3d(3) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool3d_compare4(): ms_net = AdaptiveMaxPool3d((3, None, 5)) torch_net = torch.nn.AdaptiveMaxPool3d((3, None, 5)) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) def test_adaptivemaxpool3d_compare5(): ms_net = AdaptiveMaxPool3d((None, None, 5)) torch_net = torch.nn.AdaptiveMaxPool3d((None, None, 5)) data = np.random.random((1, 3, 10, 9, 12)) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) print(ms_output.shape, torch_output.shape) assert np.allclose(ms_output.shape, torch_output.shape) if __name__ == '__main__': test_adaptiveavgpool2d_compare1() test_adaptiveavgpool2d_compare2() test_adaptiveavgpool2d_compare3() test_adaptiveavgpool2d_compare4() test_adaptiveavgpool2d_compare5() test_adaptiveavgpool1d_compare1() test_adaptiveavgpool3d_compare1() test_adaptiveavgpool3d_compare2() test_adaptiveavgpool3d_compare3() test_adaptiveavgpool3d_compare4() test_adaptiveavgpool3d_compare5() test_adaptivemaxpool2d_compare1() test_adaptivemaxpool2d_compare2() test_adaptivemaxpool2d_compare3() test_adaptivemaxpool2d_compare4() test_adaptivemaxpool2d_compare5() test_adaptivemaxpool2d_compare6() test_adaptivemaxpool1d_compare1() test_adaptivemaxpool1d_compare2() test_adaptivemaxpool3d_compare1() test_adaptivemaxpool3d_compare2() test_adaptivemaxpool3d_compare3() test_adaptivemaxpool3d_compare4() test_adaptivemaxpool3d_compare5()msadapter/testing/ut/pytorch/nn/test_batchnorm.py000066400000000000000000000416741436713554100226610ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch from mindspore import context import mindspore as ms import ms_adapter from ms_adapter.pytorch.nn import Module from ms_adapter.pytorch.nn import BatchNorm1d, BatchNorm2d, BatchNorm3d from ms_adapter.pytorch.nn import SyncBatchNorm, InstanceNorm1d, InstanceNorm2d, InstanceNorm3d from ms_adapter.pytorch import tensor context.set_context(mode=ms.GRAPH_MODE) def test_bn(): class BnModel1d(Module): def __init__(self): super(BnModel1d, self).__init__() self.bn1 = BatchNorm1d(num_features=32) self.bn2 = BatchNorm1d(32, affine=False) def forward(self, inputs): x = self.bn1(inputs) x = self.bn2(x) return x class BnModel2d(Module): def __init__(self): super(BnModel2d, self).__init__() self.bn1 = BatchNorm2d(num_features=32) self.bn2 = BatchNorm2d(32, affine=False) def forward(self, inputs): x = self.bn1(inputs) x = self.bn2(x) return x class BnModel3d(Module): def __init__(self): super(BnModel3d, self).__init__() self.bn1 = BatchNorm3d(num_features=32) self.bn2 = BatchNorm3d(32, affine=False) def forward(self, inputs): x = self.bn1(inputs) x = self.bn2(x) return x model1d = BnModel1d() model2d = BnModel2d() model3d = BnModel3d() inputs1d = tensor(np.ones(shape=(5, 32)), ms.float32) output1d = model1d(inputs1d) assert output1d.shape == (5, 32) inputs2d = tensor(np.ones(shape=(5, 32, 5, 5)), ms.float32) output2d = model2d(inputs2d) assert output2d.shape == (5, 32, 5, 5) inputs3d = tensor(np.ones(shape=(5, 32, 5, 5, 5)), ms.float32) output3d = model3d(inputs3d) assert output3d.shape == (5, 32, 5, 5, 5) class MSBnModel(Module): def __init__(self, mode, num_features, eps, momentum, affine, track_running_stats): super(MSBnModel, self).__init__() if mode == "2d": self.bn = BatchNorm2d(num_features, eps, momentum, affine, track_running_stats) elif mode == "3d": self.bn = BatchNorm3d(num_features, eps, momentum, affine, track_running_stats) else: self.bn = BatchNorm1d(num_features, eps, momentum, affine, track_running_stats) def forward(self, inputs): x = self.bn(inputs) return x class PTBnModel(torch.nn.Module): def __init__(self, mode, num_features, eps, momentum, affine, track_running_stats): super(PTBnModel, self).__init__() if mode == "2d": self.bn = torch.nn.BatchNorm2d(num_features, eps, momentum, affine, track_running_stats) elif mode == "3d": self.bn = torch.nn.BatchNorm3d(num_features, eps, momentum, affine, track_running_stats) else: self.bn = torch.nn.BatchNorm1d(num_features, eps, momentum, affine, track_running_stats) def forward(self, inputs): x = self.bn(inputs) return x def test_bn1d_momentum(): input = np.arange(5 * 3).reshape(5, 3).astype(np.float32) pt_model = PTBnModel("1d", 3, 1e-5, 0.1, True, True) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("1d", 3, 1e-5, 0.1, True, True) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn1d_affine(): input = np.arange(5 * 3).reshape(5, 3).astype(np.float32) pt_model = PTBnModel("1d", 3, 1e-5, 0.1, False, True) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("1d", 3, 1e-5, 0.1, False, True) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn1d_train(): input = np.arange(5 * 3).reshape(5, 3).astype(np.float32) pt_model = PTBnModel("1d", 3, 1e-5, 0.1, False, True) pt_model.train(mode=False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("1d", 3, 1e-5, 0.1, False, True) ms_model.train(mode=False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) np.random.seed(123) input2 = np.random.randn(5, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn1d_track_running_stats(): input = np.arange(5 * 3).reshape(5, 3).astype(np.float32) pt_model = PTBnModel("1d", 3, 1e-5, 0.5, True, False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("1d", 3, 1e-5, 0.5, True, False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) np.random.seed(123) input2 = np.random.randn(5, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) def test_bn1d_train_and_track_running_stats(): input = np.arange(5 * 3).reshape(5, 3).astype(np.float32) pt_model = PTBnModel("1d", 3, 1e-5, 0.5, True, False) pt_model.train(mode=False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("1d", 3, 1e-5, 0.5, True, False) ms_model.train(mode=False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) np.random.seed(123) input2 = np.random.randn(5, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) def test_bn2d_momentum(): input = np.arange(5 * 3 * 3 * 3).reshape(5, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("2d", 3, 1e-5, 0.1, True, True) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("2d", 3, 1e-5, 0.1, True, True) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn2d_affine(): input = np.arange(5 * 3 * 3 * 3).reshape(5, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("2d", 3, 1e-5, 0.1, False, True) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("2d", 3, 1e-5, 0.1, False, True) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn2d_train(): input = np.arange(5 * 3 * 3 * 3).reshape(5, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("2d", 3, 1e-5, 0.1, False, True) pt_model.train(mode=False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("2d", 3, 1e-5, 0.1, False, True) ms_model.train(mode=False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) np.random.seed(123) input2 = np.random.randn(5, 3, 3, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn2d_track_running_stats(): input = np.arange(5 * 3 * 3 * 3).reshape(5, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("2d", 3, 1e-5, 0.1, False, False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("2d", 3, 1e-5, 0.1, False, False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) np.random.seed(123) input2 = np.random.randn(5, 3, 3, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) def test_bn2d_train_and_track_running_stats(): input = np.arange(5 * 3 * 3 * 3).reshape(5, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("2d", 3, 1e-5, 0.1, True, False) pt_model.train(mode=False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("2d", 3, 1e-5, 0.1, True, False) ms_model.train(mode=False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) np.random.seed(123) input2 = np.random.randn(5, 3, 3, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) def test_bn3d_momentum(): input = np.arange(5 * 3 * 3 * 3 * 3).reshape(5, 3, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("3d", 3, 1e-5, 0.1, True, True) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("3d", 3, 1e-5, 0.1, True, True) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn3d_affine(): input = np.arange(5 * 3 * 3 * 3 * 3).reshape(5, 3, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("3d", 3, 1e-5, 0.1, False, True) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("3d", 3, 1e-5, 0.1, False, True) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn3d_train(): input = np.arange(5 * 3 * 3 * 3 * 3).reshape(5, 3, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("3d", 3, 1e-5, 0.1, False, True) pt_model.train(mode=False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("3d", 3, 1e-5, 0.1, False, True) ms_model.train(mode=False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) np.random.seed(123) input2 = np.random.randn(5, 3, 3, 3, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) assert np.allclose(ms_model.bn.running_mean.asnumpy(), pt_model.bn.running_mean.numpy()) def test_bn3d_track_running_stats(): input = np.arange(5 * 3 * 3 * 3 * 3).reshape(5, 3, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("3d", 3, 1e-5, 0.1, False, False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("3d", 3, 1e-5, 0.1, False, False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) np.random.seed(123) input2 = np.random.randn(5, 3, 3, 3, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) def test_bn3d_train_and_track_running_stats(): input = np.arange(5 * 3 * 3 * 3 * 3).reshape(5, 3, 3, 3, 3).astype(np.float32) pt_model = PTBnModel("3d", 3, 1e-5, 0.9, False, False) pt_model.train(mode=False) pt_input = torch.Tensor(input) pt_output = pt_model(pt_input) ms_model = MSBnModel("3d", 3, 1e-5, 0.9, False, False) ms_model.train(mode=False) ms_inputs = ms_adapter.pytorch.Tensor(input) ms_output = ms_model(ms_inputs) assert np.allclose(ms_output.asnumpy(), pt_output.detach().numpy(), atol=1e-5) np.random.seed(123) input2 = np.random.randn(5, 3, 3, 3, 3).astype(np.float32) pt_input2 = torch.Tensor(input2) ms_inputs2 = ms_adapter.pytorch.Tensor(input2) pt_output2 = pt_model(pt_input2) ms_output2 = ms_model(ms_inputs2) assert np.allclose(ms_output2.asnumpy(), pt_output2.detach().numpy(), atol=1e-5) # def test_sbn_track_runing_status(): # input = np.arange(5 * 3 * 3 * 3).reshape(5, 3, 3, 3).astype(np.float32) # pt_input = torch.Tensor(input) # ms_inputs = ms_adapter.pytorch.Tensor(input) # # pt_model = torch.nn.SyncBatchNorm(num_features=10) # # from mindspore.communication import init # init() # ms.reset_auto_parallel_context() # ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.DATA_PARALLEL) # ms_model = SyncBatchNorm(num_features=10) # # pt_output = pt_model(pt_input) # ms_output = ms_model(ms_inputs) # print(ms_output.asnumpy()) # print(pt_output.detach().numpy()) # # # def test_insn1d_track_runing_status(): # input = np.arange(5 * 3 * 3 ).reshape(5, 3, 3).astype(np.float32) # pt_input = torch.Tensor(input) # ms_inputs = ms_adapter.pytorch.Tensor(input) # # pt_model = torch.nn.InstanceNorm1d(num_features=3) # ms_model = InstanceNorm1d(num_features=3) # # pt_output = pt_model(pt_input) # ms_output = ms_model(ms_inputs) # print(ms_output.asnumpy()) # print(pt_output.detach().numpy()) # # def test_insn2d_track_runing_status(): # input = np.arange(5 * 3 * 3 *3).reshape(5, 3, 3, 3).astype(np.float32) # pt_input = torch.Tensor(input) # ms_inputs = ms_adapter.pytorch.Tensor(input) # # pt_model = torch.nn.InstanceNorm2d(num_features=3) # ms_model = InstanceNorm2d(num_features=3) # # pt_output = pt_model(pt_input) # ms_output = ms_model(ms_inputs) # print(ms_output.asnumpy()) # print(pt_output.detach().numpy()) # # # def test_insn3d_track_runing_status(): # input = np.arange(5 * 3 * 3 * 3 * 3).reshape(5, 3, 3, 3, 3).astype(np.float32) # pt_input = torch.Tensor(input) # ms_inputs = ms_adapter.pytorch.Tensor(input) # # pt_model = torch.nn.InstanceNorm3d(num_features=3) # ms_model = InstanceNorm3d(num_features=3) # # pt_output = pt_model(pt_input) # ms_output = ms_model(ms_inputs) # print(ms_output.asnumpy()) # print(pt_output.detach().numpy()) if __name__ == '__main__': test_bn() test_bn1d_momentum() test_bn1d_affine() test_bn1d_train() test_bn1d_track_running_stats() test_bn1d_train_and_track_running_stats() test_bn2d_momentum() test_bn2d_affine() test_bn2d_train() test_bn2d_track_running_stats() test_bn2d_train_and_track_running_stats() test_bn3d_momentum() test_bn3d_affine() test_bn3d_train() test_bn3d_track_running_stats() test_bn3d_train_and_track_running_stats() # test_sbn_track_runing_status() # Only Ascend is supported. # test_insn1d_track_runing_status() # Only GPU is supported. # test_insn2d_track_runing_status() # Only GPU is supported. # test_insn3d_track_runing_status() # Only GPU is supported. msadapter/testing/ut/pytorch/nn/test_channelshuffle.py000066400000000000000000000022331436713554100236550ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.GRAPH_MODE) def test_channelshuffle(): np_data1 = np.random.randn(2, 9, 3, 4) np_data2 = np.random.randn(1, 6, 12, 2, 2) torch_input1 = torch.tensor(np_data1) torch_input2 = torch.tensor(np_data2) ms_input1 = ms_torch.tensor(np_data1) ms_input2 = ms_torch.tensor(np_data2) torch_shuffle = torch.nn.ChannelShuffle(3) ms_shuffle = ms_torch.nn.ChannelShuffle(3) torch_out1 = torch_shuffle(torch_input1) torch_out2 = torch_shuffle(torch_input2) ms_out1 = ms_shuffle(ms_input1) ms_out2 = ms_shuffle(ms_input2) assert np.allclose(torch_out1.detach().numpy(), ms_out1.numpy()) assert torch_out1.detach().numpy().dtype == ms_out1.numpy().dtype assert torch_out1.shape == ms_out1.shape assert np.allclose(torch_out2.detach().numpy(), ms_out2.numpy()) assert torch_out2.detach().numpy().dtype == ms_out2.numpy().dtype assert torch_out2.shape == ms_out2.shape if __name__ == '__main__': test_channelshuffle()msadapter/testing/ut/pytorch/nn/test_contianer.py000066400000000000000000000116131436713554100226540ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.PYNATIVE_MODE) def test_module_dict(): class MyTorchModule(torch.nn.Module): def __init__(self): super(MyTorchModule, self).__init__() self.choices = torch.nn.ModuleDict({ 'conv': torch.nn.Conv2d(10, 10, 3), 'pool': torch.nn.MaxPool2d(3) }) self.activations = torch.nn.ModuleDict([ ['lrelu', torch.nn.LeakyReLU()], ['prelu', torch.nn.PReLU()] ]) def forward(self, x, choice, act): x = self.choices[choice](x) x = self.activations[act](x) return x class MyMsModule(ms_torch.nn.Module): def __init__(self): super(MyMsModule, self).__init__() self.choices = ms_torch.nn.ModuleDict({ 'conv': ms_torch.nn.Conv2d(10, 10, 3), 'pool': ms_torch.nn.MaxPool2d(3) }) self.activations = ms_torch.nn.ModuleDict([ ['lrelu', ms_torch.nn.LeakyReLU()], ['prelu', ms_torch.nn.PReLU()] ]) def forward(self, x, choice, act): x = self.choices[choice](x) x = self.activations[act](x) return x x = np.random.randint(0, 10, [1, 2, 4, 4]) torch_net = MyTorchModule() torch_out = torch_net(torch.tensor(x, dtype=torch.float32), 'pool', 'lrelu') ms_torch_net = MyMsModule() ms_torch_out = ms_torch_net(ms_torch.tensor(x, dtype=ms_torch.float32), 'pool', 'lrelu') assert np.allclose(torch_out.detach().numpy(), ms_torch_out.detach().numpy()) def test_parameter_list(): init_data = np.random.randn(10, 10) class MyTorchModule(torch.nn.Module): def __init__(self): super(MyTorchModule, self).__init__() self.params = torch.nn.ParameterList([torch.nn.Parameter(torch.tensor(init_data)) for i in range(10)]) def forward(self, x): # ParameterList can act as an iterable, or be indexed using ints for i, p in enumerate(self.params): x = self.params[i // 2].mm(x) + p.mm(x) return x class MyMsModule(ms_torch.nn.Module): def __init__(self): super(MyMsModule, self).__init__() self.params = ms_torch.nn.ParameterList([ms_torch.nn.Parameter(ms_torch.tensor(init_data)) for i in range(10)]) def forward(self, x): # ParameterList can act as an iterable, or be indexed using ints for i, p in enumerate(self.params): x = self.params[i // 2].mm(x) + p.mm(x) return x x = np.random.randn(10, 1) torch_net = MyTorchModule() ms_torch_net = MyMsModule() torch_net.params.append(torch.nn.Parameter(torch.tensor(init_data))) torch_net.params.extend([torch.nn.Parameter(torch.tensor(init_data))]) torch_out = torch_net(torch.tensor(x)) ms_torch_net.params.append(ms_torch.nn.Parameter(ms_torch.tensor(init_data))) ms_torch_net.params.extend([ms_torch.nn.Parameter(ms_torch.tensor(init_data))]) ms_torch_out = ms_torch_net(ms_torch.tensor(x)) assert np.allclose(torch_out.detach().numpy(), ms_torch_out.detach().numpy()) def test_parameter_dict(): init_data1 = np.random.randn(5, 10) init_data2 = np.random.randn(5, 10) class MyTorchModule(torch.nn.Module): def __init__(self): super(MyTorchModule, self).__init__() self.params = torch.nn.ParameterDict({ 'left': torch.nn.Parameter(torch.tensor(init_data1)), 'right': torch.nn.Parameter(torch.tensor(init_data2)) }) self.params.update({'left': torch.nn.Parameter(torch.tensor(init_data2))}) def forward(self, x, choice): x = self.params[choice].mm(x) return x class MyMsModule(ms_torch.nn.Module): def __init__(self): super(MyMsModule, self).__init__() self.params = ms_torch.nn.ParameterDict({ 'left': ms_torch.nn.Parameter(ms_torch.tensor(init_data1)), 'right': ms_torch.nn.Parameter(ms_torch.tensor(init_data2)) }) self.params.update({'left': ms_torch.nn.Parameter(ms_torch.tensor(init_data2))}) def forward(self, x, choice): x = self.params[choice].mm(x) return x x = np.random.randn(10, 1) torch_net = MyTorchModule() ms_torch_net = MyMsModule() torch_out = torch_net(torch.tensor(x), 'right') ms_torch_out = ms_torch_net(ms_torch.tensor(x), 'right') assert np.allclose(torch_out.detach().numpy(), ms_torch_out.detach().numpy()) if __name__ == '__main__': test_module_dict() test_parameter_list() test_parameter_dict() msadapter/testing/ut/pytorch/nn/test_conv.py000066400000000000000000000567341436713554100216540ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch from mindspore import context import mindspore as ms import mindspore.nn as nn import ms_adapter.pytorch as ms_pytorch from ms_adapter.pytorch.nn import Module, Parameter from ms_adapter.pytorch.nn import Conv1d, Conv2d, Conv3d from ms_adapter.pytorch.nn import ConvTranspose1d, ConvTranspose2d, ConvTranspose3d from ms_adapter.pytorch import tensor np.random.seed(0) context.set_context(mode=ms.PYNATIVE_MODE) def test_torch_ms_conv1d_padding(): input_init = np.random.randn(1, 3, 16).astype(float) weight_init = np.random.randn(64, 3, 3).astype(float) weight_init_t = np.expand_dims(weight_init, axis=2) bias_init = np.random.randn(64).astype(float) class Conv1dPadModel(torch.nn.Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv1dPadModel, self).__init__() self.conv = torch.nn.Conv1d(in_channels=3, out_channels=64, kernel_size=3, stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = torch.nn.Parameter(torch.tensor(weight_init, dtype=torch.float32)) self.conv.bias = torch.nn.Parameter(torch.tensor(bias_init, dtype=torch.float32)) def forward(self, inputs): x = self.conv(inputs) return x class Conv1dPadModelMs(Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv1dPadModelMs, self).__init__() self.conv = Conv1d(in_channels=3, out_channels=64, kernel_size=3, stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = Parameter(tensor(weight_init_t, ms_pytorch.float32)) self.conv.bias = Parameter(tensor(bias_init, ms_pytorch.float32)) def forward(self, inputs): x = self.conv(inputs) return x py_input = torch.tensor(input_init, dtype=torch.float32) ms_input = tensor(input_init, ms_pytorch.float32) # padding = 'same' py_net1 = Conv1dPadModel(padding='same') ms_net1 = Conv1dPadModelMs(padding='same') py_output1 = py_net1(py_input) ms_output1 = ms_net1(ms_input) assert(py_output1.shape == ms_output1.shape) assert np.allclose(py_output1.detach().numpy(), ms_output1.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'same', dilation = (1, 2) py_net2 = Conv1dPadModel(padding='same', dilation=1) ms_net2 = Conv1dPadModelMs(padding='same', dilation=1) py_output2 = py_net2(py_input) ms_output2 = ms_net2(ms_input) assert(py_output2.shape == ms_output2.shape) assert np.allclose(py_output2.detach().numpy(), ms_output2.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'valid' py_net3 = Conv1dPadModel(padding='valid') ms_net3 = Conv1dPadModelMs(padding='valid') py_output3 = py_net3(py_input) ms_output3 = ms_net3(ms_input) assert(py_output3.shape == ms_output3.shape) assert np.allclose(py_output3.detach().numpy(), ms_output3.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'valid', dilation = (1, 2) py_net4 = Conv1dPadModel(padding='valid', dilation=1) ms_net4 = Conv1dPadModelMs(padding='valid', dilation=1) py_output4 = py_net4(py_input) ms_output4 = ms_net4(ms_input) assert(py_output4.shape == ms_output4.shape) assert np.allclose(py_output4.detach().numpy(), ms_output4.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 1, padding_mode='zeros' py_net5 = Conv1dPadModel(padding = 1, padding_mode='zeros') ms_net5 = Conv1dPadModelMs(padding = 1, padding_mode='zeros') py_output5 = py_net5(py_input) ms_output5 = ms_net5(ms_input) assert(py_output5.shape == ms_output5.shape) assert np.allclose(py_output5.detach().numpy(), ms_output5.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = (2,2), padding_mode='zeros' py_net6 = Conv1dPadModel(padding = 2, padding_mode='zeros') ms_net6 = Conv1dPadModelMs(padding = 2, padding_mode='zeros') py_output6 = py_net6(py_input) ms_output6 = ms_net6(ms_input) assert(py_output6.shape == ms_output6.shape) assert np.allclose(py_output6.detach().numpy(), ms_output6.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 def test_torch_ms_conv2d_network(): """Test torch and ms conv cell output.""" weight1_init = np.random.randn(16, 3, 3, 3).astype(float) weight2_init = np.random.randn(32, 16, 1, 1).astype(float) weight3_init = np.random.randn(16, 32, 1, 1).astype(float) weight4_init = np.random.randn(1, 16, 3, 3).astype(float) bias1_init = np.random.randn(16).astype(float) bias3_init = np.random.randn(16).astype(float) bias4_init = np.random.randn(1).astype(float) class ConvModel(torch.nn.Module): def __init__(self): super(ConvModel, self).__init__() self.conv1 = torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3) self.conv2 = torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=1, padding=2, bias=False) self.conv3 = torch.nn.Conv2d(in_channels=32, out_channels=16, kernel_size=1, padding=6) self.conv4 = torch.nn.Conv2d(in_channels=16, out_channels=1, kernel_size=3, padding='same') self.init_param() def init_param(self): self.conv1.weight = torch.nn.Parameter(torch.tensor(weight1_init, dtype=torch.float32)) self.conv1.bias = torch.nn.Parameter(torch.tensor(bias1_init, dtype=torch.float32)) self.conv2.weight = torch.nn.Parameter(torch.tensor(weight2_init, dtype=torch.float32)) self.conv3.weight = torch.nn.Parameter(torch.tensor(weight3_init, dtype=torch.float32)) self.conv3.bias = torch.nn.Parameter(torch.tensor(bias3_init, dtype=torch.float32)) self.conv4.weight = torch.nn.Parameter(torch.tensor(weight4_init, dtype=torch.float32)) self.conv4.bias = torch.nn.Parameter(torch.tensor(bias4_init, dtype=torch.float32)) def forward(self, inputs): x = self.conv1(inputs) x = self.conv2(x) x = self.conv3(x) x = self.conv4(x) return x class ConvModelms(Module): def __init__(self): super(ConvModelms, self).__init__() self.conv1 = Conv2d(in_channels=3, out_channels=16, kernel_size=3) self.conv2 = Conv2d(in_channels=16, out_channels=32, kernel_size=1, padding=2, bias=False) self.conv3 = Conv2d(in_channels=32, out_channels=16, kernel_size=1, padding=6) self.conv4 = Conv2d(in_channels=16, out_channels=1, kernel_size=3, padding='same') self.init_param() def init_param(self): self.conv1.weight = Parameter(tensor(weight1_init, ms_pytorch.float32)) self.conv1.bias = Parameter(tensor(bias1_init, ms_pytorch.float32)) self.conv2.weight = Parameter(tensor(weight2_init, ms_pytorch.float32)) self.conv3.weight = Parameter(tensor(weight3_init, ms_pytorch.float32)) self.conv3.bias = Parameter(tensor(bias3_init, ms_pytorch.float32)) self.conv4.weight = Parameter(tensor(weight4_init, ms_pytorch.float32)) self.conv4.bias = Parameter(tensor(bias4_init, ms_pytorch.float32)) def forward(self, inputs): x = self.conv1(inputs) x = self.conv2(x) x = self.conv3(x) x = self.conv4(x) return x py_net = ConvModel() py_input = torch.tensor(np.ones(shape=(1, 3, 32, 32)), dtype=torch.float32) py_output = py_net(py_input) ms_net = ConvModelms() ms_net.train() ms_input = tensor(np.ones(shape=(1, 3, 32, 32)), ms_pytorch.float32) ms_output = ms_net(ms_input) assert (py_output.shape == ms_output.shape) assert np.allclose(py_output.detach().numpy(), ms_output.asnumpy(), atol=1e-3) def test_torch_ms_conv2d_padding(): input_init = np.random.randn(1, 3, 16, 50).astype(float) weight_init = np.random.randn(64, 3, 3, 5).astype(float) bias_init = np.random.randn(64).astype(float) class Conv2dPadModel(torch.nn.Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv2dPadModel, self).__init__() self.conv = torch.nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(3, 5), stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = torch.nn.Parameter(torch.tensor(weight_init, dtype=torch.float32)) self.conv.bias = torch.nn.Parameter(torch.tensor(bias_init, dtype=torch.float32)) def forward(self, inputs): x = self.conv(inputs) return x class Conv2dPadModelMs(Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv2dPadModelMs, self).__init__() self.conv = Conv2d(in_channels=3, out_channels=64, kernel_size=(3, 5), stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = Parameter(tensor(weight_init, ms_pytorch.float32)) self.conv.bias = Parameter(tensor(bias_init, ms_pytorch.float32)) def forward(self, inputs): x = self.conv(inputs) return x py_input = torch.tensor(input_init, dtype=torch.float32) ms_input = tensor(input_init, ms_pytorch.float32) # padding = 'same' py_net1 = Conv2dPadModel(padding='same') ms_net1 = Conv2dPadModelMs(padding='same') py_output1 = py_net1(py_input) ms_output1 = ms_net1(ms_input) assert(py_output1.shape == ms_output1.shape) assert np.allclose(py_output1.detach().numpy(), ms_output1.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'same', dilation = (1, 2) py_net2 = Conv2dPadModel(padding='same', dilation=(1, 2)) ms_net2 = Conv2dPadModelMs(padding='same', dilation=(1, 2)) py_output2 = py_net2(py_input) ms_output2 = ms_net2(ms_input) assert(py_output2.shape == ms_output2.shape) assert np.allclose(py_output2.detach().numpy(), ms_output2.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'valid' py_net3 = Conv2dPadModel(padding='valid') ms_net3 = Conv2dPadModelMs(padding='valid') py_output3 = py_net3(py_input) ms_output3 = ms_net3(ms_input) assert(py_output3.shape == ms_output3.shape) assert np.allclose(py_output3.detach().numpy(), ms_output3.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'valid', dilation = (1, 2) py_net4 = Conv2dPadModel(padding='valid', dilation=(1, 2)) ms_net4 = Conv2dPadModelMs(padding='valid', dilation=(1, 2)) py_output4 = py_net4(py_input) ms_output4 = ms_net4(ms_input) assert(py_output4.shape == ms_output4.shape) assert np.allclose(py_output4.detach().numpy(), ms_output4.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 1, padding_mode='zeros' py_net5 = Conv2dPadModel(padding = 1, padding_mode='zeros') ms_net5 = Conv2dPadModelMs(padding = 1, padding_mode='zeros') py_output5 = py_net5(py_input) ms_output5 = ms_net5(ms_input) assert(py_output5.shape == ms_output5.shape) assert np.allclose(py_output5.detach().numpy(), ms_output5.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = (2,2), padding_mode='zeros' py_net6 = Conv2dPadModel(padding = (2, 2), padding_mode='zeros') ms_net6 = Conv2dPadModelMs(padding = (2, 2), padding_mode='zeros') py_output6 = py_net6(py_input) ms_output6 = ms_net6(ms_input) assert(py_output6.shape == ms_output6.shape) assert np.allclose(py_output6.detach().numpy(), ms_output6.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 def test_torch_ms_conv3d_padding(): input_init = np.random.randn(1, 3, 50, 50, 50).astype(float) weight_init = np.random.randn(64, 3, 3, 3, 3).astype(float) bias_init = np.random.randn(64).astype(float) class Conv3dPadModel(torch.nn.Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv3dPadModel, self).__init__() self.conv = torch.nn.Conv3d(in_channels=3, out_channels=64, kernel_size=(3, 3, 3), stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = torch.nn.Parameter(torch.tensor(weight_init, dtype=torch.float32)) self.conv.bias = torch.nn.Parameter(torch.tensor(bias_init, dtype=torch.float32)) def forward(self, inputs): x = self.conv(inputs) return x class Conv3dPadModelMs(Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv3dPadModelMs, self).__init__() self.conv = Conv3d(in_channels=3, out_channels=64, kernel_size=(3, 3, 3), stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = Parameter(tensor(weight_init, ms_pytorch.float32)) self.conv.bias = Parameter(tensor(bias_init, ms_pytorch.float32)) def forward(self, inputs): x = self.conv(inputs) return x py_input = torch.tensor(input_init, dtype=torch.float32) ms_input = tensor(input_init, ms_pytorch.float32) # padding = 'same' py_net1 = Conv3dPadModel(padding='same') ms_net1 = Conv3dPadModelMs(padding='same') py_output1 = py_net1(py_input) ms_output1 = ms_net1(ms_input) assert(py_output1.shape == ms_output1.shape) assert np.allclose(py_output1.detach().numpy(), ms_output1.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'same', dilation = (1, 2) py_net2 = Conv3dPadModel(padding='same', dilation=(1, 2, 2)) ms_net2 = Conv3dPadModelMs(padding='same', dilation=(1, 2, 2)) py_output2 = py_net2(py_input) ms_output2 = ms_net2(ms_input) assert(py_output2.shape == ms_output2.shape) assert np.allclose(py_output2.detach().numpy(), ms_output2.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'valid' py_net3 = Conv3dPadModel(padding='valid') ms_net3 = Conv3dPadModelMs(padding='valid') py_output3 = py_net3(py_input) ms_output3 = ms_net3(ms_input) assert(py_output3.shape == ms_output3.shape) assert np.allclose(py_output3.detach().numpy(), ms_output3.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 'valid', dilation = (1, 2) py_net4 = Conv3dPadModel(padding='valid', dilation=(1, 2, 2)) ms_net4 = Conv3dPadModelMs(padding='valid', dilation=(1, 2, 2)) py_output4 = py_net4(py_input) ms_output4 = ms_net4(ms_input) assert(py_output4.shape == ms_output4.shape) assert np.allclose(py_output4.detach().numpy(), ms_output4.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = 1, padding_mode='zeros' py_net5 = Conv3dPadModel(padding = 1, padding_mode='zeros') ms_net5 = Conv3dPadModelMs(padding = 1, padding_mode='zeros') py_output5 = py_net5(py_input) ms_output5 = ms_net5(ms_input) assert(py_output5.shape == ms_output5.shape) assert np.allclose(py_output5.detach().numpy(), ms_output5.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 # padding = (2,2), padding_mode='zeros' py_net6 = Conv3dPadModel(padding = (2, 2, 2), padding_mode='zeros') ms_net6 = Conv3dPadModelMs(padding = (2, 2, 2), padding_mode='zeros') py_output6 = py_net6(py_input) ms_output6 = ms_net6(ms_input) assert(py_output6.shape == ms_output6.shape) assert np.allclose(py_output6.detach().numpy(), ms_output6.asnumpy(), atol=1e-05) #cpu上 atol>=1e-05 def test_torch_ms_deconv1d_padding(): weight_init = np.random.randn(3, 10, 3).astype(float) weight_init_t = np.expand_dims(weight_init, axis=2) bias_init = np.random.randn(10).astype(float) class DConv1dPadModel(torch.nn.Module): def __init__(self, stride=1, padding=0, output_padding=1): super(DConv1dPadModel, self).__init__() self.dconv = torch.nn.ConvTranspose1d(in_channels=3, out_channels=10, kernel_size=3, stride=stride, padding=padding, output_padding=output_padding) self.dconv.weight = torch.nn.Parameter(torch.tensor(weight_init, dtype=torch.float32)) self.dconv.bias = torch.nn.Parameter(torch.tensor(bias_init, dtype=torch.float32)) def forward(self, inputs): x = self.dconv(inputs) return x class DConv1dPadModelMs(Module): def __init__(self, stride=1, padding=0, output_padding=1): super(DConv1dPadModelMs, self).__init__() self.dconv = ConvTranspose1d(in_channels=3, out_channels=10, kernel_size=3, stride=stride, padding=padding, output_padding=output_padding) self.dconv.weight = Parameter(tensor(weight_init_t, ms_pytorch.float32)) self.dconv.bias = Parameter(tensor(bias_init, ms_pytorch.float32)) def forward(self, inputs): x = self.dconv(inputs) return x py_in = torch.ones(size=(10, 3, 10), dtype=torch.float32) py_net1 = DConv1dPadModel(stride=2, padding=0, output_padding=0) py_out1 = py_net1(py_in) ms_in = tensor(data=np.ones(shape=(10, 3, 10), dtype=np.float32)) ms_net1 = DConv1dPadModelMs(stride=2, padding=0, output_padding=0) ms_out1 = ms_net1(ms_in) assert (ms_out1.shape == py_out1.shape) assert np.allclose(py_out1.detach().numpy(), ms_out1.asnumpy(), atol=1e-05) py_net2 = DConv1dPadModel(stride=2, padding=2, output_padding=0) py_out2 = py_net2(py_in) ms_net2 = DConv1dPadModelMs(stride=2, padding=2, output_padding=0) ms_out2 = ms_net2(ms_in) assert (ms_out2.shape == py_out2.shape) assert np.allclose(py_out2.detach().numpy(), ms_out2.asnumpy(), atol=1e-05) py_net3 = DConv1dPadModel(stride=3, padding=2, output_padding=0) py_out3 = py_net3(py_in) ms_net3 = DConv1dPadModelMs(stride=3, padding=2, output_padding=0) ms_out3 = ms_net3(ms_in) assert (ms_out3.shape == py_out3.shape) assert np.allclose(py_out3.detach().numpy(), ms_out3.asnumpy(), atol=1e-05) def test_torch_ms_deconv2d_padding(): weight_init = np.random.randn(3, 10, 3, 3).astype(float) bias_init = np.random.randn(10).astype(float) class DConv2dPadModel(torch.nn.Module): def __init__(self, stride=(1, 1), padding=(0, 0), output_padding=(1, 1)): super(DConv2dPadModel, self).__init__() self.dconv = torch.nn.ConvTranspose2d(in_channels=3, out_channels=10, kernel_size=(3, 3), stride=stride, padding=padding, output_padding=output_padding) self.dconv.weight = torch.nn.Parameter(torch.tensor(weight_init, dtype=torch.float32)) self.dconv.bias = torch.nn.Parameter(torch.tensor(bias_init, dtype=torch.float32)) def forward(self, inputs): x = self.dconv(inputs) return x class DConv2dPadModelMs(Module): def __init__(self, stride=(1, 1), padding=(0, 0), output_padding=(1, 1)): super(DConv2dPadModelMs, self).__init__() self.dconv = ConvTranspose2d(in_channels=3, out_channels=10, kernel_size=(3, 3), stride=stride, padding=padding, output_padding=output_padding) self.dconv.weight = Parameter(tensor(weight_init, ms_pytorch.float32)) self.dconv.bias = Parameter(tensor(bias_init, ms_pytorch.float32)) def forward(self, inputs): x = self.dconv(inputs) return x py_in = torch.ones(size=(10, 3, 10, 10), dtype=torch.float32) py_net1 = DConv2dPadModel(stride=2, padding=0, output_padding=0) py_out1 = py_net1(py_in) ms_in = tensor(data=np.ones(shape=(10, 3, 10, 10), dtype=np.float32)) ms_net1 = DConv2dPadModelMs(stride=2, padding=0, output_padding=0) ms_out1 = ms_net1(ms_in) assert (ms_out1.shape == py_out1.shape) assert np.allclose(py_out1.detach().numpy(), ms_out1.asnumpy(), atol=1e-05) py_net2 = DConv2dPadModel(stride=2, padding=2, output_padding=0) py_out2 = py_net2(py_in) ms_net2 = DConv2dPadModelMs(stride=2, padding=2, output_padding=0) ms_out2 = ms_net2(ms_in) assert (ms_out2.shape == py_out2.shape) assert np.allclose(py_out2.detach().numpy(), ms_out2.asnumpy(), atol=1e-05) py_net3 = DConv2dPadModel(stride=3, padding=2, output_padding=0) py_out3 = py_net3(py_in) ms_net3 = DConv2dPadModelMs(stride=3, padding=2, output_padding=0) ms_out3 = ms_net3(ms_in) assert (ms_out3.shape == py_out3.shape) assert np.allclose(py_out3.detach().numpy(), ms_out3.asnumpy(), atol=1e-05) # def test_torch_ms_deconv3d_padding(): # weight_init = np.random.randn(3, 10, 3, 3, 3).astype(float) # bias_init = np.random.randn(10).astype(float) # # class DConv3dPadModel(torch.nn.Module): # def __init__(self, stride=(1, 1, 1), padding=(0, 0, 0), output_padding=(1, 1, 1)): # super(DConv3dPadModel, self).__init__() # self.dconv = torch.nn.ConvTranspose3d(in_channels=3, out_channels=10, kernel_size=(3, 3, 3), stride=stride, # padding=padding, output_padding=output_padding) # self.dconv.weight = torch.nn.Parameter(torch.tensor(weight_init, dtype=torch.float32)) # self.dconv.bias = torch.nn.Parameter(torch.tensor(bias_init, dtype=torch.float32)) # # def forward(self, inputs): # x = self.dconv(inputs) # return x # # class DConv3dPadModelMs(Module): # def __init__(self, stride=(1, 1, 1), padding=(0, 0, 0), output_padding=(1, 1, 1)): # super(DConv3dPadModelMs, self).__init__() # self.dconv = ConvTranspose3d(in_channels=3, out_channels=10, kernel_size=(3, 3, 3), stride=stride, # padding=padding, output_padding=output_padding) # self.dconv.weight = Parameter(tensor(weight_init, ms_pytorch.float32)) # self.dconv.bias = Parameter(tensor(bias_init, ms_pytorch.float32)) # # def forward(self, inputs): # x = self.dconv(inputs) # return x # # # py_in = torch.ones(size=(10, 3, 10, 10, 10), dtype=torch.float32) # py_net1 = DConv3dPadModel(stride=2, padding=0, output_padding=0) # py_out1 = py_net1(py_in) # # ms_in = tensor(data=np.ones(shape=(10, 3, 10, 10, 10), dtype=np.float32)) # ms_net1 = DConv3dPadModelMs(stride=2, padding=0, output_padding=0) # ms_out1 = ms_net1(ms_in) # assert (ms_out1.shape == py_out1.shape) # assert np.allclose(py_out1.detach().numpy(), ms_out1.asnumpy(), atol=1e-05) # # py_net2 = DConv3dPadModel(stride=2, padding=2, output_padding=0) # py_out2 = py_net2(py_in) # # ms_net2 = DConv3dPadModelMs(stride=2, padding=2, output_padding=0) # ms_out2 = ms_net2(ms_in) # assert (ms_out2.shape == py_out2.shape) # assert np.allclose(py_out2.detach().numpy(), ms_out2.asnumpy(), atol=1e-05) # # py_net3 = DConv3dPadModel(stride=3, padding=2, output_padding=0) # py_out3 = py_net3(py_in) # # ms_net3 = DConv3dPadModelMs(stride=3, padding=2, output_padding=0) # ms_out3 = ms_net3(ms_in) # assert (ms_out3.shape == py_out3.shape) # assert np.allclose(py_out3.detach().numpy(), ms_out3.asnumpy(), atol=1e-05) if __name__ == '__main__': test_torch_ms_conv1d_padding() test_torch_ms_conv2d_network() test_torch_ms_conv2d_padding() test_torch_ms_conv3d_padding() test_torch_ms_deconv1d_padding() test_torch_ms_deconv2d_padding() # test_torch_ms_deconv3d_padding() # TODO Unsupport CPUmsadapter/testing/ut/pytorch/nn/test_distance.py000066400000000000000000000120721436713554100224640ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch import Tensor import ms_adapter.pytorch.nn as nn import numpy as np from mindspore import context import mindspore as ms import torch context.set_context(mode=ms.GRAPH_MODE) def test_pairwise_distance1(): ms_distance = nn.PairwiseDistance(keepdim=True) torch_distance = torch.nn.PairwiseDistance(keepdim=True) data1 = np.array([[[-1, 0, 1, -2], [2, 2, 3, 4] ]]).astype(np.float32) data2 = np.array([[-1, 5, 1, 6], [-2, 12, -7, 1.4] ]).astype(np.float32) torch_input1 = torch.Tensor(data1) torch_input2 = torch.Tensor(data2) torch_out = torch_distance(torch_input1, torch_input2) ms_input1 = Tensor(data1.astype(np.float32)) ms_input2 = Tensor(data2.astype(np.float32)) ms_out = ms_distance(ms_input1, ms_input2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pairwise_distance2(): ms_distance = nn.PairwiseDistance(p=-2, eps=10) torch_distance = torch.nn.PairwiseDistance(p=-2, eps=10) data1 = np.array([[-2, -2], [-2, -2] ]).astype(np.float32) data2 = np.array([[8, 8], [8, 8] ]).astype(np.float32) torch_input1 = torch.Tensor(data1) torch_input2 = torch.Tensor(data2) torch_out = torch_distance(torch_input1, torch_input2) ms_input1 = Tensor(data1.astype(np.float32)) ms_input2 = Tensor(data2.astype(np.float32)) ms_out = ms_distance(ms_input1, ms_input2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_pairwise_distance3(): ms_distance = nn.PairwiseDistance(p=2.1, eps=1) torch_distance = torch.nn.PairwiseDistance(p=2.1, eps=1) data1 = np.array([[[[0]], [[0]], [[0]]]]).astype(np.float32) data2 = np.array([[[[2, -1], [0, 0]], [[2, -1], [0, 0]], [[2, -1], [0, 0]]]]).astype(np.float32) torch_input1 = torch.Tensor(data1) torch_input2 = torch.Tensor(data2) torch_out = torch_distance(torch_input1, torch_input2) ms_input1 = Tensor(data1.astype(np.float32)) ms_input2 = Tensor(data2.astype(np.float32)) ms_out = ms_distance(ms_input1, ms_input2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_cosine_similarity1(): ms_distance = nn.CosineSimilarity(dim=1, eps=2) torch_distance = torch.nn.CosineSimilarity(dim=1, eps=2) data1 = np.array([[[1], [-1]]]).astype(np.float32) data2 = np.array([[1, 1], [2, 1]]).astype(np.float32) torch_input1 = torch.Tensor(data1) torch_input2 = torch.Tensor(data2) torch_out = torch_distance(torch_input1, torch_input2) ms_input1 = Tensor(data1.astype(np.float32)) ms_input2 = Tensor(data2.astype(np.float32)) ms_out = ms_distance(ms_input1, ms_input2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_cosine_similarity2(): ms_distance = nn.CosineSimilarity(eps=20) torch_distance = torch.nn.CosineSimilarity(eps=20) data1 = np.array([[-1, 0, 1, -2], [2, 2, 3, 4]]).astype(np.float32) data2 = np.array([[-1, 5, 1, 6], [-2, 12, -7, 1.4]]).astype(np.float32) torch_input1 = torch.Tensor(data1) torch_input2 = torch.Tensor(data2) torch_out = torch_distance(torch_input1, torch_input2) ms_input1 = Tensor(data1.astype(np.float32)) ms_input2 = Tensor(data2.astype(np.float32)) ms_out = ms_distance(ms_input1, ms_input2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_cosine_similarity3(): ms_distance = nn.CosineSimilarity(dim=0) torch_distance = torch.nn.CosineSimilarity(dim=0) data1 = np.array([[[[1]], [[-1]], [[2]]]]).astype(np.float32) data2 = np.array([[[1, 1], [2, 1]], [[1, 3], [2, -0.5]], [[0, 0], [-1, 2]]]).astype(np.float32) torch_input1 = torch.Tensor(data1) torch_input2 = torch.Tensor(data2) torch_out = torch_distance(torch_input1, torch_input2) ms_input1 = Tensor(data1.astype(np.float32)) ms_input2 = Tensor(data2.astype(np.float32)) ms_out = ms_distance(ms_input1, ms_input2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_pairwise_distance1() test_pairwise_distance2() test_pairwise_distance3() test_cosine_similarity1() test_cosine_similarity2() test_cosine_similarity3() msadapter/testing/ut/pytorch/nn/test_dropout.py000066400000000000000000000126751436713554100223770ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import ms_adapter from ms_adapter.pytorch import tensor from mindspore import context import mindspore as ms import numpy as np import torch context.set_context(mode=ms.PYNATIVE_MODE) def _get_error_with_count_nonzero(ms_tensor, torch_tensor): nonzero_ms_output = np.count_nonzero(ms_tensor.asnumpy()) nonzero_torch_output = np.count_nonzero(torch_tensor.numpy()) if nonzero_ms_output > nonzero_torch_output: error = nonzero_ms_output - nonzero_torch_output else: error = nonzero_torch_output - nonzero_ms_output return error def test_dropout_compare1(): input = np.random.random((5, 5)).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.Dropout(p=0.5) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.Dropout(p=0.5) ms_output = ms_dropout(ms_tensor) error = _get_error_with_count_nonzero(ms_output, torch_output) kErrorThreshold = 5 # assert error < kErrorThreshold def test_dropout_compare2(): input = np.random.random((5, 5)).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.Dropout(p=0.2, inplace=True) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.Dropout(p=0.2, inplace=True) ms_output = ms_dropout(ms_tensor) error_input = _get_error_with_count_nonzero(ms_tensor, torch_tensor) error_output = _get_error_with_count_nonzero(ms_output, torch_output) kErrorThreshold = 5 assert error_input == error_output # assert error_output < kErrorThreshold def test_dropout2d_compare1(): input = np.random.random((2, 25, 1, 1)).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.Dropout2d(p=0.5) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.Dropout2d(p=0.5) ms_output = ms_dropout(ms_tensor) error = _get_error_with_count_nonzero(ms_output, torch_output) kErrorThreshold = 10 # assert error < kErrorThreshold def test_dropout2d_compare2(): input = np.random.random((2, 25, 1, 1)).astype(np.float32) torch_tensor = torch.tensor(input) ms_tensor = tensor(input) torch_dropout = torch.nn.Dropout2d(p=0.5, inplace=True) torch_output = torch_dropout(torch_tensor) ms_dropout = ms_adapter.pytorch.nn.Dropout2d(p=0.5, inplace=True) ms_output = ms_dropout(ms_tensor) error_input = _get_error_with_count_nonzero(ms_tensor, torch_tensor) error_output = _get_error_with_count_nonzero(ms_output, torch_output) kErrorThreshold = 10 assert error_input == error_output # assert error_output < kErrorThreshold def test_dropout3d_compare1(): input = np.random.random((2, 25, 1, 1, 1)).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.Dropout3d(p=0.5) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.Dropout3d(p=0.5) ms_output = ms_dropout(ms_tensor) error = _get_error_with_count_nonzero(ms_output, torch_output) kErrorThreshold = 10 # assert error < kErrorThreshold def test_dropout3d_compare2(): input = np.random.random((2, 25, 1, 1, 1)).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.Dropout3d(inplace=True) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.Dropout3d(inplace=True) ms_output = ms_dropout(ms_tensor) error_input = _get_error_with_count_nonzero(ms_tensor, torch_tensor) error_output = _get_error_with_count_nonzero(ms_output, torch_output) kErrorThreshold = 10 assert error_input == error_output # assert error_output < kErrorThreshold def test_alphadropout(): input = np.random.random([2, 3, 2, 4]).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.AlphaDropout(p=0.4, inplace=True) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.AlphaDropout(p=0.4, inplace=True) ms_output = ms_dropout(ms_tensor) assert np.allclose(ms_tensor.asnumpy(), ms_output.numpy()) def test_featurealphadropout(): input = np.random.random([2, 3, 2, 4]).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.FeatureAlphaDropout(p=0.4, inplace=True) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.FeatureAlphaDropout(p=0.4, inplace=True) ms_output = ms_dropout(ms_tensor) assert np.allclose(ms_tensor.asnumpy(), ms_output.numpy()) def test_dropout1d(): input = np.random.random([2, 3, 4]).astype(np.float32) torch_tensor = torch.tensor(input) torch_dropout = torch.nn.Dropout1d(p=0.4, inplace=True) torch_output = torch_dropout(torch_tensor) ms_tensor = tensor(input) ms_dropout = ms_adapter.pytorch.nn.Dropout1d(p=0.4, inplace=True) ms_output = ms_dropout(ms_tensor) assert np.allclose(ms_tensor.asnumpy(), ms_output.numpy()) if __name__ == '__main__': test_dropout_compare1() test_dropout_compare2() test_dropout2d_compare1() test_dropout2d_compare2() test_dropout3d_compare1() test_dropout3d_compare2() test_alphadropout() test_featurealphadropout() test_dropout1d() msadapter/testing/ut/pytorch/nn/test_flatten.py000066400000000000000000000044351436713554100223330ustar00rootroot00000000000000import numpy as np import torch from mindspore import context import mindspore as ms import ms_adapter from ms_adapter.pytorch import tensor def test_flatten_pynative(): context.set_context(mode=ms.PYNATIVE_MODE) input = np.array([[[1.2, 2.1], [2.2, 3.2], [2.0, 3.0]]]).astype(np.float32) torch_tensor = torch.tensor(input) torch_flatten = torch.nn.Flatten() torch_output = torch_flatten(torch_tensor) ms_tensor = tensor(input) ms_flatten = ms_adapter.pytorch.nn.Flatten() ms_output = ms_flatten(ms_tensor) assert torch_output.shape == ms_output.shape assert np.all(torch_output.numpy() == ms_output.asnumpy()) def test_flatten_graph(): context.set_context(mode=ms.GRAPH_MODE) input = np.array([[[1.2, 2.1], [2.2, 3.2], [2.0, 3.0]]]).astype(np.float32) torch_tensor = torch.tensor(input) torch_flatten = torch.nn.Flatten() torch_output = torch_flatten(torch_tensor) ms_tensor = tensor(input) ms_flatten = ms_adapter.pytorch.nn.Flatten() ms_output = ms_flatten(ms_tensor) assert torch_output.shape == ms_output.shape assert np.all(torch_output.numpy() == ms_output.asnumpy()) def test_unflatten_pynative(): context.set_context(mode=ms.PYNATIVE_MODE) input = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(input) torch_flatten = torch.nn.Unflatten(1, (1, 3)) torch_output = torch_flatten(torch_tensor) ms_tensor = tensor(input) ms_flatten = ms_adapter.pytorch.nn.Unflatten(1, (1, 3)) ms_output = ms_flatten(ms_tensor) assert torch_output.shape == ms_output.shape assert np.allclose(torch_output.numpy(), ms_output.asnumpy()) def test_unflatten_graph(): context.set_context(mode=ms.GRAPH_MODE) input = np.random.randn(2, 3, 4, 5).astype(np.float32) torch_tensor = torch.tensor(input) torch_flatten = torch.nn.Unflatten(1, (1, 3)) torch_output = torch_flatten(torch_tensor) ms_tensor = tensor(input) ms_flatten = ms_adapter.pytorch.nn.Unflatten(1, (1, 3)) ms_output = ms_flatten(ms_tensor) assert torch_output.shape == ms_output.shape assert np.allclose(torch_output.numpy(), ms_output.asnumpy()) if __name__ == '__main__': test_flatten_pynative() test_flatten_graph() test_unflatten_pynative() test_unflatten_graph() msadapter/testing/ut/pytorch/nn/test_fold.py000066400000000000000000000035341436713554100216210ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context def test_unfold(): np_input = np.random.randn(7, 8, 9, 10) torch_tensor = torch.tensor(np_input) torch_unfold = torch.nn.Unfold((2, 3), 1, (1, 2), (2, 1)) torch_out = torch_unfold(torch_tensor) ms_tensor = ms_torch.tensor(np_input) ms_unfold = ms_torch.nn.Unfold((2, 3), 1, (1, 2), (2, 1)) ms_out = ms_unfold(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), rtol=1e-3, atol=1e-5) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_fold(): np_input1 = np.random.randn(7, 8, 24) np_input2 = np.random.randn(18, 6) torch_tensor1 = torch.tensor(np_input1) torch_tensor2 = torch.tensor(np_input2) torch_fold1 = torch.nn.Fold((4, 5), (2, 2), 1, (1, 2), (2, 1)) torch_fold2 = torch.nn.Fold((7, 4), 3, 2, 2, 3) torch_out1 = torch_fold1(torch_tensor1) torch_out2 = torch_fold2(torch_tensor2) ms_tensor1 = ms_torch.tensor(np_input1) ms_tensor2 = ms_torch.tensor(np_input2) ms_fold1 = ms_torch.nn.Fold((4, 5), (2, 2), 1, (1, 2), (2, 1)) ms_fold2 = ms_torch.nn.Fold((7, 4), 3, 2, 2, 3) ms_out1 = ms_fold1(ms_tensor1) ms_out2 = ms_fold2(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy(), rtol=1e-3, atol=1e-5) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert ms_out1.asnumpy().shape == torch_out1.numpy().shape assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy(), rtol=1e-3, atol=1e-5) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert ms_out2.asnumpy().shape == torch_out2.numpy().shape if __name__ == '__main__': test_unfold() test_fold() msadapter/testing/ut/pytorch/nn/test_linear.py000066400000000000000000000055311436713554100221460ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch.nn import Module, Linear, Identity, Bilinear from ms_adapter.pytorch import tensor from mindspore import context import numpy as np import mindspore as ms import torch context.set_context(mode=ms.PYNATIVE_MODE) def test_linear_model(): context.set_context(mode=ms.PYNATIVE_MODE) class LinearModel(Module): def __init__(self): super(LinearModel, self).__init__() self.line1 = Linear(in_features=32, out_features=64) self.line2 = Linear(in_features=64, out_features=128, bias=False) self.line3 = Linear(in_features=128, out_features=10) def forward(self, inputs): x = self.line1(inputs) x = self.line2(x) x = self.line3(x) return x model = LinearModel() model.train() def weight_init(m): if isinstance(m, Linear): m.weight.data.normal_(0, 0.01) if m.has_bias: m.bias.data.zero_() model.apply(weight_init) inputs = tensor(np.ones(shape=(5, 32)), torch.float32) output = model(inputs) assert output.shape == (5, 10) class IdentityModel(Module): def __init__(self): super(IdentityModel, self).__init__() self.identity = Identity() def forward(self, inputs): return self.identity(inputs) def test_identity(): context.set_context(mode=context.GRAPH_MODE) net = IdentityModel() input_tensor = tensor(np.ones(shape=(5, 32)), torch.float32) output = net(input_tensor) np.testing.assert_almost_equal(output.asnumpy(), input_tensor.asnumpy()) assert id(input_tensor) != id(output) def test_bilinear_model(): context.set_context(mode=ms.PYNATIVE_MODE) class BilinearModel(Module): def __init__(self): super(BilinearModel, self).__init__() self.line1 = Bilinear(in1_features=4, in2_features=3, out_features=5) self.line2 = Bilinear(in1_features=5, in2_features=5, out_features=7, bias=False) def forward(self, inputs1, inputs2): x = self.line1(inputs1, inputs2) x = self.line2(x, x) return x model = BilinearModel() model.train() def weight_init(m): if isinstance(m, Bilinear): m.weight.data.normal_(0, 0.01) if m.has_bias: m.bias.data.zero_() model.apply(weight_init) inputs1 = tensor(np.ones(shape=(10, 4)), torch.float32) inputs2 = tensor(np.ones(shape=(10, 3)), torch.float32) output = model(inputs1, inputs2) assert output.shape == (10, 7) def test_linear_model2(): linear = Linear(64, 3) x = tensor(np.ones((1, 2, 64))) assert linear(x).shape == (1, 2, 3) if __name__ == '__main__': test_linear_model() test_identity() test_bilinear_model() test_linear_model2() msadapter/testing/ut/pytorch/nn/test_loss.py000066400000000000000000001001731436713554100216520ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch import mindspore as ms from mindspore import context, float32 import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.PYNATIVE_MODE) def test_smoothl1loss1(): a = np.array([1, 2, 3]).astype(np.float32) b = np.array([1, 2, 2]).astype(np.float32) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.SmoothL1Loss()(ms_logits, ms_labels) torch_logits = torch.tensor(a) torch_labels = torch.tensor(b) result_torch = torch.nn.SmoothL1Loss()(torch_logits, torch_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_smoothl1loss2(): a = np.array([1, 2, 3]).astype(np.float32) b = np.array([1, 2, 2]).astype(np.float32) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='sum', beta=0.5)(ms_logits, ms_labels) torch_logits = torch.tensor(a) torch_labels = torch.tensor(b) result_torch = torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='sum', beta=0.5)(torch_logits, torch_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_smoothl1loss3(): a = np.array([1, 2, 3]).astype(np.float32) b = np.array([1, 2, 2]).astype(np.float32) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='none', beta=0.5)(ms_logits, ms_labels) torch_logits = torch.tensor(a) torch_labels = torch.tensor(b) result_torch = torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='none', beta=0.5)(torch_logits, torch_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_l1loss_reduction_none(): a = np.random.randn(2).astype(float) b = np.random.randn(2).astype(float) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.L1Loss(size_average=None, reduce=None, reduction='none')(ms_logits, ms_labels) torch_input = torch.tensor(a) torch_target = torch.tensor(b) result_torch = torch.nn.L1Loss(size_average=None, reduce=None, reduction='none')(torch_input, torch_target) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_l1loss_reduction_sum(): a = np.random.randn(2, 3).astype(float) b = np.random.randn(2, 3).astype(float) torch_input = torch.tensor(a) torch_target = torch.tensor(b) result_torch = torch.nn.L1Loss(size_average=None, reduce=None, reduction='sum')(torch_input, torch_target) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.L1Loss(size_average=None, reduce=None, reduction='sum')(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_l1loss_reduction_mean(): a = np.random.randn(2, 2, 3).astype(float) b = np.random.randn(2, 2, 3).astype(float) torch_input = torch.tensor(a) torch_target = torch.tensor(b) result_torch = torch.nn.L1Loss(size_average=None, reduce=None, reduction='mean')(torch_input, torch_target) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.L1Loss(size_average=None, reduce=None, reduction='mean')(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_mseloss_reduction_none(): a = np.random.randn(2).astype(float) b = np.random.randn(2).astype(float) torch_input = torch.tensor(a) torch_target = torch.tensor(b) result_torch = torch.nn.MSELoss(size_average=None, reduce=None, reduction='none')(torch_input, torch_target) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.MSELoss(size_average=None, reduce=None, reduction='none')(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_mseloss_reduction_sum(): context.set_context(mode=ms.PYNATIVE_MODE) a = np.random.randn(2, 3).astype(float) b = np.random.randn(2, 3).astype(float) torch_input = torch.tensor(a) torch_target = torch.tensor(b) result_torch = torch.nn.MSELoss(size_average=None, reduce=None, reduction='sum')(torch_input, torch_target) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.MSELoss(size_average=None, reduce=None, reduction='sum')(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_mseloss_reduction_mean(): context.set_context(mode=ms.PYNATIVE_MODE) a = np.random.randn(2, 2, 3).astype(float) b = np.random.randn(2, 2, 3).astype(float) torch_input = torch.tensor(a) torch_target = torch.tensor(b) result_torch = torch.nn.MSELoss(size_average=None, reduce=None, reduction='mean')(torch_input, torch_target) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.MSELoss(size_average=None, reduce=None, reduction='mean')(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_cross_entropy_reduction_none(): a = np.random.randn(3, 5).astype(float) b = np.array([1, 0, 4]) c = np.array([0.1, 0.2, 0.3, 0.4, 0.5]) torch_input = torch.tensor(a, dtype=torch.float32) torch_target = torch.tensor(b, dtype=torch.int64) torch_weight = torch.tensor(c, dtype=torch.float32) torch_loss = torch.nn.CrossEntropyLoss(weight=torch_weight, reduction='none') result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(a, ms_torch.float32) ms_labels = ms_torch.tensor(b, ms_torch.int64) ms_weight = ms_torch.tensor(c, ms_torch.float32) ms_loss = ms_torch.nn.CrossEntropyLoss(weight=ms_weight, reduction='none') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_cross_entropy_reduction_sum(): a = np.random.randn(3, 5).astype(float) b = np.random.randn(3, 5).astype(float) torch_input = torch.tensor(a, dtype=torch.float32) torch_target = torch.tensor(b, dtype=torch.float32) torch_loss = torch.nn.CrossEntropyLoss(weight=None, reduction='sum', label_smoothing=1.0) result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(a, ms_torch.float32) ms_labels = ms_torch.tensor(b, ms_torch.float32) ms_loss = ms_torch.nn.CrossEntropyLoss(weight=None, reduction='sum', label_smoothing=1.0) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_cross_entropy_reduction_mean(): a = np.random.randn(3, 5).astype(float) b = np.array([1, 0, 4]) torch_input = torch.tensor(a, dtype=torch.float32) torch_target = torch.tensor(b, dtype=torch.int64) torch_loss = torch.nn.CrossEntropyLoss(weight=None, ignore_index=-1, reduction='mean', label_smoothing=0.1) result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(a, ms_torch.float32) ms_labels = ms_torch.tensor(b, ms_torch.int64) ms_loss = ms_torch.nn.CrossEntropyLoss(weight=None, ignore_index=-1, reduction='mean', label_smoothing=0.1) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_nllloss_reduction(): a = np.array([[0.1, 0.2, 0.3], [0.5, 0.7, 0.9], [0.5, 0.7, 0.9]]) b = np.array([0, 1, 0]) c = np.array([1.0, 2.0, 3.0]) torch_input = torch.tensor(a, dtype=torch.float32) torch_target = torch.tensor(b, dtype=torch.int64) torch_weight = torch.tensor(c, dtype=torch.float32) ms_logits = ms_torch.tensor(a, ms_torch.float32) ms_labels = ms_torch.tensor(b, ms_torch.int64) ms_weight = ms_torch.tensor(c, ms_torch.float32) torch_loss = torch.nn.NLLLoss(weight=None, ignore_index=-1, reduction='mean') result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.NLLLoss(weight=None, ignore_index=-1, reduction='mean') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape torch_loss = torch.nn.NLLLoss(weight=torch_weight, reduction='sum') result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.NLLLoss(weight=ms_weight, reduction='sum') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape torch_loss = torch.nn.NLLLoss(weight=torch_weight, reduction='none') result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.NLLLoss(weight=ms_weight, reduction='none') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_kl_div_loss_reduction(): a = np.array([0.2, 0.7, 0.1]) b = np.array([0., 1., 0.]) torch_input = torch.tensor(a) torch_target = torch.tensor(b) torch_loss = torch.nn.KLDivLoss(reduction='mean', log_target=False) result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) ms_loss = ms_torch.nn.KLDivLoss(reduction='mean', log_target=False) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_bce_loss(): a = np.array([[0.1, 0.2, 0.3], [0.5, 0.7, 0.9]]) b = np.array([[0, 1, 0], [0, 0, 1]]) c = np.array([[1.0, 2.0, 3.0], [4.0, 3.3, 2.2]]) torch_input = torch.tensor(a, dtype=torch.float32) torch_target = torch.tensor(b, dtype=torch.float32) torch_weight = torch.tensor(c, dtype=torch.float32) ms_logits = ms_torch.tensor(a, ms_torch.float32) ms_labels = ms_torch.tensor(b, ms_torch.float32) ms_weight = ms_torch.tensor(c, ms_torch.float32) torch_loss = torch.nn.BCELoss(weight=None, reduction='none') result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.BCELoss(weight=None, reduction='none') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape torch_loss = torch.nn.BCELoss(weight=torch_weight, reduction='sum') result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.BCELoss(weight=ms_weight, reduction='sum') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape torch_loss = torch.nn.BCELoss(weight=torch_weight, reduction='mean') result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.BCELoss(weight=ms_weight, reduction='mean') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_binary_cross_entropy_with_logits(): a = np.array([[-0.8, 1.2, 0.7], [-0.1, -0.4, 0.7]]) b = np.array([[0.3, 0.8, 1.2], [-0.6, 0.1, 2.2]]) c = np.array([1.0, 1.0, 1.0]) d = np.array([0.1, 0.1, 0.1]) torch_input = torch.tensor(a, dtype=torch.float32) torch_target = torch.tensor(b, dtype=torch.float32) torch_weight = torch.tensor(c, dtype=torch.float32) torch_pos_weight = torch.tensor(d, dtype=torch.float32) ms_logits = ms_torch.tensor(a, ms_torch.float32) ms_labels = ms_torch.tensor(b, ms_torch.float32) ms_weight = ms_torch.tensor(c, ms_torch.float32) ms_pos_weight = ms_torch.tensor(d, ms_torch.float32) torch_loss = torch.nn.BCEWithLogitsLoss(weight=torch_weight, reduction='none', pos_weight=torch_pos_weight) result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.BCEWithLogitsLoss(weight=ms_weight, reduction='none', pos_weight=ms_pos_weight) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape torch_loss = torch.nn.BCEWithLogitsLoss(weight=torch_weight, reduction='sum', pos_weight=None) result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.BCEWithLogitsLoss(weight=ms_weight, reduction='sum', pos_weight=None) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape torch_loss = torch.nn.BCEWithLogitsLoss(weight=None, reduction='mean', pos_weight=torch_pos_weight) result_torch = torch_loss(torch_input, torch_target) ms_loss = ms_torch.nn.BCEWithLogitsLoss(weight=None, reduction='mean', pos_weight=ms_pos_weight) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.shape == result_torch.shape def test_huber_loss_mean(): input = np.random.randn(3, 5).astype(np.float64) target = np.random.randn(3, 5).astype(np.float64) torch_input = torch.tensor(input) torch_target = torch.tensor(target) torch_loss = torch.nn.HuberLoss(reduction='mean', delta=1.0) result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(input) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.HuberLoss(reduction='mean', delta=1.0) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_huber_loss_sum(): input = np.random.randn(3, 5).astype(np.float32) target = np.random.randn(3, 5).astype(np.float32) torch_input = torch.tensor(input) torch_target = torch.tensor(target) torch_loss = torch.nn.HuberLoss(reduction='sum', delta=1.0) result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(input) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.HuberLoss(reduction='sum', delta=1.0) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_huber_loss_none(): input = np.random.randn(3, 5).astype(np.float32) target = np.random.randn(3, 5).astype(np.float32) torch_input = torch.tensor(input) torch_target = torch.tensor(target) torch_loss = torch.nn.HuberLoss(reduction='none', delta=1.0) result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(input) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.HuberLoss(reduction='none', delta=1.0) result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape #TODO #mindspore.ops.soft_margin_loss not support cpu until now ''' def test_soft_margin_loss_mean(): input = np.random.randn(3, 5).astype(np.float32) target = np.random.randn(3, 5).astype(np.float32) torch_input = torch.tensor(input) torch_target = torch.tensor(target) torch_loss = torch.nn.SoftMarginLoss(reduction='mean') result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(input) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.SoftMarginLoss(reduction='mean') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) def test_soft_margin_loss_sum(): input = np.random.randn(3, 5).astype(np.float32) target = np.random.randn(3, 5).astype(np.float32) torch_input = torch.tensor(input) torch_target = torch.tensor(target) torch_loss = torch.nn.SoftMarginLoss(reduction='sum') result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(input) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.SoftMarginLoss(reduction='sum') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) def test_soft_margin_loss_none(): input = np.random.randn(3, 5).astype(np.float32) target = np.random.randn(3, 5).astype(np.float32) torch_input = torch.tensor(input) torch_target = torch.tensor(target) torch_loss = torch.nn.SoftMarginLoss(reduction='none') result_torch = torch_loss(torch_input, torch_target) ms_logits = ms_torch.tensor(input) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.SoftMarginLoss(reduction='none') result_ms = ms_loss(ms_logits, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) ''' def test_cosine_embedding_loss_none(): input1 = np.random.randn(3, 5).astype(np.float32) input2 = np.random.randn(3, 5).astype(np.float32) target = np.array([1, -1, 1]).astype(np.float32) torch_input1 = torch.tensor(input1) torch_input2 = torch.tensor(input2) torch_target = torch.tensor(target) torch_loss = torch.nn.CosineEmbeddingLoss(reduction='none', margin=0.5) result_torch = torch_loss(torch_input1, torch_input2, torch_target) ms_logits1 = ms_torch.tensor(input1) ms_logits2 = ms_torch.tensor(input2) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.CosineEmbeddingLoss(reduction='none', margin=0.5) result_ms = ms_loss(ms_logits1, ms_logits2, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_cosine_embedding_loss_sum(): input1 = np.random.randn(3, 5).astype(np.float32) input2 = np.random.randn(3, 5).astype(np.float32) target = np.array([1, -1, 1]).astype(np.float32) torch_input1 = torch.tensor(input1) torch_input2 = torch.tensor(input2) torch_target = torch.tensor(target) torch_loss = torch.nn.CosineEmbeddingLoss(reduction='sum', margin=0.5) result_torch = torch_loss(torch_input1, torch_input2, torch_target) ms_logits1 = ms_torch.tensor(input1) ms_logits2 = ms_torch.tensor(input2) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.CosineEmbeddingLoss(reduction='sum', margin=0.5) result_ms = ms_loss(ms_logits1, ms_logits2, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_cosine_embedding_loss_mean(): input1 = np.random.randn(3, 5).astype(np.float32) input2 = np.random.randn(3, 5).astype(np.float32) target = np.array([1, -1, 1]).astype(np.float32) torch_input1 = torch.tensor(input1) torch_input2 = torch.tensor(input2) torch_target = torch.tensor(target) torch_loss = torch.nn.CosineEmbeddingLoss(reduction='mean', margin=0.5) result_torch = torch_loss(torch_input1, torch_input2, torch_target) ms_logits1 = ms_torch.tensor(input1) ms_logits2 = ms_torch.tensor(input2) ms_labels = ms_torch.tensor(target) ms_loss = ms_torch.nn.CosineEmbeddingLoss(reduction='mean', margin=0.5) result_ms = ms_loss(ms_logits1, ms_logits2, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_cosine_triplet_margin_loss_none(): anchor = np.random.randn(100, 128).astype(np.float32) positive = np.random.randn(100, 128).astype(np.float32) negative = np.random.randn(100, 128).astype(np.float32) torch_input1 = torch.tensor(anchor) torch_input2 = torch.tensor(positive) torch_target = torch.tensor(negative) torch_triplet_loss = torch.nn.TripletMarginLoss(margin=1.0, p=2, reduction='none') result_torch = torch_triplet_loss(torch_input1, torch_input2, torch_target) ms_logits1 = ms_torch.tensor(anchor) ms_logits2 = ms_torch.tensor(positive) ms_labels = ms_torch.tensor(negative) ms_loss = ms_torch.nn.TripletMarginLoss(margin=1.0, p=2, reduction='none') result_ms = ms_loss(ms_logits1, ms_logits2, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy(), atol=1e-5) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_cosine_triplet_margin_loss_sum(): anchor = np.random.randn(100, 128).astype(np.float32) positive = np.random.randn(100, 128).astype(np.float32) negative = np.random.randn(100, 128).astype(np.float32) torch_input1 = torch.tensor(anchor) torch_input2 = torch.tensor(positive) torch_target = torch.tensor(negative) torch_triplet_loss = torch.nn.TripletMarginLoss(margin=1.0, p=2, reduction='sum') result_torch = torch_triplet_loss(torch_input1, torch_input2, torch_target) ms_logits1 = ms_torch.tensor(anchor) ms_logits2 = ms_torch.tensor(positive) ms_labels = ms_torch.tensor(negative) ms_loss = ms_torch.nn.TripletMarginLoss(margin=1.0, p=2, reduction='sum') result_ms = ms_loss(ms_logits1, ms_logits2, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_cosine_triplet_margin_loss_mean(): anchor = np.random.randn(100, 128).astype(np.float32) positive = np.random.randn(100, 128).astype(np.float32) negative = np.random.randn(100, 128).astype(np.float32) torch_input1 = torch.tensor(anchor) torch_input2 = torch.tensor(positive) torch_target = torch.tensor(negative) torch_triplet_loss = torch.nn.TripletMarginLoss(margin=1.0, p=2) result_torch = torch_triplet_loss(torch_input1, torch_input2, torch_target) ms_logits1 = ms_torch.tensor(anchor) ms_logits2 = ms_torch.tensor(positive) ms_labels = ms_torch.tensor(negative) ms_loss = ms_torch.nn.TripletMarginLoss(margin=1.0, p=2) result_ms = ms_loss(ms_logits1, ms_logits2, ms_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_multi_margin_loss_none(): x = np.array([[0.1, 0.2, 0.4, 0.8]]) y = np.array([3]) torch_input1 = torch.tensor(x) torch_input2 = torch.tensor(y, dtype=torch.int64) torch_loss = torch.nn.MultiMarginLoss(p=2, reduction='none') result_torch = torch_loss(torch_input1, torch_input2) ms_logits1 = ms_torch.tensor(x) ms_logits2 = ms_torch.tensor(y, dtype=ms.int64) ms_loss = ms_torch.nn.MultiMarginLoss(p=2, reduction='none') result_ms = ms_loss(ms_logits1, ms_logits2) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_multi_margin_loss_weight(): x = np.array([[0.1, 0.2, 0.4, 0.8]]) y = np.array([3]) weight = np.array([0.2, 0.3, 0.4, 0.1]) torch_input1 = torch.tensor(x) torch_input2 = torch.tensor(y, dtype=torch.int64) torch_weight = torch.tensor(weight) torch_loss = torch.nn.MultiMarginLoss(weight=torch_weight) result_torch = torch_loss(torch_input1, torch_input2) ms_logits1 = ms_torch.tensor(x) ms_logits2 = ms_torch.tensor(y, dtype=ms.int64) ms_weight = ms_torch.tensor(weight) ms_loss = ms_torch.nn.MultiMarginLoss(weight=ms_weight) result_ms = ms_loss(ms_logits1, ms_logits2) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) assert result_ms.asnumpy().dtype == result_torch.numpy().dtype assert result_ms.shape == result_torch.shape def test_poisson_nll_loss(): np_data = np.random.randn(5, 2) np_target = np.random.randn(5, 2) torch_input = torch.tensor(np_data, requires_grad=True) torch_target = torch.tensor(np_target) ms_input = ms_torch.tensor(np_data, requires_grad=True) ms_target = ms_torch.tensor(np_target) torch_loss = torch.nn.PoissonNLLLoss(reduction="sum") ms_loss = ms_torch.nn.PoissonNLLLoss(reduction="sum") torch_out = torch_loss(torch_input, torch_target) ms_out = ms_loss(ms_input, ms_target) assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype assert torch_out.shape == ms_out.shape def test_gaussian_nll_loss(): np_data = np.random.randn(5, 2) np_target = np.random.randn(5, 2) np_var = np.ones((5, 2)) torch_input = torch.tensor(np_data, requires_grad=True) torch_target = torch.tensor(np_target) torch_var = torch.tensor(np_var, requires_grad=True) ms_input = ms_torch.tensor(np_data, requires_grad=True) ms_target = ms_torch.tensor(np_target) ms_var = ms_torch.tensor(np_var, requires_grad=True) torch_loss = torch.nn.GaussianNLLLoss() ms_loss = ms_torch.nn.GaussianNLLLoss() torch_out = torch_loss(torch_input, torch_target, torch_var) ms_out = ms_loss(ms_input, ms_target, ms_var) assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype assert torch_out.shape == ms_out.shape def test_hinge_embedding_loss(): np_data = np.random.randn(5, 2) np_target = np.sign(np.random.randn(5, 2)) torch_input = torch.tensor(np_data, requires_grad=True) torch_target = torch.tensor(np_target) ms_input = ms_torch.tensor(np_data, requires_grad=True) ms_target = ms_torch.tensor(np_target) torch_loss = torch.nn.HingeEmbeddingLoss(reduction="none") ms_loss = ms_torch.nn.HingeEmbeddingLoss(reduction="none") torch_out = torch_loss(torch_input, torch_target) ms_out = ms_loss(ms_input, ms_target) assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype assert torch_out.shape == ms_out.shape # def test_multilabel_margin_loss(): # np_data = np.array([[0.1, 0.2, 0.4, 0.8], [0.2, 0.3, 0.5, 0.7]]) # np_target = np.array([[1, 2, 0, 3], [2, 3, -1, 1]]) # # torch_input = torch.tensor(np_data, requires_grad=True) # torch_target = torch.tensor(np_target, dtype=torch.int64) # ms_input = ms_torch.tensor(np_data, requires_grad=True) # ms_target = ms_torch.tensor(np_target, dtype=ms.int64) # # torch_loss = torch.nn.MultiLabelMarginLoss(reduction="none") # ms_loss = ms_torch.nn.MultiLabelMarginLoss(reduction="none") # # torch_out = torch_loss(torch_input, torch_target) # ms_out = ms_loss(ms_input, ms_target) # # assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) # assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype # assert torch_out.shape == ms_out.shape def test_multilabel_soft_margin_loss(): np_data = np.random.randn(5, 2) np_target = np.sign(np.random.randn(5, 2)) torch_input = torch.tensor(np_data, requires_grad=True) torch_target = torch.tensor(np_target) ms_input = ms_torch.tensor(np_data, requires_grad=True) ms_target = ms_torch.tensor(np_target) torch_loss = torch.nn.MultiLabelSoftMarginLoss(reduction="none") ms_loss = ms_torch.nn.MultiLabelSoftMarginLoss(reduction="none") torch_out = torch_loss(torch_input, torch_target) ms_out = ms_loss(ms_input, ms_target) assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype assert torch_out.shape == ms_out.shape def test_triplet_margin_with_distance_loss(): np_anc = np.random.randn(5, 6, 7) np_pos = np.random.randn(5, 6, 7) np_neg = np.random.randn(5, 6, 7) t_anc = torch.tensor(np_anc) t_pos = torch.tensor(np_pos) t_neg = torch.tensor(np_neg) ms_anc = ms_torch.tensor(np_anc) ms_pos = ms_torch.tensor(np_pos) ms_neg = ms_torch.tensor(np_neg) t_loss = torch.nn.TripletMarginWithDistanceLoss(distance_function=torch.nn.functional.cosine_similarity, swap=True) ms_loss = ms_torch.nn.TripletMarginWithDistanceLoss(distance_function=ms_torch.nn.functional.cosine_similarity, swap=True) torch_out = t_loss(t_anc, t_pos, t_neg) ms_out = ms_loss(ms_anc, ms_pos, ms_neg) assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype assert torch_out.shape == ms_out.shape def test_ctc_loss(): np_data = np.random.randn(24, 2, 10) np_target = np.random.rand(2, 10) * 10 np_input_length = np.array([8, 10]) np_target_length = np.array([5, 6]) torch_input = torch.tensor(np_data, requires_grad=True) torch_target = torch.tensor(np_target) torch_input_length = torch.tensor(np_input_length) torch_target_length = torch.tensor(np_target_length) ms_input = ms_torch.tensor(np_data, requires_grad=True) ms_target = ms_torch.tensor(np_target) ms_input_length = ms_torch.tensor(np_input_length) ms_target_length = ms_torch.tensor(np_target_length) torch_loss = torch.nn.CTCLoss(reduction="none") ms_loss = ms_torch.nn.CTCLoss(reduction="none") torch_out = torch_loss(torch_input, torch_target, torch_input_length, torch_target_length) ms_out = ms_loss(ms_input, ms_target, ms_input_length, ms_target_length) assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype assert torch_out.shape == ms_out.shape def test_margin_ranking_loss(): np_data1 = np.random.randn(5, 2) np_data2 = np.random.randn(5, 2) np_target = np.sign(np.random.randn(5, 2)) torch_input1 = torch.tensor(np_data1, requires_grad=True) torch_input2 = torch.tensor(np_data2, requires_grad=True) torch_target = torch.tensor(np_target) ms_input1 = ms_torch.tensor(np_data1, requires_grad=True) ms_input2 = ms_torch.tensor(np_data2, requires_grad=True) ms_target = ms_torch.tensor(np_target) torch_loss = torch.nn.MarginRankingLoss(reduction="none") ms_loss = ms_torch.nn.MarginRankingLoss(reduction="none") torch_out = torch_loss(torch_input1, torch_input2, torch_target) ms_out = ms_loss(ms_input1, ms_input2, ms_target) assert np.allclose(torch_out.detach().numpy(), ms_out.numpy()) assert torch_out.detach().numpy().dtype == ms_out.numpy().dtype assert torch_out.shape == ms_out.shape if __name__ == '__main__': test_smoothl1loss1() test_smoothl1loss2() test_smoothl1loss3() test_l1loss_reduction_none() test_l1loss_reduction_sum() test_l1loss_reduction_mean() test_mseloss_reduction_none() test_mseloss_reduction_sum() test_mseloss_reduction_mean() test_cross_entropy_reduction_none() test_cross_entropy_reduction_sum() test_cross_entropy_reduction_mean() test_bce_loss() test_binary_cross_entropy_with_logits() test_nllloss_reduction() test_huber_loss_mean() test_huber_loss_sum() test_huber_loss_none() ''' test_soft_margin_loss_none() test_soft_margin_loss_mean() test_soft_margin_loss_sum() ''' test_cosine_embedding_loss_none() test_cosine_embedding_loss_sum() test_cosine_embedding_loss_mean() test_cosine_triplet_margin_loss_none() test_cosine_triplet_margin_loss_sum() test_cosine_triplet_margin_loss_mean() test_multi_margin_loss_none() test_multi_margin_loss_weight() test_poisson_nll_loss() test_gaussian_nll_loss() test_hinge_embedding_loss() # test_multilabel_margin_loss() test_multilabel_soft_margin_loss() test_triplet_margin_with_distance_loss() test_ctc_loss() test_margin_ranking_loss() msadapter/testing/ut/pytorch/nn/test_module.py000066400000000000000000000112421436713554100221550ustar00rootroot00000000000000import torch import ms_adapter.pytorch as ms_torch import ms_adapter.pytorch.nn as nn from ms_adapter.pytorch.nn.modules.pooling import MaxPool2d def test_add_modules(): layer1 = MaxPool2d(1, 2) layer2 = MaxPool2d(1, 2) layer1.add_module("layer2", layer2) def test_named_children1(): layer1 = MaxPool2d(1, 2) layer2 = MaxPool2d(1, 2) layer1.add_module("layer2", layer2) aa = layer1.named_children() def test_modules(): layer1 = MaxPool2d(1, 2) layer2 = MaxPool2d(1, 2) layer1.add_module("layer2", layer2) aa = layer1.modules() def test_named_parameters(): layer1 = MaxPool2d(1, 2) layer2 = MaxPool2d(1, 2) layer1.add_module("layer2", layer2) aa = layer1.named_parameters() for k, v in aa: print(k, v) def test_named_children2(): layer1 = MaxPool2d(1, 2) layer2 = MaxPool2d(1, 2) layer1.add_module("layer2", layer2) for k, v in layer1.named_children(): print(k, v) def test_modulelist(): conv = nn.Conv2d(100, 20, 3) bn = nn.BatchNorm2d(20) relu = nn.ReLU() modulelist = nn.ModuleList([bn]) modulelist.insert(0, conv) modulelist.append(relu) modulelist.extend([relu, relu]) print(modulelist) def test_apply_and_sequential(): def init_weights(m): print(m) if type(m) == nn.Linear: m.weight.fill_(1.0) print(m.weight) net = nn.Sequential(nn.Conv2d(100, 20, 3), nn.Linear(2, 2), nn.Linear(2, 2)) net.append(nn.Linear(2, 2)) net.apply(init_weights) class MyTorchNet(torch.nn.Module): def __init__(self): super(MyTorchNet, self).__init__() self.conv1 = torch.nn.Conv2d(3,32,3,1,1) self.relu1 = torch.nn.ReLU() self.max_pooling1 = torch.nn.MaxPool2d(2,1) self.conv2 = torch.nn.Conv2d(3,32,2,1,1) self.relu2 = torch.nn.ReLU() self.max_pooling2 = torch.nn.MaxPool2d(2,1) self.dense1 = torch.nn.Linear(32*3*3,128) self.dense2 = torch.nn.Linear(128,10) def forward(self,x): x = self.conv1(x) x = self.relu1(x) x = self.max_pooling1(x) x = self.conv2(x) x = self.relu2(x) x = self.max_pooling2(x) x = self.dense1(x) x = self.dense2(x) return x class MyAdapterNet(ms_torch.nn.Module): def __init__(self): super(MyAdapterNet, self).__init__() self.conv1 = ms_torch.nn.Conv2d(3,32,3,1,1) self.relu1 = ms_torch.nn.ReLU() self.max_pooling1 = ms_torch.nn.MaxPool2d(2,1) self.conv2 = ms_torch.nn.Conv2d(3,32,2,1,1) self.relu2 = ms_torch.nn.ReLU() self.max_pooling2 = ms_torch.nn.MaxPool2d(2,1) self.dense1 = ms_torch.nn.Linear(32*3*3,128) self.dense2 = ms_torch.nn.Linear(128,10) def forward(self,x): x = self.conv1(x) x = self.relu1(x) x = self.max_pooling1(x) x = self.conv2(x) x = self.relu2(x) x = self.max_pooling2(x) x = self.dense1(x) x = self.dense2(x) return x def test_module_state_dict(): torch_model = MyTorchNet() ms_torch_model = MyAdapterNet() # print(torch_model.state_dict()) # print(ms_torch_model.state_dict()) def test_module_register(): class MyTorchModel(torch.nn.Module): def __init__(self): super(MyTorchModel, self).__init__() self.register_buffer('mybuffer1', torch.tensor([2, 2])) self.register_parameter('my_param1', torch.nn.Parameter(torch.zeros(3, 3))) self.param1 = torch.ones(3, 3) def forward(self, x): return x class MyAdapterModel(ms_torch.nn.Module): def __init__(self): super(MyAdapterModel, self).__init__() self.register_buffer('mybuffer1', ms_torch.tensor([2, 2])) self.register_parameter('my_param1', ms_torch.nn.Parameter(ms_torch.zeros(3, 3))) self.param1 = ms_torch.ones(3, 3) def forward(self, x): return x torch_model = MyTorchModel() ms_mymodel = MyAdapterModel() assert torch_model.my_param1.size() == ms_mymodel.my_param1.size() assert torch_model.mybuffer1.size() == ms_mymodel.mybuffer1.size() assert torch_model.param1.size() == ms_mymodel.param1.size() def test_named_modules(): l = nn.Linear(2, 2) net = nn.Sequential(l, l) str_ms = () for m in net.named_modules(prefix='aaa'): str_ms += m print(str_ms) if __name__ == '__main__': test_named_children1() test_add_modules() test_named_children2() test_modules() test_named_parameters() test_modulelist() test_apply_and_sequential() test_module_state_dict() test_module_register() test_named_modules() msadapter/testing/ut/pytorch/nn/test_normalization.py000066400000000000000000000071321436713554100235610ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import numpy as np import torch from mindspore import context import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np context.set_context(mode=ms.PYNATIVE_MODE) def test_layer_norm1(): batch, sentence_length, embedding_dim = 20, 5, 10 data = np.random.randn(batch, sentence_length, embedding_dim).astype(np.float32) torch_embedding = torch.tensor(data) torch_layer_norm = torch.nn.LayerNorm(embedding_dim) torch_output = torch_layer_norm(torch_embedding) ms_torch_embedding = ms_torch.tensor(data) ms_torch_layer_norm = ms_torch.nn.LayerNorm(embedding_dim) ms_torch_output = ms_torch_layer_norm(ms_torch_embedding) assert np.allclose(ms_torch_output.asnumpy(), torch_output.detach().numpy(), atol=1e-3) def test_layer_norm2(): N, C, H, W = 20, 5, 10, 10 data = np.random.randn(N, C, H, W).astype(np.float32) torch_input = torch.tensor(data) torch_layer_norm = torch.nn.LayerNorm([C, H, W]) torch_output = torch_layer_norm(torch_input) ms_torch_input = ms_torch.tensor(data) ms_torch_layer_norm = ms_torch.nn.LayerNorm([C, H, W]) ms_torch_output = ms_torch_layer_norm(ms_torch_input) assert np.allclose(ms_torch_output.asnumpy(), torch_output.detach().numpy(), atol=1e-5) def test_group_norm1(): N, C, H, W = 20, 5, 10, 10 data = np.random.randn(N, C, H, W).astype(np.float32) torch_input = torch.tensor(data) torch_layer_norm = torch.nn.GroupNorm(num_groups=5, num_channels=5) torch_output = torch_layer_norm(torch_input) ms_torch_input = ms_torch.tensor(data) ms_torch_layer_norm = ms_torch.nn.GroupNorm(num_groups=5, num_channels=5) ms_torch_output = ms_torch_layer_norm(ms_torch_input) assert np.allclose(ms_torch_output.asnumpy(), torch_output.detach().numpy(), atol=1e-5) def test_group_norm2(): N, C, H, W = 20, 10, 10, 10 data = np.random.randn(N, C, H, W).astype(np.float32) torch_input = torch.tensor(data) torch_layer_norm = torch.nn.GroupNorm(num_groups=5, num_channels=10, affine=False) torch_output = torch_layer_norm(torch_input) ms_torch_input = ms_torch.tensor(data) ms_torch_layer_norm = ms_torch.nn.GroupNorm(num_groups=5, num_channels=10, affine=False) ms_torch_output = ms_torch_layer_norm(ms_torch_input) assert np.allclose(ms_torch_output.asnumpy(), torch_output.detach().numpy(), atol=1e-5) def local_response_norm1(): data = np.random.random((4, 4, 4)).astype(np.float32)*10 ms_input = ms_torch.tensor(data) ms_torch_layer_norm = ms_torch.nn.LocalResponseNorm(size=3, alpha=0.01) ms_out = ms_torch_layer_norm(ms_input) torch_input = torch.tensor(data) torch_layer_norm = torch.nn.LocalResponseNorm(size=3, alpha=0.01) torch_out = torch_layer_norm(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def local_response_norm2(): data = np.random.random((2, 4, 4, 3, 2)).astype(np.float32)*10 ms_input = ms_torch.tensor(data) ms_torch_layer_norm = ms_torch.nn.LocalResponseNorm(3, 0.01) ms_out = ms_torch_layer_norm(ms_input) torch_input = torch.tensor(data) torch_layer_norm = torch.nn.LocalResponseNorm(3, 0.01) torch_out = torch_layer_norm(torch_input) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_layer_norm1() test_layer_norm2() test_group_norm1() test_group_norm2() local_response_norm1() local_response_norm2() msadapter/testing/ut/pytorch/nn/test_padding.py000066400000000000000000000216351436713554100223050ustar00rootroot00000000000000import numpy as np import torch import mindspore as ms from mindspore import context import ms_adapter.pytorch as ms_pytorch context.set_context(mode=ms.PYNATIVE_MODE) def test_constant_pad_1d(): padding = 3 value = 3.5 pt_input_2d = torch.ones(2, 3) pt_pad_fun1 = torch.nn.ConstantPad1d(padding, value) pt_pad_out1 = pt_pad_fun1(pt_input_2d) ms_input_2d = ms_pytorch.ones(2, 3) ms_pad_fun1 = ms_pytorch.nn.ConstantPad1d(padding, value) ms_pad_out1 = ms_pad_fun1(ms_input_2d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) padding = (3, 1) value = 2.0 pt_input_3d = torch.ones(2, 3, 4) pt_pad_fun2 = torch.nn.ConstantPad1d(padding, value) pt_pad_out2 = pt_pad_fun2(pt_input_3d) ms_input_3d = ms_pytorch.ones(2, 3, 4) ms_pad_fun2 = ms_pytorch.nn.ConstantPad1d(padding, value) ms_pad_out2 = ms_pad_fun2(ms_input_3d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_constant_pad_2d(): padding = 2 value = 2.5 pt_input_3d = torch.ones(2, 3, 4) pt_pad_fun1 = torch.nn.ConstantPad2d(padding, value) pt_pad_out1 = pt_pad_fun1(pt_input_3d) ms_input_3d = ms_pytorch.ones(2, 3, 4) ms_pad_fun1 = ms_pytorch.nn.ConstantPad2d(padding, value) ms_pad_out1 = ms_pad_fun1(ms_input_3d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) padding = (-1, 1, 0, 1) value = 2.0 pt_input_4d = torch.ones(1, 2, 3, 4) pt_pad_fun2 = torch.nn.ConstantPad2d(padding, value) pt_pad_out2 = pt_pad_fun2(pt_input_4d) ms_input_4d = ms_pytorch.ones(1, 2, 3, 4) ms_pad_fun2 = ms_pytorch.nn.ConstantPad2d(padding, value) ms_pad_out2 = ms_pad_fun2(ms_input_4d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_constant_pad_3d(): padding = 1 value = 2.5 pt_input_4d = torch.ones(2, 1, 3, 4) pt_pad_fun1 = torch.nn.ConstantPad3d(padding, value) pt_pad_out1 = pt_pad_fun1(pt_input_4d) ms_input_4d = ms_pytorch.ones(2, 1, 3, 4) ms_pad_fun1 = ms_pytorch.nn.ConstantPad3d(padding, value) ms_pad_out1 = ms_pad_fun1(ms_input_4d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) padding = (3, 3, 6, 6, 0, 1) value = 2.0 pt_input_5d = torch.ones(1, 2, 1, 3, 4) pt_pad_fun2 = torch.nn.ConstantPad3d(padding, value) pt_pad_out2 = pt_pad_fun2(pt_input_5d) ms_input_5d = ms_pytorch.ones(1, 2, 1, 3, 4) ms_pad_fun2 = ms_pytorch.nn.ConstantPad3d(padding, value) ms_pad_out2 = ms_pad_fun2(ms_input_5d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_reflection_pad_1d(): padding = 3 pt_input_2d = torch.ones(2, 6) pt_pad_fun1 = torch.nn.ReflectionPad1d(padding) pt_pad_out1 = pt_pad_fun1(pt_input_2d) ms_input_2d = ms_pytorch.ones(2, 6) ms_pad_fun1 = ms_pytorch.nn.ReflectionPad1d(padding) ms_pad_out1 = ms_pad_fun1(ms_input_2d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) padding = (2, 1) pt_input_3d = torch.ones(2, 3, 4) pt_pad_fun2 = torch.nn.ReflectionPad1d(padding) pt_pad_out2 = pt_pad_fun2(pt_input_3d) ms_input_3d = ms_pytorch.ones(2, 3, 4) ms_pad_fun2 = ms_pytorch.nn.ReflectionPad1d(padding) ms_pad_out2 = ms_pad_fun2(ms_input_3d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_reflection_pad_2d(): padding = 2 pt_input_3d = torch.ones(2, 3, 3) pt_pad_fun1 = torch.nn.ReflectionPad2d(padding) pt_pad_out1 = pt_pad_fun1(pt_input_3d) ms_input_3d = ms_pytorch.ones(2, 3, 3) ms_pad_fun1 = ms_pytorch.nn.ReflectionPad2d(padding) ms_pad_out1 = ms_pad_fun1(ms_input_3d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) padding = (1, 1, 2, 0) pt_input_4d = torch.ones(1, 2, 3, 4) pt_pad_fun2 = torch.nn.ReflectionPad2d(padding) pt_pad_out2 = pt_pad_fun2(pt_input_4d) ms_input_4d = ms_pytorch.ones(1, 2, 3, 4) ms_pad_fun2 = ms_pytorch.nn.ReflectionPad2d(padding) ms_pad_out2 = ms_pad_fun2(ms_input_4d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_reflection_pad_3d(): padding = 2 pt_input_4d = torch.ones(2, 3, 3, 4) pt_pad_fun1 = torch.nn.ReflectionPad3d(padding) pt_pad_out1 = pt_pad_fun1(pt_input_4d) ms_input_4d = ms_pytorch.ones(2, 3, 3, 4) ms_pad_fun1 = ms_pytorch.nn.ReflectionPad3d(padding) ms_pad_out1 = ms_pad_fun1(ms_input_4d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) padding = (1, 1, 2, 0, 3, 2) pt_input_5d = torch.ones(1, 2, 6, 4, 5) pt_pad_fun2 = torch.nn.ReflectionPad3d(padding) pt_pad_out2 = pt_pad_fun2(pt_input_5d) ms_input_5d = ms_pytorch.ones(1, 2, 6, 4, 5) ms_pad_fun2 = ms_pytorch.nn.ReflectionPad3d(padding) ms_pad_out2 = ms_pad_fun2(ms_input_5d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_zero_pad_2d(): padding = 2 pt_input_3d = torch.ones(1, 3, 3) pt_pad_fun1 = torch.nn.ZeroPad2d(padding) pt_pad_out1 = pt_pad_fun1(pt_input_3d) ms_input_3d = ms_pytorch.ones(1, 3, 3) ms_pad_fun1 = ms_pytorch.nn.ZeroPad2d(padding) ms_pad_out1 = ms_pad_fun1(ms_input_3d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) padding = (1, 1, 2, 0) pt_input_4d = torch.ones(1, 1, 3, 3) pt_pad_fun2 = torch.nn.ZeroPad2d(padding) pt_pad_out2 = pt_pad_fun2(pt_input_4d) ms_input_4d = ms_pytorch.ones(1, 1, 3, 3) ms_pad_fun2 = ms_pytorch.nn.ZeroPad2d(padding) ms_pad_out2 = ms_pad_fun2(ms_input_4d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_replication_pad_1d(): padding1 = 3 padding2 = (2, 1) pt_input_3d = torch.ones(2, 3, 3) ms_input_3d = ms_pytorch.ones(2, 3, 3) pt_pad_fun1 = torch.nn.ReplicationPad1d(padding1) ms_pad_fun1 = ms_pytorch.nn.ReplicationPad1d(padding1) pt_pad_out1 = pt_pad_fun1(pt_input_3d) ms_pad_out1 = ms_pad_fun1(ms_input_3d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) pt_pad_fun2 = torch.nn.ReplicationPad1d(padding2) pt_pad_out2 = pt_pad_fun2(pt_input_3d) ms_pad_fun2 = ms_pytorch.nn.ReplicationPad1d(padding2) ms_pad_out2 = ms_pad_fun2(ms_input_3d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_replication_pad_2d(): padding1 = 3 padding2 = (2, 1, 1, 2) pt_input_4d = torch.ones(1, 1, 3, 3) ms_input_4d = ms_pytorch.ones(1, 1, 3, 3) pt_pad_fun1 = torch.nn.ReplicationPad2d(padding1) ms_pad_fun1 = ms_pytorch.nn.ReplicationPad2d(padding1) pt_pad_out1 = pt_pad_fun1(pt_input_4d) ms_pad_out1 = ms_pad_fun1(ms_input_4d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) pt_pad_fun2 = torch.nn.ReplicationPad2d(padding2) pt_pad_out2 = pt_pad_fun2(pt_input_4d) ms_pad_fun2 = ms_pytorch.nn.ReplicationPad2d(padding2) ms_pad_out2 = ms_pad_fun2(ms_input_4d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) def test_replication_pad_3d(): padding1 = 3 padding2 = (2, 1, 1, 2, 3, 2) pt_input_4d = torch.ones(2, 3, 8, 32, 48) ms_input_4d = ms_pytorch.ones(2, 3, 8, 32, 48) pt_pad_fun1 = torch.nn.ReplicationPad3d(padding1) ms_pad_fun1 = ms_pytorch.nn.ReplicationPad3d(padding1) pt_pad_out1 = pt_pad_fun1(pt_input_4d) ms_pad_out1 = ms_pad_fun1(ms_input_4d) assert (pt_pad_out1.shape == ms_pad_out1.shape) assert np.allclose(pt_pad_out1.numpy(), ms_pad_out1.asnumpy()) pt_pad_fun2 = torch.nn.ReplicationPad3d(padding2) pt_pad_out2 = pt_pad_fun2(pt_input_4d) ms_pad_fun2 = ms_pytorch.nn.ReplicationPad3d(padding2) ms_pad_out2 = ms_pad_fun2(ms_input_4d) assert (pt_pad_out2.shape == ms_pad_out2.shape) assert np.allclose(pt_pad_out2.numpy(), ms_pad_out2.asnumpy()) if __name__ == '__main__': test_constant_pad_1d() test_constant_pad_2d() test_constant_pad_3d() test_reflection_pad_1d() test_reflection_pad_2d() test_reflection_pad_3d() test_zero_pad_2d() test_replication_pad_1d() test_replication_pad_2d() test_replication_pad_3d() msadapter/testing/ut/pytorch/nn/test_parameter.py000066400000000000000000000104751436713554100226570ustar00rootroot00000000000000import numpy as np import torch from mindspore import context import mindspore as ms import ms_adapter.pytorch as ms_torch from ms_adapter.pytorch.nn import Parameter context.set_context(mode=ms.PYNATIVE_MODE) def test_parameter_data1(): input = np.array([[[1.2, 2.1], [2.2, 3.2], [2.0, 3.0]]]).astype(np.float32) torch_tensor = torch.tensor(input) torch_parameter = torch.nn.parameter.Parameter(data=torch_tensor) torch_parameter.data.fill_(0) ms_tensor = ms_torch.tensor(input) ms_parameter = Parameter(data=ms_tensor) ms_parameter.data.fill_(0) assert np.allclose(torch_parameter.data.numpy(), ms_parameter.data.numpy()) return def test_parameter_data2(): input = np.array([[[1.2, 2.1], [2.2, 3.2], [2.0, 3.0]]]).astype(np.float32) torch_tensor = torch.tensor(input) torch_parameter = torch.nn.parameter.Parameter(data=torch_tensor, requires_grad=False) torch_parameter = -torch_parameter ms_tensor = ms_torch.tensor(input) ms_parameter = Parameter(data=ms_tensor, requires_grad=False) ms_parameter = -ms_parameter assert np.allclose(torch_parameter.numpy(), ms_parameter.numpy()) return def test_parameter_no_grad(): input = np.array([[[1.2, 2.1], [2.2, 3.2], [2.0, 3.0]]]).astype(np.float32) torch_tensor = torch.tensor(input) torch_parameter = torch.nn.parameter.Parameter(data=torch_tensor, requires_grad=False) torch_parameter.fill_(0) ms_tensor = ms_torch.tensor(input) ms_parameter = Parameter(data=ms_tensor, requires_grad=False) ms_parameter.fill_(0) assert np.allclose(torch_parameter.numpy(), ms_parameter.numpy()) return def test_requires_grad_(): input_init = np.random.randn(1, 3, 16).astype(float) weight_init = np.random.randn(64, 3, 3).astype(float) weight_init_t = np.expand_dims(weight_init, axis=2) bias_init = np.random.randn(64).astype(float) class Conv1dPadModel(torch.nn.Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv1dPadModel, self).__init__() self.conv = torch.nn.Conv1d(in_channels=3, out_channels=64, kernel_size=3, stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = torch.nn.Parameter(torch.tensor(weight_init, dtype=torch.float32)) self.conv.bias = torch.nn.Parameter(torch.tensor(bias_init, dtype=torch.float32)) self.conv.bias.requires_grad_(False) def forward(self, inputs): x = self.conv(inputs) return torch.sum(x) class Conv1dPadModelMs(ms_torch.nn.Module): def __init__(self, stride=1, padding=0, dilation=1, padding_mode='zeros'): super(Conv1dPadModelMs, self).__init__() self.conv = ms_torch.nn.Conv1d(in_channels=3, out_channels=64, kernel_size=3, stride=stride, padding=padding, dilation=dilation, padding_mode=padding_mode) self.conv.weight = Parameter(ms_torch.tensor(weight_init_t, ms_torch.float32)) self.conv.bias = Parameter(ms_torch.tensor(bias_init, ms_torch.float32)) self.conv.bias.requires_grad_(False) def forward(self, inputs): x = self.conv(inputs) return ms_torch.sum(x) torch_net = Conv1dPadModel() py_input = torch.tensor(input_init, dtype=torch.float32) torch_out = torch_net(py_input) ms_net = Conv1dPadModelMs() ms_input = ms_torch.tensor(input_init, ms_torch.float32) ms_out = ms_net(ms_input) if ms.get_context('device_target') == 'Ascend': assert np.allclose(ms_out.numpy(), torch_out.detach().numpy(), rtol=1e-3, atol=1e-2) else: assert np.allclose(ms_out.numpy(), torch_out.detach().numpy()) torch_out.backward() grad_out = torch_net.conv.weight.grad ms_grad = ms.grad(ms_net, grad_position=None, weights=ms_net.trainable_params())(ms_input) assert len(ms_grad) == 1 ms_grad = ms.ops.squeeze(ms_grad[0]) if ms.get_context('device_target') == 'Ascend': assert np.allclose(grad_out.numpy(), ms_grad.numpy(), rtol=1e-3, atol=1e-3) else: assert np.allclose(grad_out.numpy(), ms_grad.numpy()) if __name__ == '__main__': test_parameter_data1() test_parameter_data2() test_parameter_no_grad() test_requires_grad_() msadapter/testing/ut/pytorch/nn/test_pixelshuffle.py000066400000000000000000000015331436713554100233700ustar00rootroot00000000000000import ms_adapter.pytorch as ms_torch import numpy as np import torch def test_pixelshuffle(): data = np.random.randn(1, 9 ,4, 4).astype(np.float32) pixel_shuffle = torch.nn.PixelShuffle(3) input = torch.tensor(data) torch_output = pixel_shuffle(input) pixel_shuffle = ms_torch.nn.PixelShuffle(3) input = ms_torch.tensor(data) ms_output = pixel_shuffle(input) assert np.allclose(torch_output.numpy(), ms_output.numpy()) def test_pixelunshuffle(): data = np.random.randn(1, 1 ,12, 12).astype(np.float32) pixel_unshuffle = torch.nn.PixelUnshuffle(3) input = torch.tensor(data) torch_output = pixel_unshuffle(input) pixel_unshuffle = ms_torch.nn.PixelUnshuffle(3) input = ms_torch.tensor(data) ms_output = pixel_unshuffle(input) assert np.allclose(torch_output.numpy(), ms_output.numpy()) msadapter/testing/ut/pytorch/nn/test_pooling.py000066400000000000000000000446041436713554100223470ustar00rootroot00000000000000import numpy as np import torch from mindspore import Tensor from ms_adapter.pytorch.nn import MaxPool1d, MaxPool2d, MaxPool3d, AvgPool1d, AvgPool2d, AvgPool3d, \ AdaptiveAvgPool1d, AdaptiveAvgPool2d, AdaptiveAvgPool3d, AdaptiveMaxPool1d, AdaptiveMaxPool2d, AdaptiveMaxPool3d, \ LPPool1d, LPPool2d, FractionalMaxPool2d, FractionalMaxPool3d import mindspore as ms ms.context.set_context(mode=ms.GRAPH_MODE) def test_maxpool1d_compare1(): ms_net = MaxPool1d(4, 2, padding=2) torch_net = torch.nn.MaxPool1d(4, 2, padding=2) data = np.random.random([4, 5, 6]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxpool2d_compare1(): ms_net = MaxPool2d(3, 2) torch_net = torch.nn.MaxPool2d(3, 2) data = np.random.random([4, 5, 6, 7]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxpool2d_compare2(): ms_net = MaxPool2d(4, 2, padding=2) torch_net = torch.nn.MaxPool2d(4, 2, padding=2) data = np.random.random([6, 7, 8, 9]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxpool2d_compare3(): ms_net = MaxPool2d((3, 5), (3, 1), padding=0) torch_net = torch.nn.MaxPool2d((3, 5), (3, 1), padding=0) data = np.random.random([1, 32, 9, 9]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxpool2d_compare4(): ms_net = MaxPool2d(kernel_size=2, ceil_mode=True) torch_net = torch.nn.MaxPool2d(kernel_size=2, ceil_mode=True) data = np.zeros((1, 1, 5, 5)) for i in range(5): data[0][0][i] = i * np.ones(5) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxpool2d_compare5(): ms_net = MaxPool2d(3, 2, ceil_mode=True) torch_net = torch.nn.MaxPool2d(3, 2, ceil_mode=True) data = np.random.random([4, 5, 6, 7]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxpool2d_compare6(): ms_net = MaxPool2d(4, 2, padding=2, ceil_mode=True) torch_net = torch.nn.MaxPool2d(4, 2, padding=2, ceil_mode=True) data = np.random.random([6, 7, 8, 9]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # TODO: ms.ops.pad to support "value = float(-inf)" ''' def test_maxpool2d_compare7(): ms_net = MaxPool2d(4, 2, padding=2, ceil_mode=True) torch_net = torch.nn.MaxPool2d(4, 2, padding=2, ceil_mode=True) data = np.random.randn(6, 7, 8, 9) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) ''' def test_maxpool3d_compare1(): ms_net = MaxPool3d(4, 2, padding=2) torch_net = torch.nn.MaxPool3d(4, 2, padding=2) data = np.random.random([4, 5, 6, 6, 6]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avgpool1d_compare(): ms_net = AvgPool1d(6, stride=2, padding=3) torch_net = torch.nn.AvgPool1d(6, stride=2, padding=3) data = np.random.random([5, 3, 8]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avgpool2d_compare1(): ms_net = AvgPool2d(2, stride=2, padding=0) torch_net = torch.nn.AvgPool2d(2, stride=2, padding=0) data = np.random.random([5, 3, 11, 11]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avgpool2d_compare2(): ms_net = AvgPool2d(6, stride=2, padding=3) torch_net = torch.nn.AvgPool2d(6, stride=2, padding=3) data = np.random.random([5, 3, 8, 8]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avgpool2d_compare3(): ms_net = AvgPool2d((3, 5), (3, 1), padding=0) torch_net = torch.nn.AvgPool2d((3, 5), (3, 1), padding=0) data = np.random.random([1, 32, 9, 9]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_avgpool3d_compare(): ms_net = AvgPool3d(6, stride=2, padding=3) torch_net = torch.nn.AvgPool3d(6, stride=2, padding=3) data = np.random.random([5, 3, 8, 8, 8]) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_adaptive_maxpool1d_compare1(): ms_net = AdaptiveMaxPool1d(3) torch_net = torch.nn.AdaptiveMaxPool1d(3) data = np.random.randn(1, 32, 7).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data.astype(np.float32)) ms_output = ms_net(ms_input) assert (torch_output.shape == ms_output.shape) # TODO: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_adaptive_maxpool1d_compare2(): ms_net = AdaptiveMaxPool1d(3, True) torch_net = torch.nn.AdaptiveMaxPool1d(3, True) data = np.arange(3 * 8).reshape(1, 3, 8).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output, ms_indices = ms_net(ms_input) torch_output, torch_indices = torch_net(torch_input) assert (torch_output.shape == ms_output.shape) # TODO: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # TODO: assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) def test_adaptive_maxpool2d_compare1(): ms_net = AdaptiveMaxPool2d((3, 10)) torch_net = torch.nn.AdaptiveMaxPool2d((3, 10)) data = np.random.randn(32, 9, 9).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output = ms_net(ms_input) torch_output = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_adaptive_maxpool2d_compare2(): ms_net = AdaptiveMaxPool2d(3, True) torch_net = torch.nn.AdaptiveMaxPool2d(3, True) data = np.random.randn(2, 32, 6, 9).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output, ms_indices = ms_net(ms_input) torch_output, torch_indices = torch_net(torch_input) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) def test_adaptive_maxpool3d_compare1(): ms_net = AdaptiveMaxPool3d(3) torch_net = torch.nn.AdaptiveMaxPool3d(3) data = np.random.randn(1, 32, 7, 9, 5).astype(np.float32) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input) ms_input = Tensor(data.astype(np.float32)) ms_output = ms_net(ms_input) assert (torch_output.shape == ms_output.shape) # TODO: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_adaptive_avgpool1d_compare1(): ms_net = AdaptiveAvgPool1d(3) torch_net = torch.nn.AdaptiveAvgPool1d(3) data = np.arange(3 * 8).reshape(1, 3, 8).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output= ms_net(ms_input) torch_output= torch_net(torch_input) assert (torch_output.shape == ms_output.shape) # TODO: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_adaptive_avgpool2d_compare1(): ms_net = AdaptiveAvgPool2d(3) torch_net = torch.nn.AdaptiveAvgPool2d(3) data = np.random.randn(2, 32, 6, 9).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output= ms_net(ms_input) torch_output= torch_net(torch_input) assert (torch_output.shape == ms_output.shape) # TODO: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_adaptive_avgpool3d_compare1(): ms_net = AdaptiveAvgPool3d(3) torch_net = torch.nn.AdaptiveAvgPool3d(3) data = np.random.randn(1, 32, 7, 9, 5).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output= ms_net(ms_input) torch_output= torch_net(torch_input) assert (torch_output.shape == ms_output.shape) # TODO: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # def test_fractional_maxpool2d_compare1(): # ms_net = FractionalMaxPool2d(1, output_ratio=(0.5, 0.5)) # torch_net = torch.nn.FractionalMaxPool2d(1, output_ratio=(0.5, 0.5)) # # data = np.array([0.3220, 0.9545, 0.7879, 0.0975, 0.3698, # 0.5135, 0.5740, 0.3435, 0.1895, 0.8764, # 0.9581, 0.4760, 0.9014, 0.8522, 0.3664, # 0.4980, 0.9673, 0.9879, 0.6988, 0.9022, # 0.9304, 0.1558, 0.0153, 0.1559, 0.9852]).reshape([1, 1, 5, 5]) # ms_input = Tensor(data.astype(np.float32)) # torch_input = torch.Tensor(data) # # ms_output = ms_net(ms_input) # torch_output = torch_net(torch_input) # # assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # # def test_fractional_maxpool2d_compare2(): # _random_samples = np.array([[[0.8, 0.8]]]) # ms_random_samples = Tensor(_random_samples.astype(np.float32)) # torch_random_samples = torch.Tensor(_random_samples) # ms_net = FractionalMaxPool2d(kernel_size=1, output_size=(2, 2), _random_samples=ms_random_samples, # return_indices=True) # torch_net = torch.nn.FractionalMaxPool2d(kernel_size=1, output_size=(2, 2), _random_samples=torch_random_samples, # return_indices=True) # # data = np.array([0.3220, 0.9545, 0.7879, 0.0975, 0.3698, # 0.5135, 0.5740, 0.3435, 0.1895, 0.8764, # 0.9581, 0.4760, 0.9014, 0.8522, 0.3664, # 0.4980, 0.9673, 0.9879, 0.6988, 0.9022, # 0.9304, 0.1558, 0.0153, 0.1559, 0.9852]).reshape([1, 1, 5, 5]) # # ms_input = Tensor(data.astype(np.float32)) # torch_input = torch.Tensor(data) # # ms_output, ms_indices= ms_net(ms_input) # torch_output, torch_indices = torch_net(torch_input) # # assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) # # def test_fractional_maxpool3d_compare1(): # _random_samples = np.array([0.7, 0.7, 0.7]).reshape([1, 1, 3]) # ms_random_samples = Tensor(_random_samples.astype(np.float32)) # torch_random_samples = torch.Tensor(_random_samples) # ms_net = FractionalMaxPool3d(kernel_size=(1.0, 1.0, 1.0), output_size=(1, 1, 3), _random_samples=ms_random_samples, # return_indices=True) # torch_net = torch.nn.FractionalMaxPool3d(kernel_size=(1, 1, 1), output_size=(1, 1, 3), _random_samples=torch_random_samples, # return_indices=True) # # data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]).reshape([1, 1, 2, 2, 4]) # # ms_input = Tensor(data.astype(np.float32)) # torch_input = torch.Tensor(data) # # ms_output, ms_indices = ms_net(ms_input) # torch_output, torch_indices = torch_net(torch_input) # # assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) # # def test_fractional_maxpool3d_compare2(): # _random_samples = np.array([0.7, 0.7, 0.7]).reshape([1, 1, 3]) # ms_random_samples = Tensor(_random_samples.astype(np.float32)) # torch_random_samples = torch.Tensor(_random_samples) # ms_net = FractionalMaxPool3d(kernel_size=(1.0, 1.0, 1.0), output_ratio=(0.5, 0.5, 0.5), _random_samples=ms_random_samples, # return_indices=True) # torch_net = torch.nn.FractionalMaxPool3d(kernel_size=(1, 1, 1), output_ratio=(0.5, 0.5, 0.5), _random_samples=torch_random_samples, # return_indices=True) # # data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]).reshape([1, 1, 2, 2, 4]) # # ms_input = Tensor(data.astype(np.float32)) # torch_input = torch.Tensor(data) # # ms_output, ms_indices = ms_net(ms_input) # torch_output, torch_indices = torch_net(torch_input) # assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # assert np.allclose(ms_indices.asnumpy(), torch_indices.numpy()) def test_lppool1d_compare1(): ms_net = LPPool1d(norm_type=1, kernel_size=3, stride=1) torch_net = torch.nn.LPPool1d(norm_type=1, kernel_size=3, stride=1) data = np.random.randn(2, 3, 4).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output= ms_net(ms_input) torch_output= torch_net(torch_input) assert (torch_output.shape == ms_output.shape) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_lppool2d_compare1(): ms_net = LPPool2d(norm_type=1, kernel_size=3, stride=1) torch_net = torch.nn.LPPool2d(norm_type=1, kernel_size=3, stride=1) data = np.random.randn(2, 3, 4, 5).astype(np.float32) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) ms_output= ms_net(ms_input) torch_output= torch_net(torch_input) assert (torch_output.shape == ms_output.shape) if ms.context.get_context('device_target') == 'Ascend': assert np.allclose(ms_output.asnumpy(), torch_output.numpy(), atol=1e-3) else: assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) if __name__ == '__main__': test_maxpool1d_compare1() test_maxpool2d_compare1() test_maxpool2d_compare2() test_maxpool2d_compare3() test_maxpool2d_compare4() test_maxpool2d_compare5() test_maxpool2d_compare6() # test_maxpool2d_compare7() test_maxpool3d_compare1() test_avgpool1d_compare() test_avgpool2d_compare1() test_avgpool2d_compare2() test_avgpool2d_compare3() test_avgpool3d_compare() test_adaptive_maxpool1d_compare1() test_adaptive_maxpool1d_compare2() test_adaptive_maxpool2d_compare1() test_adaptive_maxpool2d_compare2() test_adaptive_maxpool3d_compare1() test_adaptive_avgpool1d_compare1() test_adaptive_avgpool2d_compare1() test_adaptive_avgpool3d_compare1() test_lppool1d_compare1() test_lppool2d_compare1() # test_fractional_maxpool2d_compare1() # test_fractional_maxpool2d_compare2() # test_fractional_maxpool3d_compare1() # test_fractional_maxpool3d_compare2() msadapter/testing/ut/pytorch/nn/test_rnn.py000066400000000000000000000333221436713554100214700ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import ms_adapter import ms_adapter.pytorch as ms_torch from ms_adapter.pytorch import tensor from mindspore import context import mindspore as ms import numpy as np import torch #context.set_context(mode=ms.PYNATIVE_MODE) context.set_context(mode=ms.GRAPH_MODE) def test_rnn1(): input = np.random.random((5, 3, 10)).astype(np.float32) h0 = np.random.random((4, 3, 20)).astype(np.float32) torch_tensor = torch.tensor(input) torch_h0 = torch.tensor(h0) torch_rnn = torch.nn.RNN(input_size=10, hidden_size=20, num_layers=2, nonlinearity='tanh', bias=True, batch_first=False, dropout=0.6, bidirectional=True) torch_output, torch_h = torch_rnn(torch_tensor, torch_h0) ms_tensor = tensor(input) ms_h0 = tensor(h0) ms_rnn = ms_adapter.pytorch.nn.RNN(input_size=10, hidden_size=20, num_layers=2, nonlinearity='tanh', bias=True, batch_first=False, dropout=0.6, bidirectional=True) ms_output, ms_h = ms_rnn(ms_tensor, ms_h0) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_h.shape == ms_h.shape assert torch_output.shape == ms_output.shape def test_rnn2(): input = np.random.random((3, 5, 10)).astype(np.float32) h0 = np.random.random((2, 3, 20)).astype(np.float32) torch_tensor = torch.tensor(input) torch_h0 = torch.tensor(h0) torch_rnn = torch.nn.RNN(input_size=10, hidden_size=20, num_layers=2, nonlinearity='relu', bias=False, batch_first=True, dropout=0.6, bidirectional=False) torch_output, torch_h = torch_rnn(torch_tensor, torch_h0) ms_tensor = tensor(input) ms_h0 = tensor(h0) ms_rnn = ms_adapter.pytorch.nn.RNN(input_size=10, hidden_size=20, num_layers=2, nonlinearity='relu', bias=False, batch_first=True, dropout=0.6, bidirectional=False) ms_output, ms_h = ms_rnn(ms_tensor, ms_h0) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_h.shape == ms_h.shape assert torch_output.shape == ms_output.shape def test_rnn3(): input = np.random.random((5, 3, 10)).astype(np.float32) torch_tensor = torch.tensor(input) torch_rnn = torch.nn.RNN(input_size=10, hidden_size=20, num_layers=2, nonlinearity='tanh', bias=True, batch_first=False, dropout=0.6, bidirectional=True) torch_output, torch_h = torch_rnn(torch_tensor) ms_tensor = tensor(input) ms_rnn = ms_adapter.pytorch.nn.RNN(input_size=10, hidden_size=20, num_layers=2, nonlinearity='tanh', bias=True, batch_first=False, dropout=0.6, bidirectional=True) ms_output, ms_h = ms_rnn(ms_tensor) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_h.shape == ms_h.shape assert torch_output.shape == ms_output.shape def test_gru1(): input = np.random.random((5, 3, 10)).astype(np.float32) h0 = np.random.random((4, 3, 32)).astype(np.float32) torch_tensor = torch.tensor(input) torch_h0 = torch.tensor(h0) torch_gru = torch.nn.GRU(input_size=10, hidden_size=32, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) torch_output, torch_h = torch_gru(torch_tensor, torch_h0) ms_tensor = tensor(input) ms_h0 = tensor(h0) ms_gru = ms_adapter.pytorch.nn.GRU(input_size=10, hidden_size=32, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) ms_output, ms_h = ms_gru(ms_tensor, ms_h0) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_h.shape == ms_h.shape assert torch_output.shape == ms_output.shape def test_gru2(): input = np.random.random((3, 5, 10)).astype(np.float32) h0 = np.random.random((2, 3, 32)).astype(np.float32) torch_tensor = torch.tensor(input) torch_h0 = torch.tensor(h0) torch_gru = torch.nn.GRU(input_size=10, hidden_size=32, num_layers=2, bias=False, batch_first=True, dropout=0.6, bidirectional=False) torch_output, torch_h = torch_gru(torch_tensor, torch_h0) ms_tensor = tensor(input) ms_h0 = tensor(h0) ms_gru = ms_adapter.pytorch.nn.GRU(input_size=10, hidden_size=32, num_layers=2, bias=False, batch_first=True, dropout=0.6, bidirectional=False) ms_output, ms_h = ms_gru(ms_tensor, ms_h0) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_h.shape == ms_h.shape assert torch_output.shape == ms_output.shape def test_gru3(): input = np.random.random((5, 3, 10)).astype(np.float32) torch_tensor = torch.tensor(input) torch_gru = torch.nn.GRU(input_size=10, hidden_size=32, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) torch_output, torch_h = torch_gru(torch_tensor) ms_tensor = tensor(input) ms_gru = ms_adapter.pytorch.nn.GRU(input_size=10, hidden_size=32, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) ms_output, ms_h = ms_gru(ms_tensor) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_h.shape == ms_h.shape assert torch_output.shape == ms_output.shape def test_lstm1(): input = np.random.random((5, 3, 10)).astype(np.float32) h0 = np.random.random((4, 3, 20)).astype(np.float32) c0 = np.random.random((4, 3, 20)).astype(np.float32) torch_tensor = torch.tensor(input) torch_h0 = torch.tensor(h0) torch_c0 = torch.tensor(c0) torch_gru = torch.nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) torch_output, torch_hc = torch_gru(torch_tensor, (torch_h0, torch_c0)) ms_tensor = tensor(input) ms_h0 = tensor(h0) ms_c0 = tensor(c0) ms_gru = ms_adapter.pytorch.nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) ms_output, ms_hc = ms_gru(ms_tensor, (ms_h0, ms_c0)) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_hc[0].shape == ms_hc[0].shape assert torch_hc[1].shape == ms_hc[1].shape assert torch_output.shape == ms_output.shape def test_lstm2(): input = np.random.random((3, 5, 10)).astype(np.float32) h0 = np.random.random((2, 3, 20)).astype(np.float32) c0 = np.random.random((2, 3, 20)).astype(np.float32) torch_tensor = torch.tensor(input) torch_h0 = torch.tensor(h0) torch_c0 = torch.tensor(c0) torch_gru = torch.nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bias=False, batch_first=True, dropout=0.6, bidirectional=False) torch_output, torch_hc = torch_gru(torch_tensor, (torch_h0, torch_c0)) ms_tensor = tensor(input) ms_h0 = tensor(h0) ms_c0 = tensor(c0) ms_gru = ms_adapter.pytorch.nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bias=False, batch_first=True, dropout=0.6, bidirectional=False) ms_output, ms_hc = ms_gru(ms_tensor, (ms_h0, ms_c0)) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_hc[0].shape == ms_hc[0].shape assert torch_hc[1].shape == ms_hc[1].shape assert torch_output.shape == ms_output.shape def test_lstm3(): input = np.random.random((5, 3, 10)).astype(np.float32) torch_tensor = torch.tensor(input) torch_gru = torch.nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) torch_output, torch_hc = torch_gru(torch_tensor) ms_tensor = tensor(input) ms_gru = ms_adapter.pytorch.nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bias=True, batch_first=False, dropout=0.6, bidirectional=True) ms_output, ms_hc = ms_gru(ms_tensor) # Here we can not assert torch_output == ms_output, because the `w` and `b` in RNN are initial randomly assert torch_hc[0].shape == ms_hc[0].shape assert torch_hc[1].shape == ms_hc[1].shape assert torch_output.shape == ms_output.shape def test_rnncell1(): rnn = torch.nn.RNNCell(10, 20) input = torch.randn(6, 3, 10) hx = torch.randn(3, 20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.RNNCell(10, 20) input = ms_torch.randn(6, 3, 10) hx = ms_torch.randn(3, 20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_rnncell2(): rnn = torch.nn.RNNCell(10, 20) input = torch.randn(6, 10) hx = torch.randn(20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.RNNCell(10, 20) input = ms_torch.randn(6, 10) hx = ms_torch.randn(20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_rnncell3(): rnn = torch.nn.RNNCell(10, 20) input = torch.randn(6, 3, 10) output = [] for i in range(6): hx = rnn(input[i]) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.RNNCell(10, 20) input = ms_torch.randn(6, 3, 10) output = [] for i in range(6): hx = rnn(input[i]) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_grucell1(): rnn = torch.nn.GRUCell(10, 20) input = torch.randn(6, 3, 10) hx = torch.randn(3, 20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.GRUCell(10, 20) input = ms_torch.randn(6, 3, 10) hx = ms_torch.randn(3, 20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_grucell2(): rnn = torch.nn.GRUCell(10, 20) input = torch.randn(6, 10) hx = torch.randn(20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.GRUCell(10, 20) input = ms_torch.randn(6, 10) hx = ms_torch.randn(20) output = [] for i in range(6): hx = rnn(input[i], hx) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_grucell3(): rnn = torch.nn.GRUCell(10, 20) input = torch.randn(6, 3, 10) output = [] for i in range(6): hx = rnn(input[i]) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.GRUCell(10, 20) input = ms_torch.randn(6, 3, 10) output = [] for i in range(6): hx = rnn(input[i]) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_lstmcell1(): rnn = torch.nn.LSTMCell(10, 20) # (input_size, hidden_size) input = torch.randn(2, 3, 10) # (time_steps, batch, input_size) hx = torch.randn(3, 20) # (batch, hidden_size) cx = torch.randn(3, 20) output = [] for i in range(input.size()[0]): hx, cx = rnn(input[i], (hx, cx)) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.LSTMCell(10, 20) # (input_size, hidden_size) input = ms_torch.randn(2, 3, 10) # (time_steps, batch, input_size) hx = ms_torch.randn(3, 20) # (batch, hidden_size) cx = ms_torch.randn(3, 20) output = [] for i in range(input.shape[0]): hx, cx = rnn(input[i], (hx, cx)) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_lstmcell2(): rnn = torch.nn.LSTMCell(10, 20) # (input_size, hidden_size) input = torch.randn(2, 10) # (time_steps, batch, input_size) hx = torch.randn(20) # (batch, hidden_size) cx = torch.randn(20) output = [] for i in range(input.shape[0]): hx, cx = rnn(input[i], (hx, cx)) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.LSTMCell(10, 20) # (input_size, hidden_size) input = ms_torch.randn(2, 10) # (time_steps, batch, input_size) hx = ms_torch.randn(20) # (batch, hidden_size) cx = ms_torch.randn(20) output = [] for i in range(input.shape[0]): hx, cx = rnn(input[i], (hx, cx)) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape def test_lstmcell3(): rnn = torch.nn.LSTMCell(10, 20) # (input_size, hidden_size) input = torch.randn(2, 3, 10) # (time_steps, batch, input_size) output = [] for i in range(input.size()[0]): hx, cx = rnn(input[i]) output.append(hx) torch_output = torch.stack(output, dim=0) rnn = ms_torch.nn.LSTMCell(10, 20) # (input_size, hidden_size) input = ms_torch.randn(2, 3, 10) # (time_steps, batch, input_size) output = [] for i in range(input.shape[0]): hx, cx = rnn(input[i]) output.append(hx) ms_output = ms_torch.stack(output, dim=0) assert torch_output.shape == ms_output.shape if __name__ == '__main__': test_rnn1() test_rnn2() test_rnn3() test_gru1() test_gru2() test_gru3() test_lstm1() test_lstm2() test_lstm3() test_rnncell1() test_rnncell2() test_rnncell3() test_grucell1() test_grucell2() test_grucell3() test_lstmcell1() test_lstmcell2() test_lstmcell3() msadapter/testing/ut/pytorch/nn/test_sequential.py000066400000000000000000000020621436713554100230420ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- from ms_adapter.pytorch import nn from collections import OrderedDict # Using Sequential to create a small model. When `model` is run, # input will first be passed to `Conv2d(1,20,5)`. The output of # `Conv2d(1,20,5)` will be used as the input to the first # `ReLU`; the output of the first `ReLU` will become the input # for `Conv2d(20,64,5)`. Finally, the output of # `Conv2d(20,64,5)` will be used as input to the second `ReLU` model1 = nn.Sequential( nn.Conv2d(1,20,5), nn.ReLU(), nn.Conv2d(20,64,5), nn.ReLU() ) print(model1) # Using Sequential with OrderedDict. This is functionally the # same as the above code model2 = nn.Sequential(OrderedDict([ ('conv1', nn.Conv2d(1,20,5)), ('relu1', nn.ReLU()), ('conv2', nn.Conv2d(20,64,5)), ('relu2', nn.ReLU()) ])) print(model2) model3 = nn.Sequential( [nn.Conv2d(1,20,5), nn.ReLU(), nn.Conv2d(20,64,5), nn.ReLU()] ) print(model3) msadapter/testing/ut/pytorch/nn/test_smooth_l1_loss.py000066400000000000000000000034621436713554100236420ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_smoothl1loss1(): a = np.array([1, 2, 3]).astype(np.float32) b = np.array([1, 2, 2]).astype(np.float32) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.SmoothL1Loss()(ms_logits, ms_labels) torch_logits = torch.tensor(a) torch_labels = torch.tensor(b) result_torch = torch.nn.SmoothL1Loss()(torch_logits, torch_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) def test_smoothl1loss2(): a = np.array([1, 2, 3]).astype(np.float32) b = np.array([1, 2, 2]).astype(np.float32) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='sum', beta=0.5)(ms_logits, ms_labels) torch_logits = torch.tensor(a) torch_labels = torch.tensor(b) result_torch = torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='sum', beta=0.5)(torch_logits, torch_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) def test_smoothl1loss3(): a = np.array([1, 2, 3]).astype(np.float32) b = np.array([1, 2, 2]).astype(np.float32) ms_logits = ms_torch.tensor(a) ms_labels = ms_torch.tensor(b) result_ms = ms_torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='none', beta=0.5)(ms_logits, ms_labels) torch_logits = torch.tensor(a) torch_labels = torch.tensor(b) result_torch = torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='none', beta=0.5)(torch_logits, torch_labels) assert np.allclose(result_ms.asnumpy(), result_torch.numpy()) msadapter/testing/ut/pytorch/nn/test_sparse.py000066400000000000000000000123171436713554100221710ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) class TrainNet(ms.nn.Cell): def __init__(self, net): super(TrainNet, self).__init__() self.net = net self.net.set_grad() def construct(self, index): out = self.net(index) loss = out - ms.ops.ones_like(out) loss = loss.sum() loss = ms.ops.square(loss) return loss def test_embedding(): index_np = np.array([[1, 2, 3], [4, 5, 0]]).astype(np.int32) ms_index = ms_torch.tensor(index_np) net = ms_torch.nn.Embedding(10, 2) result_ms = net(ms_index) train_net = TrainNet(net) train_net.set_grad() grad_fn = ms.value_and_grad(train_net, grad_position=None, weights=train_net.trainable_params()) _, grads = grad_fn(ms_index) assert not np.allclose(grads[0][1].asnumpy(), ms.ops.ZerosLike()(grads[0][1]).asnumpy()) torch_index = torch.tensor(index_np) result_torch = torch.nn.Embedding(10, 2)(torch_index) assert result_ms.asnumpy().shape == result_torch.detach().numpy().shape def test_embedding_with_weight(): index_np = np.array([[1, 2, 3], [0, 2, 1]]).astype(np.int32) weight_np = np.array([[2.3, 4.5], [6.0, 7.1], [3.5, 6.], [8.9, 4]]).astype(np.float32) ms_index = ms_torch.tensor(index_np) ms_weight = ms_torch.tensor(weight_np) net = ms_torch.nn.Embedding(4, 2, _weight=ms_weight) result_ms = net(ms_index) train_net = TrainNet(net) train_net.set_grad() grad_fn = ms.value_and_grad(train_net, grad_position=None, weights=train_net.trainable_params()) _, grads = grad_fn(ms_index) assert not np.allclose(grads[0][1].asnumpy(), ms.ops.ZerosLike()(grads[0][1]).asnumpy()) torch_index = torch.tensor(index_np) torch_weight = torch.tensor(weight_np) result_torch = torch.nn.Embedding(4, 2, _weight=torch_weight)(torch_index) assert result_ms.asnumpy().shape == result_torch.detach().numpy().shape assert np.allclose(result_ms.asnumpy(), result_torch.detach().numpy()) def test_embedding_from_pretrained(): index_np = np.array([[1, 2, 3], [0, 2, 1]]).astype(np.int32) weight_np = np.array([[2.3, 4.5], [6.0, 7.1], [3.5, 6.], [8.9, 4]]).astype(np.float32) ms_index = ms_torch.tensor(index_np) ms_weight = ms_torch.tensor(weight_np) net = ms_torch.nn.Embedding.from_pretrained(ms_weight) result_ms = net(ms_index) train_net = TrainNet(net) train_net.set_grad() grad_fn = ms.value_and_grad(train_net, grad_position=None, weights=train_net.trainable_params()) _, grads = grad_fn(ms_index) assert not grads torch_index = torch.tensor(index_np) torch_weight = torch.tensor(weight_np) result_torch = torch.nn.Embedding.from_pretrained(torch_weight)(torch_index) assert result_ms.asnumpy().shape == result_torch.detach().numpy().shape assert np.allclose(result_ms.asnumpy(), result_torch.detach().numpy()) # TODO: padding_idx not support yet ''' def test_embedding_with_padding_idx(): # TODO: Pynative gradient has some problem, they are all zeros. #context.set_context(mode=ms.PYNATIVE_MODE) context.set_context(mode=ms.GRAPH_MODE) index_np = np.array([[1, 2, 3], [4, 5, 0]]).astype(np.int32) ms_index = ms_torch.tensor(index_np) net = ms_torch.nn.Embedding(10, 2, padding_idx=1) result_ms = net(ms_index) train_net = TrainNet(net) train_net.set_grad() grad_fn = ms.value_and_grad(train_net, grad_position=None, weights=train_net.trainable_params()) _, grads = grad_fn(ms_index) assert np.allclose(grads[0][1].asnumpy(), ms.ops.ZerosLike()(grads[0][1]).asnumpy()) torch_index = torch.tensor(index_np) result_torch = torch.nn.Embedding(10, 2, padding_idx=1)(torch_index) assert result_ms.asnumpy().shape == result_torch.detach().numpy().shape def test_embedding_with_weight_with_padding_idx(): # TODO: Pynative gradient has some problem, they are all zeros. #context.set_context(mode=ms.PYNATIVE_MODE) context.set_context(mode=ms.GRAPH_MODE) index_np = np.array([[1, 2, 3], [0, 2, 1]]).astype(np.int32) weight_np = np.array([[2.3, 4.5], [6.0, 7.1], [3.5, 6.], [8.9, 4]]) ms_index = ms_torch.tensor(index_np) ms_weight = ms_torch.tensor(weight_np) net = ms_torch.nn.Embedding(4, 2, _weight=ms_weight, padding_idx=1) result_ms = net(ms_index) train_net = TrainNet(net) train_net.set_grad() grad_fn = ms.value_and_grad(train_net, grad_position=None, weights=train_net.trainable_params()) _, grads = grad_fn(ms_index) assert np.allclose(grads[0][1].asnumpy(), ms.ops.ZerosLike()(grads[0][1]).asnumpy()) torch_index = torch.tensor(index_np) torch_weight = torch.tensor(weight_np) result_torch = torch.nn.Embedding(4, 2, _weight=torch_weight, padding_idx=1)(torch_index) assert result_ms.asnumpy().shape == result_torch.detach().numpy().shape assert np.allclose(result_ms.asnumpy(), result_torch.detach().numpy()) ''' if __name__ == '__main__': test_embedding() test_embedding_with_weight() test_embedding_from_pretrained() #test_embedding_with_padding_idx() #test_embedding_with_weight_with_padding_idx()msadapter/testing/ut/pytorch/nn/test_transformer.py000066400000000000000000000000001436713554100232200ustar00rootroot00000000000000msadapter/testing/ut/pytorch/nn/test_unpooling.py000066400000000000000000000114561436713554100227110ustar00rootroot00000000000000import numpy as np import torch from mindspore import Tensor from ms_adapter.pytorch.nn import MaxUnpool1d, MaxUnpool2d, MaxUnpool3d import mindspore as ms ms.context.set_context(mode=ms.GRAPH_MODE) def test_maxunpool1d_compare1(): kernel_size, stride, padding = 4, 2, 2 ms_net = MaxUnpool1d(kernel_size, stride, padding) torch_net = torch.nn.MaxUnpool1d(kernel_size, stride, padding) B, N, C = 4, 5, 6 data = np.random.random([B, N, C]) indices_range = (C - 1) * stride + kernel_size - 2 * padding indices = np.random.choice(indices_range - 1, size=(1, 1, C), replace=False) indices = indices.repeat(B, 0).repeat(N, 1) ms_input = Tensor(data) ms_indices = Tensor(indices) torch_input = torch.Tensor(data) torch_indices = torch.Tensor(indices).type(torch.int64) torch_output = torch_net(torch_input, torch_indices) ms_output = ms_net(ms_input, ms_indices) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_maxunpool2d_compare1(): kernel_size, stride = 3, 2 ms_net = MaxUnpool2d(kernel_size, stride) torch_net = torch.nn.MaxUnpool2d(kernel_size, stride) B, N, H, W = 4, 5, 6, 7 data = np.random.random([B, N, H, W]) indices_range = (H - 1) * stride + kernel_size indices_range = ((W - 1) * stride - 1 + kernel_size) * indices_range indices = np.random.choice(indices_range - 1, size=(1, 1, H * W), replace=False) indices = indices.repeat(B, 0).repeat(N, 1).reshape(B, N, H, W) ms_input = Tensor(data) ms_indices = Tensor(indices) torch_input = torch.Tensor(data) torch_indices = torch.Tensor(indices).type(torch.int64) torch_output = torch_net(torch_input, torch_indices) ms_output = ms_net(ms_input, ms_indices) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool2d_compare2(): kernel_size, stride, padding = 4, 2, 2 ms_net = MaxUnpool2d(kernel_size, stride, padding) torch_net = torch.nn.MaxUnpool2d(kernel_size, stride, padding) B, N, H, W = 6, 7, 8, 9 data = np.random.random([B, N, H, W]) indices_range = (H - 1) * stride + kernel_size - 2 * padding indices_range = ((W - 1) * stride - 1 + kernel_size - 2 * padding) * indices_range indices = np.random.choice(indices_range - 1, size=(1, 1, H * W), replace=False) indices = indices.repeat(B, 0).repeat(N, 1).reshape(B, N, H, W) ms_input = Tensor(data) ms_indices = Tensor(indices) torch_input = torch.Tensor(data) torch_indices = torch.Tensor(indices).type(torch.int64) torch_output = torch_net(torch_input, torch_indices) ms_output = ms_net(ms_input, ms_indices) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool2d_compare3(): kernel_size, stride, padding = (3, 5), (3, 1), 0 ms_net = MaxUnpool2d(kernel_size, stride, padding) torch_net = torch.nn.MaxUnpool2d(kernel_size, stride, padding) B, N, H, W = 1, 32, 9, 9 data = np.random.random([B, N, H, W]) indices_range = (H - 1) * stride[0] + kernel_size[0] indices_range = ((W - 1) * stride[1] - 1 + kernel_size[1]) * indices_range indices = np.random.choice(indices_range - 1, size=(1, 1, H * W), replace=False) indices = indices.repeat(B, 0).repeat(N, 1).reshape(B, N, H, W) ms_input = Tensor(data) ms_indices = Tensor(indices) torch_input = torch.Tensor(data) torch_indices = torch.Tensor(indices).type(torch.int64) torch_output = torch_net(torch_input, torch_indices) ms_output = ms_net(ms_input, ms_indices) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_maxunpool3d_compare1(): kernel_size, stride, padding = 4, 2, 2 ms_net = MaxUnpool3d(kernel_size, stride, padding) torch_net = torch.nn.MaxUnpool3d(kernel_size, stride, padding) B, C, D, H, W = 4, 5, 6, 6, 6 data = np.random.random([B, C, D, H, W]) indices_range = (D - 1) * stride + kernel_size - 2 * padding indices_range = ((H - 1) * stride - 1 + kernel_size - 2 * padding) * indices_range indices_range = ((W - 1) * stride - 1 + kernel_size - 2 * padding) * indices_range indices = np.random.choice(indices_range - 1, size=(1, 1,D * H * W), replace=False) indices = indices.repeat(B, 0).repeat(C, 1).reshape(B, C, D, H, W) ms_input = Tensor(data) ms_indices = Tensor(indices) torch_input = torch.Tensor(data) torch_indices = torch.Tensor(indices).type(torch.int64) ms_input = Tensor(data.astype(np.float32)) torch_input = torch.Tensor(data) torch_output = torch_net(torch_input, torch_indices) ms_output = ms_net(ms_input, ms_indices) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) if __name__ == '__main__': test_maxunpool1d_compare1() test_maxunpool2d_compare1() test_maxunpool2d_compare2() test_maxunpool2d_compare3() test_maxunpool3d_compare1() msadapter/testing/ut/pytorch/nn/test_upsampling.py000066400000000000000000000044301436713554100230500ustar00rootroot00000000000000import numpy as np import torch import mindspore as ms import ms_adapter.pytorch as ms_torch from ms_adapter.pytorch.nn import Upsample def test_upsample1(): ms.context.set_context(mode=ms.PYNATIVE_MODE) tensor = np.arange(1, 5).reshape((1, 1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_upsample = torch.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) torch_output = torch_upsample(torch_tensor) ms_tensor = ms_torch.tensor(tensor) ms_upsample = Upsample(scale_factor=2, mode="bilinear", align_corners=True) ms_output = ms_upsample(ms_tensor) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_upsample2(): ms.context.set_context(mode=ms.PYNATIVE_MODE) tensor = np.arange(1, 5).reshape((1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(tensor) torch_upsample = torch.nn.Upsample(size=4, mode="linear", align_corners=False) torch_output = torch_upsample(torch_tensor) ms_tensor = ms_torch.tensor(tensor) ms_upsample = Upsample(size=4, mode="linear", align_corners=False) ms_output = ms_upsample(ms_tensor) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_upsampling_nearest2d(): input = np.arange(1, 5).reshape((1, 1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(input) torch_func = torch.nn.UpsamplingNearest2d(scale_factor=2) torch_output = torch_func(torch_tensor) ms_torch_tensor = ms_torch.tensor(input) ms_torch_func = ms_torch.nn.UpsamplingNearest2d(scale_factor=2) ms_torch_output = ms_torch_func(ms_torch_tensor) assert np.allclose(ms_torch_output.numpy(), torch_output.numpy()) def test_upsampling_bilinear2d(): input = np.arange(1, 5).reshape((1, 1, 2, 2)).astype(np.float32) torch_tensor = torch.tensor(input) torch_func = torch.nn.UpsamplingBilinear2d(scale_factor=2) torch_output = torch_func(torch_tensor) ms_torch_tensor = ms_torch.tensor(input) ms_torch_func = ms_torch.nn.UpsamplingBilinear2d(scale_factor=2) ms_torch_output = ms_torch_func(ms_torch_tensor) assert np.allclose(ms_torch_output.numpy(), torch_output.numpy()) if __name__ == '__main__': test_upsample1() test_upsample2() test_upsampling_nearest2d() test_upsampling_bilinear2d() msadapter/testing/ut/pytorch/tensor/000077500000000000000000000000001436713554100201565ustar00rootroot00000000000000msadapter/testing/ut/pytorch/tensor/test_magic_methods.py000066400000000000000000000347361436713554100244070ustar00rootroot00000000000000import numpy as np import mindspore as ms import torch from mindspore import context import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.PYNATIVE_MODE) def test_neg(): class Net(ms_torch.nn.Module): def forward(self, x): return -x net = Net() x = ms_torch.tensor([1.0, 2.0]) out = net(x) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-1.0, -2.0]) def test_pos(): class Net(ms_torch.nn.Module): def forward(self, x): return +x net = Net() x = ms_torch.tensor([-2.0, 3.0]) out = net(x) assert (type(out) is ms_torch.Tensor) assert np.allclose(out.numpy(), [-2.0, 3.0]) def test_invert(): class Net(ms_torch.nn.Module): def forward(self, x): return ~x net = Net() x = ms_torch.tensor([False, True]) out = net(x) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [True, False]) y = ms_torch.tensor([5, -3]) out = net(y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-6, 2]) def test_round(): class Net(ms_torch.nn.Module): def forward(self, x): return round(x) net = Net() x = ms_torch.tensor([0.7, 1.4]) out = net(x) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [1.0, 1.0]) def test_abs(): class Net(ms_torch.nn.Module): def forward(self, x): return abs(x) net = Net() x = ms_torch.tensor([-0.1, -1.0]) out = net(x) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [0.1, 1.0]) def test_add(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x+y net = Net() x = ms_torch.tensor([-0.1, -1.0]) y = 10.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [9.9, 9.0]) def test_radd(): class Net(ms_torch.nn.Module): def forward(self, x, y): return y+x net = Net() x = ms_torch.tensor([-1, 0]) y = 1 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [0, 1]) def test_iadd(): class Net(ms_torch.nn.Module): def forward(self, x, y): x += y return x net = Net() x = ms_torch.tensor([-1, 0]) y = 1 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [0, 1]) def test_and(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x & y net = Net() x = ms_torch.tensor([-1, 0]) y = 1 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [1, 0]) def test_xor(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x ^ y net = Net() x = ms_torch.tensor([-1, 0]) y = 1 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2, 1]) def test_or(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x | y net = Net() x = ms_torch.tensor([-1, 0]) y = 1 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-1, 1]) def test_sub(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x-y net = Net() x = ms_torch.tensor([-1.7, 2.9]) y = 1.3 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-3.0, 1.6]) def test_rsub(): class Net(ms_torch.nn.Module): def forward(self, x, y): return y - x net = Net() x = ms_torch.tensor([-1.7, 2.9]) y = 1.3 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [3.0, -1.6]) def test_isub(): class Net(ms_torch.nn.Module): def forward(self, x, y): x-=y return x net = Net() x = ms_torch.tensor([-1.7, 2.9]) y = 1.3 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-3.0, 1.6]) def test_mul(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x*y net = Net() x = ms_torch.tensor([-1.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 4.0]) def test_rmul(): class Net(ms_torch.nn.Module): def forward(self, x, y): return y*x net = Net() x = ms_torch.tensor([-1.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 4.0]) def test_imul(): class Net(ms_torch.nn.Module): def forward(self, x, y): x *= y return x net = Net() x = ms_torch.tensor([-1.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 4.0]) def test_truediv(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x/y net = Net() x = ms_torch.tensor([-4.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 1.0]) def test_rtruediv(): class Net(ms_torch.nn.Module): def forward(self, x, y): return y/x net = Net() x = ms_torch.tensor([-4.0, 2.0]) y = 8.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 4.0]) def test_floordiv(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x//y net = Net() x = ms_torch.tensor([-4.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 1.0]) def test_rfloordiv(): class Net(ms_torch.nn.Module): def forward(self, x, y): return y//x net = Net() x = ms_torch.tensor([-4.0, 2.0]) y = 8.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 4.0]) def test_ifloordiv(): class Net(ms_torch.nn.Module): def forward(self, x, y): x //=y return x net = Net() x = ms_torch.tensor([-4.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [-2.0, 1.0]) def test_mod(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x%y net = Net() x = ms_torch.tensor([4.0, 2.0]) y = 3.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [1.0, 2.0]) def test_rmod(): class Net(ms_torch.nn.Module): def forward(self, x, y): return y%x net = Net() x = ms_torch.tensor([4.0, 2.0]) y = 5.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [1.0, 1.0]) def test_imod(): class Net(ms_torch.nn.Module): def forward(self, x, y): x %= y return x net = Net() x = ms_torch.tensor([4.0, 2.0]) y = 3.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [1.0, 2.0]) def test_pow(): class Net(ms_torch.nn.Module): def forward(self, x, y): return pow(x, y) net = Net() x = ms_torch.tensor([-1.0, 2.0]) y = 2 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [1.0, 4.0]) def test_rpow(): class Net(ms_torch.nn.Module): def forward(self, x, y): return y**x net = Net() x = ms_torch.tensor([2, 1]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [4.0, 2.0]) def test_lt(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x < y net = Net() x = ms_torch.tensor([3.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [False, False]) def test_le(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x <= y net = Net() x = ms_torch.tensor([3.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [False, True]) def test_gt(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x > y net = Net() x = ms_torch.tensor([3.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [True, False]) def test_ge(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x >= y net = Net() x = ms_torch.tensor([3.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [True, True]) def test_eq(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x == y net = Net() x = ms_torch.tensor([3.0, 2.0]) y = 2.0 out = net(x, y) # assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode # assert np.allclose(out.numpy(), [False, True]) def test_ne(): class Net(ms_torch.nn.Module): def forward(self, x, y): return x != y net = Net() x = ms_torch.tensor([3.0, 2.0]) y = 2.0 out = net(x, y) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [True, False]) def test_getitem(): class Net(ms_torch.nn.Module): def forward(self, x): return x[0] class Net_bool_index(ms_torch.nn.Module): def forward(self, x, index): return x[index] class Net_bool_index_torch(torch.nn.Module): def forward(self, x, index): return x[index] net = Net() net2 = Net_bool_index() net2_torch = Net_bool_index_torch() x = ms_torch.tensor([3.0, 2.0]) out = net(x) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), 3.0) x2_torch = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) x2_ms = ms_torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) x3_torch = torch.tensor([[1, 2], [3, 4]]) x3_ms = ms_torch.tensor([[1, 2], [3, 4]]) index1 = ([[[True, False], [True, False]], [[True, False], [True, False]]], [[True, True], [True, False]], [True, False]) index2 = (([True, True], [True, True]), (True, True, True, True), ([True, False], [True, False]), (True, False, True, False), ([True, False]), (True, False), (True, True, True), (True, True, False), True) for i in range(len(index1)): torch_out = net2_torch(x2_torch, torch.tensor(index1[i])) ms_out = net2(x2_ms, ms_torch.tensor(index1[i])) assert np.allclose(torch_out.numpy(), ms_out.numpy()) assert torch_out.numpy().shape == ms_out.numpy().shape assert (type(ms_out) is ms_torch.Tensor) # Incorrect in Graph mode for i in range(len(index2)): torch_out = net2_torch(x3_torch, index2[i]) ms_out = net2(x3_ms, index2[i]) if torch_out.shape[0] != 0: assert np.allclose(torch_out.numpy(), ms_out.numpy()) assert torch_out.numpy().shape == ms_out.numpy().shape # Incorrect when torch_out.shape[0] is 0 assert (type(ms_out) is ms_torch.Tensor) # Incorrect in Graph mode index3 = [False, False, False, False] x4_torch = torch.tensor([[0.3, 0.2], [0.3, 0.1], [0.2, 0.3], [0.1, 0.3]]) x4_ms = ms_torch.tensor([[0.3, 0.2], [0.3, 0.1], [0.2, 0.3], [0.1, 0.3]]) torch_out = net2_torch(x4_torch, torch.tensor(index3)) ms_out = net2(x4_ms, ms_torch.tensor(index3)) assert np.allclose(torch_out.numpy(), ms_out.numpy()) assert torch_out.numpy().shape == ms_out.numpy().shape x5_torch = torch.tensor([[[0.3], [0.2]], [[0.3], [0.1]], [[0.2], [0.3]], [[0.1], [0.3]]]) x5_ms = ms_torch.tensor([[[0.3], [0.2]], [[0.3], [0.1]], [[0.2], [0.3]], [[0.1], [0.3]]]) torch_out = net2_torch(x5_torch, torch.tensor(index3)) ms_out = net2(x5_ms, ms_torch.tensor(index3)) assert np.allclose(torch_out.numpy(), ms_out.numpy()) assert torch_out.numpy().shape == ms_out.numpy().shape x6_torch = torch.tensor([[3, 2, 4, 1], [3, 1, 4, 0], [2, 3, 4, 1]]) x6_ms = ms_torch.tensor([[3, 2, 4, 1], [3, 1, 4, 0], [2, 3, 4, 1]]) index6_torch = torch.tensor([False, False, False]) index6_ms = ms_torch.tensor([False, False, False]) torch_out = x6_torch[index6_torch] ms_out = x6_ms[index6_ms] torch_out_ = torch_out[:, [0, 2]] ms_out_ = ms_out[:, [0, 2]] assert np.allclose(torch_out_.numpy(), ms_out_.numpy()) assert torch_out_.numpy().shape == ms_out_.numpy().shape def test_setitem(): class Net(ms_torch.nn.Module): def forward(self, x): x[1]=1.0 return x net = Net() x = ms_torch.tensor([3.0, 2.0]) out = net(x) assert (type(out) is ms_torch.Tensor) # Incorrect in Graph mode assert np.allclose(out.numpy(), [3.0, 1.0]) if __name__ == '__main__': test_neg() test_pos() test_invert() test_round() test_abs() test_add() test_iadd() test_radd() test_and() test_xor() test_or() test_sub() test_rsub() test_isub() test_mul() test_rmul() test_imul() test_truediv() test_rtruediv() test_floordiv() test_rfloordiv() test_ifloordiv() test_mod() test_rmod() test_imod() test_pow() test_rpow() test_lt() test_le() test_gt() test_ge() test_eq() test_ne() test_getitem() test_setitem() msadapter/testing/ut/pytorch/tensor/test_storage.py000066400000000000000000000011121436713554100232260ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import pickle import numpy as np import torch import ms_adapter.pytorch as ms_torch import mindspore as ms from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_bytestorage(): a = pickle.dumps((1, 2, 3)) storage = torch.ByteStorage.from_buffer(a) torch_out = torch.ByteTensor(storage) storage = ms_torch.ByteStorage.from_buffer(a) ms_out = ms_torch.ByteTensor(storage) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype msadapter/testing/ut/pytorch/tensor/test_tensor.py000066400000000000000000002324561436713554100231150ustar00rootroot00000000000000from cgi import test import ms_adapter.pytorch as pytorch import numpy as np import torch from mindspore import context import mindspore as ms context.set_context(mode=ms.PYNATIVE_MODE) def test_add_(): value = np.array([[1.0, 1.0],[1.0,1.0]]).astype(np.float32) py_tensor = torch.tensor(value) py_tensor.add_(1.5) ms_tensor = pytorch.tensor(value) ms_tensor.add_(1.5) assert np.allclose(ms_tensor.numpy(), py_tensor.numpy()) def test_mul(): tensor = np.array([[1.0, 1.0],[1.0,1.0]]) ms_tensor = pytorch.tensor(tensor) torch_tensor = torch.tensor(tensor) ms_result = ms_tensor.mul(2.0) torch_result = torch_tensor.mul(2.0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) tensor = np.random.random((3, 3)).astype(np.float32) value = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_value = torch.tensor(value) ms_value = pytorch.tensor(value) torch_output = torch_tensor.mul(torch_value) ms_output = ms_tensor.mul(ms_value) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_mul_(): tensor = np.array([[1.0, 1.0],[1.0,1.0]]) ms_tensor = pytorch.tensor(tensor) torch_tensor = torch.tensor(tensor) ms_tensor.mul_(2.0) torch_tensor.mul_(2.0) assert np.allclose(ms_tensor.numpy(), torch_tensor.numpy()) def test_fill_(): input = np.array([[1.0, 1.0],[1.0,1.0]]).astype(np.float32) py_tensor = torch.tensor(input) py_tensor.fill_(3.0) ms_tensor = pytorch.tensor(input) ms_tensor.fill_(3.0) assert np.allclose(ms_tensor.asnumpy(), py_tensor.numpy()) def test_size_with_dim(): input = np.random.random((3, 4, 5)).astype(np.float32) torch_tensor = torch.tensor(input) ms_tensor = pytorch.tensor(input) torch_size = torch_tensor.size(dim=1) ms_size = ms_tensor.size(dim=1) assert np.allclose(ms_size, torch_size) def test_size(): input = np.random.random((3, 4, 5)).astype(np.float32) torch_tensor = torch.tensor(input) ms_tensor = pytorch.tensor(input) torch_size = torch_tensor.size() ms_size = ms_tensor.size() assert np.allclose(ms_size, torch_size) def test_zero_(): input = np.array([[1.0, 1.0],[1.0,1.0]]) py_tensor = torch.tensor(input) py_tensor.zero_() ms_tensor = pytorch.tensor(input) ms_tensor.zero_() assert np.allclose(ms_tensor.numpy(), py_tensor.numpy()) def test_normal_(): input = np.array([[1.0, 1.0],[1.0,1.0]]) py_tensor = torch.tensor(input) py_tensor.normal_(1.0, 0.5) ms_tensor = pytorch.tensor(input) ms_tensor.normal_(1.0, 0.5) # assert np.allclose(ms_tensor.numpy(), py_tensor.numpy()) # seed is diff # def test_uniform_(): # input = np.array([[1.0, 1.0],[1.0,1.0]]) # py_tensor = torch.tensor(input) # py_tensor.uniform_(0.5, 1.0) # # ms_tensor = pytorch.tensor(input) # ms_tensor.uniform_(0.5, 1.0) # assert np.allclose(ms_tensor.numpy(), py_tensor.numpy()) # seed is diff # def test_erfinv_(): # # only support on Ascend # tensor = pytorch.tensor(np.array([[1.0, 1.0],[1.0,1.0]])) # tensor.erfinv_() # print(tensor) def test_clamp_(): input = np.array([[1.0, 5.0],[10.0, 20.0]]).astype(np.float32) py_tensor = torch.tensor(input) py_tensor.clamp_(min=5.0) ms_tensor = pytorch.tensor(input) ms_tensor.clamp_(min=5.0) assert np.allclose(ms_tensor.numpy(), py_tensor.numpy()) py_tensor_2 = torch.tensor(input) py_tensor_2.clamp_(max=10.0) ms_tensor_2 = pytorch.tensor(input) ms_tensor_2.clamp_(max=10.0) assert np.allclose(ms_tensor_2.numpy(), py_tensor_2.numpy()) py_tensor_3 = torch.tensor(input) py_tensor_3.clamp_(min=10.0, max=5.0) ms_tensor_3 = pytorch.tensor(input) ms_tensor_3.clamp_(min=10.0, max=5.0) assert np.allclose(ms_tensor_3.numpy(), py_tensor_3.numpy()) def test_permute1(): tensor = np.random.random((3, 4, 4)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.permute(1,0,2) ms_output = ms_tensor.permute(1,0,2) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_permute2(): tensor = np.random.random((3, 4, 4)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.permute((1,0,2)) ms_output = ms_tensor.permute((1,0,2)) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_copy_(): tensor = np.random.random((3, 4, 4)).astype(np.float32) src = np.random.random((3, 4, 4)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_src = torch.tensor(src) ms_src = pytorch.tensor(src) torch_tensor.copy_(torch_src) ms_tensor.copy_(ms_src) assert np.allclose(ms_tensor.asnumpy(), torch_tensor.numpy()) def test_expand(): tensor = np.random.random((1, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.expand(4, -1) ms_output = ms_tensor.expand(4, -1) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) torch_tensor = torch.tensor([True, False]) ms_tensor = pytorch.tensor([True, False]) torch_output = torch_tensor.expand(2, 2) ms_output = ms_tensor.expand(2, 2) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_sigmoid(): tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.sigmoid() ms_output = ms_tensor.sigmoid() assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_sigmoid_(): tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_tensor.sigmoid_() ms_tensor.sigmoid_() assert np.allclose(ms_tensor.asnumpy(), torch_tensor.numpy()) def test_float(): tensor = np.random.random((3, 3)).astype(np.int) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.float() ms_output = ms_tensor.float() assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_flip(): tensor = np.random.random((3, 3)).astype(np.int) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.flip(0) ms_output = ms_tensor.flip(0) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) torch_output = torch_tensor.flip([0,1]) ms_output = ms_tensor.flip([0,1]) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_sign(): tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.sign() ms_output = ms_tensor.sign() assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_div(): tensor = np.random.random((3, 3)).astype(np.float32) value = np.random.random((3, 3)).astype(np.float32) torch_value = torch.tensor(value) ms_value = pytorch.tensor(value) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.div(torch_value) ms_output = ms_tensor.div(ms_value) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) def test_div2(): tensor = np.array([3, 3]).astype(np.float32) value = np.array([2, 2]).astype(np.float32) torch_value = torch.tensor(value) ms_value = pytorch.tensor(value) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.div(torch_value, rounding_mode='floor') ms_output = ms_tensor.div(ms_value, rounding_mode='floor') assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) # TODO # mindspore.ops.trunc has some problem on code gate version of mindspore ''' def test_div3(): tensor = np.array([-3, 3]).astype(np.float32) value = np.array([2, 2]).astype(np.float32) torch_value = torch.tensor(value) ms_value = pytorch.tensor(value) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.div(torch_value, rounding_mode='trunc') ms_output = ms_tensor.div(ms_value, rounding_mode='trunc') assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) ''' def test_max(): tensor = np.random.random((3, 3)).astype(np.int32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.max() ms_output = ms_tensor.max() assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().dtype == torch_output.numpy().dtype tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.max(dim = 1) ms_output = ms_tensor.max(dim = 1) assert np.allclose(ms_output[0].asnumpy(), torch_output[0].numpy()) assert np.allclose(ms_output[1].asnumpy(), torch_output[1].numpy()) assert ms_output[0].asnumpy().dtype == torch_output[0].numpy().dtype tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.max(1, True) ms_output = ms_tensor.max(1, True) assert np.allclose(ms_output[0].asnumpy(), torch_output[0].numpy()) assert np.allclose(ms_output[1].asnumpy(), torch_output[1].numpy()) assert ms_output[0].asnumpy().dtype == torch_output[0].numpy().dtype tensor = np.random.random((3, 3)) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.max(1, True) ms_output = ms_tensor.max(1, True) assert np.allclose(ms_output[0].asnumpy(), torch_output[0].numpy()) assert np.allclose(ms_output[1].asnumpy(), torch_output[1].numpy()) assert ms_output[0].asnumpy().dtype == torch_output[0].numpy().dtype def test_numel(): tensor = np.zeros((1, 2, 3, 4, 5)) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.numel() ms_output = ms_tensor.numel() assert np.allclose(ms_output, torch_output) def test_sum(): tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.sum() ms_output = ms_tensor.sum() assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().dtype == torch_output.numpy().dtype tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.sum(dim = 1) ms_output = ms_tensor.sum(dim = 1) assert np.allclose(ms_output.asnumpy(), torch_output.numpy()) assert ms_output.asnumpy().dtype == torch_output.numpy().dtype def test_split(): tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.split(1, 0) ms_output = ms_tensor.split(1, 0) for i in range(3): assert np.allclose(ms_output[i].asnumpy(), torch_output[i].numpy()) assert ms_output[i].asnumpy().dtype == torch_output[i].numpy().dtype def test_numpy(): tensor = np.random.random((3, 3)).astype(np.float32) torch_tensor_np = torch.tensor(tensor).numpy() ms_tensor_np = pytorch.tensor(tensor).numpy() assert np.allclose(ms_tensor_np, torch_tensor_np) assert ms_tensor_np.dtype == torch_tensor_np.dtype def test_ndimension(): x = np.ones((1, 2, 3)) torch_tensor = torch.tensor(x) ms_tensor = pytorch.tensor(x) torch_output = torch_tensor.ndimension() ms_output = ms_tensor.ndimension() assert np.allclose(ms_output, torch_output) def test_pow(): x = np.random.random((2, 3)).astype(np.float32) exponent_float = 2.0 torch_tensor = torch.tensor(x) ms_tensor = pytorch.tensor(x) torch_output = torch_tensor.pow(exponent_float) ms_output = ms_tensor.pow(exponent_float) assert np.allclose(ms_output.numpy(), torch_output.numpy()) x_tensor = np.random.random((1, 4)).astype(np.float32) exponent_tensor = np.arange(1.0, 5.0) torch_tensor = torch.tensor(x_tensor) torch_exp = torch.tensor(exponent_tensor) torch_output = torch_tensor.pow(torch_exp) ms_tensor = pytorch.tensor(x_tensor) ms_exp = pytorch.tensor(exponent_tensor) ms_output = ms_tensor.pow(ms_exp) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_repeat(): torch_x = torch.tensor([1, 2, 3]) torch_out = torch_x.repeat(4, 2) ms_x = pytorch.tensor([1, 2, 3]) ms_out = ms_x.repeat(4, 2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) def test_repeat_interleave(): pytorch_x = torch.tensor([[1, 2, 3], [2, 4, 6]]) pytorch_out = pytorch_x.repeat_interleave(2) ms_x = pytorch.tensor([[1, 2, 3], [2, 4, 6]]) ms_out = ms_x.repeat_interleave(2) assert np.allclose(pytorch_out.numpy(), ms_out.asnumpy()) pytorch_y = torch.tensor([[1, 2], [3, 4]]) pytorch_out = pytorch_y.repeat_interleave(3, dim=1) ms_y = pytorch.tensor([[1, 2], [3, 4]]) ms_out = ms_y.repeat_interleave(3, dim=1) assert np.allclose(pytorch_out.numpy(), ms_out.asnumpy()) pytorch_out = pytorch_y.repeat_interleave(torch.tensor([1, 2]), dim=0) ms_out = ms_y.repeat_interleave(pytorch.tensor([1, 2]), dim=0) assert np.allclose(pytorch_out.numpy(), ms_out.asnumpy()) def test_reshape(): x = np.ones((2, 3, 4)) torch_tensor = torch.tensor(x) torch_output = torch_tensor.reshape(2,3,4) ms_tensor = pytorch.tensor(x) ms_output = ms_tensor.reshape(2,3,4) assert np.allclose(ms_output.numpy(), torch_output.numpy()) torch_output = torch_tensor.reshape((6,4)) ms_output = ms_tensor.reshape((6,4)) assert np.allclose(ms_output.numpy(), torch_output.numpy()) torch_output = torch_tensor.new().reshape((0,1,2,3)) ms_output = ms_tensor.new().reshape((0,1,2,3)) assert np.allclose(ms_output.size(), torch_output.size()) def test_reshape_as(): x = [[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]] y = np.arange(6).reshape(3,2) torch_x = torch.tensor(x, dtype=torch.float32) torch_y = torch.tensor(y) torch_out = torch_x.reshape_as(torch_y) ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_y = pytorch.tensor(y) ms_out = ms_x.reshape_as(ms_y) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_arcsinh(): x = np.array([-5.0, 1.5, 3.0, 100.0]) torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.arcsinh() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.arcsinh() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_arctanh(): x = np.array([-0.5, 0.3, 0.0, 0.99]) torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.arctanh() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.arctanh() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_det(): x = [[[-4.5, -1.5], [7.0, 6.0]], [[2.5, 0.5], [3.0, 9.0]]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.det() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.det() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_negative(): x = [1, 2, -1, 2, 0, -3.5] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.negative() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.negative() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_abs(): x = [-2, -3, 0, -0, -0.0, -1.5, 2.5] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.abs() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.abs() assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_ndim(): x = [[[[[1], [2]]], [[[3], [4]]]]] torch_x = torch.tensor(x) torch_out = torch_x.ndim ms_x = pytorch.tensor(x) ms_out = ms_x.ndim assert ms_out == torch_out def test_amax(): x = np.random.randn(1, 4, 4) torch_x = torch.tensor(x) torch_out = torch_x.amax(None, True) ms_x = pytorch.tensor(x) ms_out = ms_x.amax(None, True) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_amin(): x = np.random.randn(1, 4, 4) torch_x = torch.tensor(x) torch_out = torch_x.amin(-1, False) ms_x = pytorch.tensor(x) ms_out = ms_x.amin(-1, False) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_as_strided(): x = np.arange(0,16).reshape(4, 4).astype(np.int64)*2 torch_x = torch.tensor(x) torch_out = torch_x.as_strided((2, 2, 3), (1, 2, 4), 1) ms_x = pytorch.tensor(x) ms_out = ms_x.as_strided((2, 2, 3), (1, 2, 4), 1) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_bmm(): x = np.random.randn(3, 4, 5).astype(np.float32) y = np.random.randn(3, 5, 2).astype(np.float32) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.bmm(torch_y) ms_x = pytorch.tensor(x) ms_y = pytorch.tensor(y) ms_out = ms_x.bmm(ms_y) assert np.allclose(ms_out.numpy(), torch_out.numpy(), rtol=1e-4, atol=1e-6) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_clamp(): x = np.random.randn(3, 3).astype(np.float64) torch_x = torch.tensor(x) torch_out = torch_x.clamp(0.3, 0.6) ms_x = pytorch.tensor(x) ms_out = ms_x.clamp(0.3, 0.6) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_dim(): x = [[[[[1], [2]]], [[[3], [4]]]]] torch_x = torch.tensor(x) torch_out = torch_x.dim() ms_x = pytorch.tensor(x) ms_out = ms_x.dim() assert ms_out == torch_out def test_expand_as(): x = np.random.randn(3, 3) y = np.random.randn(2, 2, 3, 3) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.expand_as(torch_y) ms_x = pytorch.tensor(x) ms_y = pytorch.tensor(y) ms_out = ms_x.expand_as(ms_y) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_item(): x = np.random.randn(1, 1).astype(np.half) torch_x = torch.tensor(x) torch_out = torch_x.item() ms_x = pytorch.tensor(x) ms_out = ms_x.item() assert ms_out == torch_out assert type(ms_out)==type(torch_out) def test_log(): x = [[0, -1, 1, 2, 2.71828, 3.5]] torch_x = torch.tensor(x) torch_out = torch_x.log() ms_x = pytorch.tensor(x) ms_out = ms_x.log() assert np.allclose(ms_out.numpy(), torch_out.numpy(), rtol=1e-4, atol=1e-5, equal_nan=True) def test_log2(): x = [[0, -1, 1, 2, 2.71828, 3.5]] torch_x = torch.tensor(x) torch_out = torch_x.log2() ms_x = pytorch.tensor(x) ms_out = ms_x.log2() assert np.allclose(ms_out.numpy(), torch_out.numpy(), rtol=1e-4, atol=1e-5, equal_nan=True) def test_matmul(): x = np.random.randn(3, 4, 5).astype(np.float32) y = np.random.randn(5, 2).astype(np.float32) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.matmul(torch_y) ms_x = pytorch.tensor(x) ms_y = pytorch.tensor(y) ms_out = ms_x.matmul(ms_y) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_squeeze(): x = np.random.randn(1, 2, 3, 1) torch_x = torch.tensor(x) torch_out = torch_x.squeeze(-1) ms_x = pytorch.tensor(x) ms_out = ms_x.squeeze(-1) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_stride(): x = np.arange(0,32).reshape(2, 4, 4).astype(np.half) torch_x = torch.tensor(x) torch_out = torch_x.stride(1) ms_x = pytorch.tensor(x) ms_out = ms_x.stride(1) assert ms_out == torch_out assert type(ms_out) == type(torch_out) def test_sub(): x = np.random.randn(3, 4).astype(np.half) y = np.random.randn(3, 4).astype(np.half) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.sub(torch_y, alpha=2) ms_x = pytorch.tensor(x) ms_y = pytorch.tensor(y) ms_out = ms_x.sub(ms_y, alpha=2) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_others_tensor(): tensor = pytorch.ByteTensor() tensor = pytorch.ByteTensor(3, 5) assert tensor.shape == (3, 5) assert tensor.dtype == pytorch.uint8 tensor = pytorch.ByteTensor([1, 2, 3]) assert tensor.shape == (3,) assert tensor.dtype == pytorch.uint8 tensor = pytorch.CharTensor() tensor = pytorch.CharTensor(3, 5) assert tensor.shape == (3, 5) assert tensor.dtype == pytorch.int8 tensor = pytorch.CharTensor([1, 2, 3]) assert tensor.shape == (3,) assert tensor.dtype == pytorch.int8 tensor = pytorch.ShortTensor() tensor = pytorch.ShortTensor(3, 5) assert tensor.shape == (3, 5) assert tensor.dtype == pytorch.int16 tensor = pytorch.ShortTensor([1, 2, 3]) assert tensor.shape == (3,) assert tensor.dtype == pytorch.int16 tensor = pytorch.IntTensor() tensor = pytorch.IntTensor(3, 5) assert tensor.shape == (3, 5) assert tensor.dtype == pytorch.int32 tensor = pytorch.IntTensor([1, 2, 3]) assert tensor.shape == (3,) assert tensor.dtype == pytorch.int32 tensor = pytorch.HalfTensor() tensor = pytorch.HalfTensor(3, 5) assert tensor.shape == (3, 5) assert tensor.dtype == pytorch.float16 tensor = pytorch.HalfTensor([1, 2, 3]) assert tensor.shape == (3,) assert tensor.dtype == pytorch.float16 tensor = pytorch.FloatTensor() tensor = pytorch.FloatTensor(3, 5) assert tensor.shape == (3, 5) assert tensor.dtype == pytorch.float32 tensor = pytorch.FloatTensor([1, 2, 3]) assert tensor.shape == (3,) assert tensor.dtype == pytorch.float32 tensor = pytorch.DoubleTensor() tensor = pytorch.DoubleTensor(3, 5) assert tensor.shape == (3, 5) assert tensor.dtype == pytorch.float64 tensor = pytorch.DoubleTensor([1, 2, 3]) assert tensor.shape == (3,) assert tensor.dtype == pytorch.float64 def test_is_floating_point(): x = [1, 2, -1, 2, 0, -3.5] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.is_floating_point() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.is_floating_point() assert np.allclose(ms_out, torch_out) torch_x = torch.tensor(x, dtype=torch.int32) torch_out = torch_x.is_floating_point() ms_x = pytorch.tensor(x, dtype=pytorch.int32) ms_out = ms_x.is_floating_point() assert np.allclose(ms_out, torch_out) def test_unbind(): x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.unbind(1) ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.unbind(1) for ms_out_, torch_out_ in zip(ms_out, torch_out): assert np.allclose(ms_out_.numpy(), torch_out_.numpy()) def test_unsqueeze(): x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.unsqueeze(1) ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.unsqueeze(1) assert np.allclose(ms_out.numpy(), torch_out.numpy()) # def test_is_signed(): # x = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] # torch_x = torch.tensor(x, dtype=torch.float32) # torch_out = torch_x.is_signed() # ms_x = pytorch.tensor(x, dtype=pytorch.float32) # ms_out = ms_x.is_signed() # np.allclose(ms_out, torch_out) def test_transpose1(): tensor = np.random.random((3, 4, 4)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.transpose(0, 2) ms_output = ms_tensor.transpose(0, 2) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_transpose2(): tensor = np.random.random((1, 3, 4, 4)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.transpose(1, 2) ms_output = ms_tensor.transpose(1, 2) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_floor(): tensor = np.random.random((1, 3, 4, 4)).astype(np.float32) torch_tensor = torch.tensor(tensor) ms_tensor = pytorch.tensor(tensor) torch_output = torch_tensor.floor() ms_output = ms_tensor.floor() assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_isfinite(): x = [1, float('inf'), 2, float('-inf'), float('nan')] torch_x = torch.tensor(x) torch_out = torch_x.isfinite() ms_x = pytorch.tensor(x) ms_out = ms_x.isfinite() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_isnan(): x = [1, float('nan'), 2] torch_x = torch.tensor(x) torch_out = torch_x.isnan() ms_x = pytorch.tensor(x) ms_out = ms_x.isnan() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_expand(): x = [[1], [2], [3]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.expand(3, 4) ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.expand(3, 4) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_clone(): def fun(x): b = x return 2 * x + 3 * b x = [8] torch_x = torch.tensor(x, dtype=torch.float32, requires_grad=True) torch_out = fun(torch_x) ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = fun(ms_x) assert np.allclose(ms_out.asnumpy(), torch_out.detach().numpy()) torch_out.backward() torch_grad = torch_x.grad ms_grad = ms.grad(fun)(ms_x) assert np.allclose(torch_grad.detach().numpy(), ms_grad.asnumpy()) def test_detach(): def fun(x): b = x.detach() return 2 * x + 3 * b x = [8] torch_x = torch.tensor(x, dtype=torch.float32, requires_grad=True) torch_out = fun(torch_x) ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = fun(ms_x) assert np.allclose(ms_out.asnumpy(), torch_out.detach().numpy()) torch_out.backward() torch_grad = torch_x.grad ms_grad = ms.grad(fun)(ms_x) assert np.allclose(torch_grad.detach().numpy(), ms_grad.asnumpy()) def test_new_zeros(): x = [[1], [2], [3]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.new_zeros((3, 5), dtype=torch.int32) ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_out = ms_x.new_zeros((3, 5), dtype=pytorch.int32) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_sort(): x = [[0, 1, 2, 0, 1, 2], [5, 4, 3, 5, 4, 3]] torch_x = torch.tensor(x, dtype=torch.float32) torch_x_int = torch.tensor(x, dtype=torch.int64) torch_out1, torch_out2 = torch_x.sort() torch_out1_int, torch_out2_int = torch_x_int.sort() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_x_int = pytorch.tensor(x, dtype=pytorch.int64) ms_out1, ms_out2 = ms_x.sort() ms_out1_int, ms_out2_int = ms_x_int.sort() assert np.allclose(torch_out1.numpy(), ms_out1.numpy()) assert np.allclose(ms_out2.numpy(), torch_out2.numpy()) assert np.allclose(torch_out1_int.numpy(), ms_out1_int.numpy()) assert np.allclose(ms_out2_int.numpy(), torch_out2_int.numpy()) assert torch_out1_int.numpy().dtype == ms_out1_int.numpy().dtype assert torch_out2_int.numpy().dtype == ms_out2_int.numpy().dtype torch_out1, torch_out2 = torch_x.sort(0, True) ms_out1, ms_out2 = ms_x.sort(0, True) assert np.allclose(torch_out1.numpy(), ms_out1.numpy()) assert np.allclose(ms_out2.numpy(), torch_out2.numpy()) def test_msort(): x = [[0, 1, 2, 0, 1, 2], [5, 4, 3, 5, 4, 3]] torch_x = torch.tensor(x, dtype=torch.float32) torch_x_int = torch.tensor(x, dtype=torch.int64) torch_out = torch_x.msort() torch_out_int = torch_x_int.msort() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_x_int = pytorch.tensor(x, dtype=pytorch.int64) ms_out = ms_x.msort() ms_out_int = ms_x_int.msort() assert np.allclose(torch_out.numpy(), ms_out.numpy()) assert np.allclose(torch_out_int.numpy(), ms_out_int.numpy()) assert torch_out_int.numpy().dtype == ms_out_int.numpy().dtype def test_argsort(): x = [[0, 1, 2, 0, 1, 2], [5, 4, 3, 5, 4, 3]] torch_x = torch.tensor(x, dtype=torch.float32) torch_x_int = torch.tensor(x, dtype=torch.int64) torch_out = torch_x.argsort(-1, True) torch_out_int = torch_x_int.argsort() ms_x = pytorch.tensor(x, dtype=pytorch.float32) ms_x_int = pytorch.tensor(x, dtype=pytorch.int64) ms_out = ms_x.argsort(-1, True) ms_out_int = ms_x_int.argsort() assert np.allclose(torch_out.numpy(), ms_out.numpy()) assert np.allclose(torch_out_int.numpy(), ms_out_int.numpy()) assert torch_out_int.numpy().dtype == ms_out_int.numpy().dtype def test_sqrt(): x = [2.0755, 1.0226, 0.0831, 0.4806] torch_x = torch.tensor(x, dtype=torch.float64) torch_out = torch_x.sqrt() ms_x = pytorch.tensor(x, dtype=pytorch.float64) ms_out = ms_x.sqrt() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_rsqrt1(): x = [2.0755, 1.0226, 0.0831, 0.4806] torch_x = torch.tensor(x, dtype=torch.float64) torch_out = torch_x.rsqrt() ms_x = pytorch.tensor(x, dtype=pytorch.float64) ms_out = ms_x.rsqrt() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_rsqrt2(): x = [2, 1, 3, 6] torch_x = torch.tensor(x, dtype=torch.int64) torch_out = torch_x.rsqrt() ms_x = pytorch.tensor(x, dtype=pytorch.int64) ms_out = ms_x.rsqrt() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_rsqrt3(): x = [True, False, True, True] torch_x = torch.tensor(x, dtype=torch.bool) torch_out = torch_x.rsqrt() ms_x = pytorch.tensor(x, dtype=pytorch.bool) ms_out = ms_x.rsqrt() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_to(): ms_x = pytorch.tensor([[1], [2], [3]], dtype=pytorch.float32) ms_y = pytorch.tensor([2], dtype=pytorch.int32) assert ms_x.dtype == pytorch.float32 ms_x = ms_x.to(ms_y) assert ms_x.dtype == pytorch.int32 ms_x = ms_x.to("cpu") assert ms_x.dtype == pytorch.int32 ms_x = ms_x.to("cpu", pytorch.float32) assert ms_x.dtype == pytorch.float32 ms_x = ms_x.to(other=ms_y) assert ms_x.dtype == pytorch.int32 ms_x = ms_x.to(dtype=pytorch.float32) assert ms_x.dtype == pytorch.float32 ms_x = ms_x.to(device="CUDA", dtype=pytorch.int32) assert ms_x.dtype == pytorch.int32 ms_x = ms_x.to("CUDA", dtype=pytorch.float32) assert ms_x.dtype == pytorch.float32 ms_x = ms_x.to(pytorch.int32) assert ms_x.dtype == pytorch.int32 ms_x = ms_x.to(device="CPU") assert ms_x.dtype == pytorch.int32 def test_resize(): x = np.array([[1, 2, 3], [4, 5, 6]]) torch_x = torch.tensor(x) torch_out = torch_x.resize(3, 2) ms_x = pytorch.tensor(x) ms_out = ms_x.resize(3, 2) assert np.allclose(torch_out.numpy(), ms_out.numpy()) torch_out = torch_x.new().resize_(2, 2) ms_out = ms_x.new().resize_(2, 2) assert np.allclose(torch_out.size(), ms_out.size()) def test_resize_as(): # In torch.Tensor.resize_as(input), input.numel() should be equal to origin tensor's # So here, we don't test the need padding case a = np.array([[1, 2, 3],[4, 5, 6]]) b = np.array([[1, 2],[3, 4], [5, 6]]) c = np.array([1, 2, 3, 4, 5, 6]) ms_a = pytorch.tensor(a) ms_b = pytorch.tensor(b) ms_c = pytorch.tensor(c) ms_out1 = ms_a.resize_as(ms_b) ms_out2 = ms_a.resize_as(ms_c) torch_a = torch.tensor(a) torch_b = torch.tensor(b) torch_c = torch.tensor(c) torch_out1 = torch_a.resize_as(torch_b) torch_out2 = torch_a.resize_as(torch_c) assert np.allclose(torch_out1.numpy(), ms_out1.numpy()) assert np.allclose(torch_out2.numpy(), ms_out2.numpy()) torch_out = ms_c.new().resize_as_(ms_b) ms_out = torch_c.new().resize_as_(torch_b) assert np.allclose(torch_out.size(), ms_out.size()) def test_index_fill(): x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) index = np.array([0, 2]) ms_x = pytorch.tensor(x) ms_index = pytorch.tensor(index) ms_out = ms_x.index_fill(1, ms_index, -1) torch_x = torch.tensor(x) torch_index = torch.tensor(index, dtype=torch.int64) torch_out = torch_x.index_fill(1, torch_index, -1) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_data(): x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) def func(x): return x.data * x.data ms_x = pytorch.tensor(x) ms_x_data = ms_x.data grad_func = ms.ops.grad(func) grads = grad_func(ms_x) torch_x = torch.tensor(x) torch_x_data = torch_x.data assert np.allclose(ms_x_data.numpy(), torch_x_data.numpy()) assert np.allclose(grads.asnumpy(), np.zeros((3, 3))) def test_new(): x = np.array([1, 2]) ms_x = pytorch.tensor(x) torch_x = torch.tensor(x) ms_out = ms_x.new(3, 4, 5) torch_out = torch_x.new(3, 4, 5) assert ms_out.shape == torch_out.shape assert ms_out.numpy().dtype == torch_out.numpy().dtype ms_out = ms_x.new([1, 2, 3, 4]) torch_out = torch_x.new([1, 2, 3, 4]) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype ms_out = ms_x.new() torch_out = torch_x.new() assert ms_out.shape == torch_out.shape assert ms_out.numpy().dtype == torch_out.numpy().dtype ms_out = ms_x.new(x) torch_out = torch_x.new(x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_cuda(): x = pytorch.tensor([1]) x.cuda(device=None, non_blocking=False, memory_format=None) def test_t1(): x = np.random.randn(3, 4) torch_x = torch.tensor(x) ms_x = pytorch.tensor(x) torch_out = torch_x.t() ms_out = ms_x.t() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_t2(): x = np.random.randn(4) torch_x = torch.tensor(x) ms_x = pytorch.tensor(x) torch_out = torch_x.t() ms_out = ms_x.t() assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_T(): x = np.random.randn(2,3,4) torch_x = torch.tensor(x) ms_x = pytorch.tensor(x) torch_out = torch_x.T ms_out = ms_x.T assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_mean1(): ms_tensor = pytorch.tensor([[1., 2, 3], [1, 2, 3]]) ms_result = ms_tensor.mean(dim=0) torch_tensor = torch.tensor([[1., 2, 3], [1, 2, 3]]) torch_result = torch_tensor.mean(dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_mean2(): ms_tensor = pytorch.tensor([[1., 2, 3], [1, 2, 3]]) ms_result = ms_tensor.mean(dim=(0, -1)) torch_tensor = torch.tensor([[1., 2, 3], [1, 2, 3]]) torch_result = torch_tensor.mean(dim=(0, -1)) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_prod1(): ms_tensor = pytorch.tensor([[1., 2, 3], [1, 2, 3]]) ms_result = ms_tensor.prod(dim=0) torch_tensor = torch.tensor([[1., 2, 3], [1, 2, 3]]) torch_result = torch_tensor.prod(dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_index_select(): data = np.random.randn(3, 4 ,5) x_torch = torch.tensor(data) indices = torch.tensor([0, 2]) torch_out = x_torch.index_select(1, indices) x_ms = pytorch.tensor(data) indices = pytorch.tensor([0, 2]) ms_out = x_ms.index_select(1, indices) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def nonzero(): np_array = np.random.random((2,3,4)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.nonzero() ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.nonzero() assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def bool(): np_array = np.random.random((2, 3, 4)).astype(np.float32) torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.bool() ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.bool() assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_std(): x = np.random.randn(2, 2, 3, 4) x2 = np.random.randn(2, 1, 3, 4) #x2 = np.array([[[[1-1j, -2j]], [[3+4j, 5]]]]) torch_x = torch.tensor(x) torch_x2 = torch.tensor(x2) ms_x = pytorch.tensor(x) ms_x2 = pytorch.tensor(x2) torch_out1 = torch_x.std(-2, False) ms_out1 = ms_x.std(-2, False) torch_out2 = torch_x2.std((3, 1), True, True) ms_out2 = ms_x2.std((3, 1), True, True) torch_out3 = torch_x.std() ms_out3 = ms_x.std() assert np.allclose(ms_out1.numpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.numpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.numpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_exp(): x = np.random.randn(2, 2, 3, 4) torch_x = torch.tensor(x) torch_out = torch_x.exp() ms_x = pytorch.tensor(x) ms_out = ms_x.exp() assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_masked_fill(): x = np.random.randn(2, 2, 3, 4) mask = np.random.randn(2, 2, 3, 4) * 2 mask = mask.astype(np.int16).astype(np.bool8) torch_x = torch.tensor(x) torch_mask = torch.tensor(mask) torch_out = torch_x.masked_fill(torch_mask, 2) ms_x = pytorch.tensor(x) ms_mask = pytorch.tensor(mask) ms_out = ms_x.masked_fill(ms_mask, 2) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_tolist(): ms_tensor = pytorch.tensor([[1., 2, 3], [1, 2, 3]]) ms_result = ms_tensor.tolist() torch_tensor = torch.tensor([[1., 2, 3], [1, 2, 3]]) torch_result = torch_tensor.tolist() assert np.allclose(ms_result, torch_result) ms_tensor = ms_tensor[0, 0] ms_result = ms_tensor.tolist() torch_tensor = torch_tensor[0, 0] torch_result = torch_tensor.tolist() assert ms_result == torch_result def test_chunk(): x1 = np.arange(11) x2 = np.arange(12) x3 = np.arange(13) x4 = np.arange(24).reshape(2, 12) torch_x1 = torch.tensor(x1) torch_x2 = torch.tensor(x2) torch_x3 = torch.tensor(x3) torch_x4 = torch.tensor(x4) torch_out1 = torch_x1.chunk(6) torch_out2 = torch_x2.chunk(6) torch_out3 = torch_x3.chunk(6) torch_out4 = torch_x4.chunk(6, 1) torch_out = [torch_out1, torch_out2, torch_out3, torch_out4] ms_x1 = pytorch.tensor(x1) ms_x2 = pytorch.tensor(x2) ms_x3 = pytorch.tensor(x3) ms_x4 = pytorch.tensor(x4) ms_out1 = ms_x1.chunk(6) ms_out2 = ms_x2.chunk(6) ms_out3 = ms_x3.chunk(6) ms_out4 = ms_x4.chunk(6, 1) ms_out = [ms_out1, ms_out2, ms_out3, ms_out4] for i in range(4): assert len(ms_out[i]) == len(torch_out[i]) for j in range(len(ms_out[i])): assert np.allclose(ms_out[i][j].numpy(), torch_out[i][j].numpy()) assert np.allclose(ms_out[i][j].shape, torch_out[i][j].shape) def test_cumsum(): x = [-0.8286, -0.4890, 0.5155, 0.8443, 0.1865, -0.1752, -2.0595, 0.1850, -1.1571, -0.4243] torch_x = torch.tensor(x) torch_out = torch_x.cumsum(0) ms_x = pytorch.tensor(x) ms_out = ms_x.cumsum(0) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_absolute(): x = [-1, -2, 3] torch_x = torch.tensor(x) torch_out = torch_x.absolute_() ms_x = pytorch.tensor(x) ms_out = ms_x.absolute_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_acos(): x = [ 0.3348, -0.5889, 0.2005, -0.1584] torch_x = torch.tensor(x) torch_out = torch_x.acos_() ms_x = pytorch.tensor(x) ms_out = ms_x.acos_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_addcdiv(): t = np.random.rand(1, 3) t1 = np.random.rand(3, 1) t2 = np.random.rand(1, 3) torch_t = torch.tensor(t) torch_t1 = torch.tensor(t1) torch_t2 = torch.tensor(t2) torch_out = torch_t.addcdiv(torch_t1, torch_t2, value=0.1) ms_t = pytorch.tensor(t) ms_t1 = pytorch.tensor(t1) ms_t2 = pytorch.tensor(t2) ms_out = ms_t.addcdiv(ms_t1, ms_t2, value=0.1) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_asinh(): x = [ 0.1606, -1.4267, -1.0899, -1.0250 ] torch_x = torch.tensor(x) torch_out = torch_x.asinh_() ms_x = pytorch.tensor(x) ms_out = ms_x.asinh_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_atanh(): x = [-0.9385, 0.2968, -0.8591, -0.1871] torch_x = torch.tensor(x) torch_out = torch_x.atanh_() ms_x = pytorch.tensor(x) ms_out = ms_x.atanh_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_gather(): input = [[1, 2], [3, 4]] index = [[0, 0], [1, 1]] torch_input = torch.tensor(input) torch_index = torch.tensor(index) torch_out = torch.gather(torch_input, 0, torch_index) ms_input = pytorch.tensor(input) ms_index = pytorch.tensor(index) ms_out = pytorch.gather(ms_input, 0, ms_index) assert np.allclose(torch_out.size(), ms_out.size()) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_fmod(): x = [-3., -2, -1, 1, 2, 3] torch_x = torch.tensor(x) torch_out = torch_x.fmod(2) ms_x = pytorch.tensor(x) ms_out = ms_x.fmod(2) assert np.allclose(torch_out.numpy(), ms_out.numpy()) x = [1, 2, 3, 4, 5] divisor = [-1.5, 1.0, -1.0, -2.0, 2.5] torch_x = torch.tensor(x) tourch_divisor = torch.tensor(divisor) torch_out = torch_x.fmod(tourch_divisor) ms_x = pytorch.tensor(x) ms_divisor = pytorch.tensor(divisor) ms_out = ms_x.fmod(ms_divisor) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_absolute(): x = [-1, -2, 3] torch_x = torch.tensor(x) torch_out = torch_x.absolute_() ms_x = pytorch.tensor(x) ms_out = ms_x.absolute_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_acos(): x = [ 0.3348, -0.5889, 0.2005, -0.1584] torch_x = torch.tensor(x) torch_out = torch_x.acos_() ms_x = pytorch.tensor(x) ms_out = ms_x.acos_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_addcdiv(): t = np.random.rand(1, 3) t1 = np.random.rand(3, 1) t2 = np.random.rand(1, 3) torch_t = torch.tensor(t) torch_t1 = torch.tensor(t1) torch_t2 = torch.tensor(t2) torch_out = torch_t.addcdiv(torch_t1, torch_t2, value=0.1) ms_t = pytorch.tensor(t) ms_t1 = pytorch.tensor(t1) ms_t2 = pytorch.tensor(t2) ms_out = ms_t.addcdiv(ms_t1, ms_t2, value=0.1) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_asinh(): x = [ 0.1606, -1.4267, -1.0899, -1.0250 ] torch_x = torch.tensor(x) torch_out = torch_x.asinh_() ms_x = pytorch.tensor(x) ms_out = ms_x.asinh_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_atanh(): x = [-0.9385, 0.2968, -0.8591, -0.1871] torch_x = torch.tensor(x) torch_out = torch_x.atanh_() ms_x = pytorch.tensor(x) ms_out = ms_x.atanh_() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_argmin(): x = np.random.randn(4, 4) torch_x = torch.tensor(x) torch_out = torch_x.argmin(-1, False) ms_x = pytorch.tensor(x) ms_out = ms_x.argmin(-1, False) assert np.allclose(ms_out.numpy(), torch_out.numpy()) # TODO: assert ms_out.dtype == torch_out.dtype # torch_out = torch_x.argmin(0, True) # ms_out = ms_x.argmin(0, True) # assert np.allclose(ms_out.numpy(), torch_out.numpy()) # torch_out = torch_x.argmin() # ms_out = ms_x.argmin() # assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_argmax(): x = np.random.randn(1, 4, 4) torch_x = torch.tensor(x) torch_out = torch_x.argmax(-1, False) ms_x = pytorch.tensor(x) ms_out = ms_x.argmax(-1, False) assert np.allclose(ms_out.numpy(), torch_out.numpy()) # TODO: assert ms_out.asnumpy().dtype == torch_out.numpy().dtype # torch_out = torch_x.argmin(0, True) # ms_out = ms_x.argmin(0, True) # assert np.allclose(ms_out.numpy(), torch_out.numpy()) # torch_out = torch_x.argmin() # ms_out = ms_x.argmin() # assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_baddbmm(): x = np.ones([1, 3, 3]) b1 = np.ones([1, 3, 4]) b2 = np.ones([1, 4, 3]) torch_input = torch.tensor(x, dtype=torch.float32) torch_batch1 = torch.tensor(b1, dtype=torch.float32) torch_batch2 = torch.tensor(b2, dtype=torch.float32) torch_out = torch_input.baddbmm(torch_batch1, torch_batch2) ms_torch_input = pytorch.tensor(x, dtype=pytorch.float32) ms_torch_batch1 = pytorch.tensor(b1, dtype=pytorch.float32) ms_torch_batch2 = pytorch.tensor(b2, dtype=pytorch.float32) ms_torch_out = ms_torch_input.baddbmm(ms_torch_batch1, ms_torch_batch2) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_topk(): x = np.random.randn(10).astype(np.float32) torch_x = torch.tensor(x) torch_v, torch_i = torch_x.topk(3) ms_x = pytorch.tensor(x) ms_v, ms_i = ms_x.topk(3) assert np.allclose(torch_v.numpy(), ms_v.numpy()) assert np.allclose(torch_i.numpy(), ms_i.numpy()) def test_slice(): x = np.ones([2, 3]) torch_input = torch.tensor(x, dtype=torch.float32) ms_torch_input = pytorch.tensor(x, dtype=pytorch.float32) torch_input[0][0] = 2 ms_torch_input[0][0] = 2 assert np.allclose(torch_input.numpy(), ms_torch_input.numpy()) mask = [True, False] torch_out = torch_input[mask] ms_torch_out = ms_torch_input[mask] assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_maximum(): a = [1, 2, -1] b = [3, 0, 4] torch_a = torch.tensor(a) torch_b = torch.tensor(b) torch_out = torch_a.maximum(torch_b) ms_torch_a = pytorch.tensor(a) ms_torch_b = pytorch.tensor(b) ms_torch_out = ms_torch_a.maximum(ms_torch_b) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) a = [float('nan'), float('inf'), -float('inf'), 3.14] b = [-1, 0, 1, 5] torch_a = torch.tensor(a) torch_b = torch.tensor(b) torch_out = torch_a.maximum(torch_b) ms_torch_a = pytorch.tensor(a) ms_torch_b = pytorch.tensor(b) ms_torch_out = ms_torch_a.maximum(ms_torch_b) # NAN is diff # assert np.allclose(torch_out.numpy(), ms_torch_out.numpy(), equal_nan=True) def test_minimum(): a = [1, 2, -1] b = [3, 0, 4] torch_a = torch.tensor(a) torch_b = torch.tensor(b) torch_out = torch_a.minimum(torch_b) ms_torch_a = pytorch.tensor(a) ms_torch_b = pytorch.tensor(b) ms_torch_out = ms_torch_a.minimum(ms_torch_b) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) a = [float('nan'), float('inf'), -float('inf'), 3.14] b = [-1, 0, 1, 5] torch_a = torch.tensor(a) torch_b = torch.tensor(b) torch_out = torch_a.minimum(torch_b) ms_torch_a = pytorch.tensor(a) ms_torch_b = pytorch.tensor(b) ms_torch_out = ms_torch_a.minimum(ms_torch_b) # NAN is diff # assert np.allclose(torch_out.numpy(), ms_torch_out.numpy(), equal_nan=True) def test_multiply(): x = [0.2015, -0.4255, 2.6087] y = 100 torch_x = torch.tensor(x) torch_out = torch_x.multiply(y) ms_x = pytorch.tensor(x) ms_out = ms_x.multiply(y) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_neg(): x = [0.0090, -0.2262, -0.0682, -0.2866, 0.3940] torch_x = torch.tensor(x) torch_out = torch_x.neg() ms_x = pytorch.tensor(x) ms_out = ms_x.neg() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_ravel(): x = [[[1, 2],[3, 4]], [[5, 6], [7, 8]]] torch_x = torch.tensor(x) torch_out = torch_x.ravel() ms_x = pytorch.tensor(x) ms_out = ms_x.ravel() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_select(): x = np.random.rand(1,2,3,4) torch_x = torch.tensor(x) torch_out = torch_x.select(dim=1, index=1) ms_x = pytorch.tensor(x) ms_out = ms_x.select(dim=1, index=1) assert np.allclose(torch_out.shape, ms_out.shape) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_square(): x = [-2.0755, 1.0226, 0.0831, 0.4806] torch_x = torch.tensor(x) torch_out = torch_x.square() ms_x = pytorch.tensor(x) ms_out = ms_x.square() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_broadcast_to(): x = [1, 2, 3] torch_x = torch.tensor(x) torch_out = torch_x.broadcast_to((3,3)) ms_x = pytorch.tensor(x) ms_out = ms_x.broadcast_to((3,3)) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_divide(): x = [ 0.3810, 1.2774, -0.2972, -0.3719, 0.4637] y = 0.5 torch_x = torch.tensor(x) torch_out = torch_x.divide(y) ms_x = pytorch.tensor(x) ms_out = ms_x.divide(y) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_mm(): x = np.random.rand(2, 3) y = np.random.rand(3, 4) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.mm(torch_y) ms_x = pytorch.tensor(x) ms_y = pytorch.tensor(y) ms_out = ms_x.mm(ms_y) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_view(): torch_x = torch.Tensor(2,3) torch_out = torch_x.view((3,2)) ms_torch_x = pytorch.Tensor(2,3) ms_torch_out = ms_torch_x.view((3,2)) assert np.allclose(torch_out.size(), ms_torch_out.size()) torch_out = torch_x.new().view((-1,2,3)) ms_torch_out = ms_torch_x.new().view((-1,2,3)) assert np.allclose(torch_out.size(), ms_torch_out.size()) def test_view_as(): torch_x = torch.Tensor(2, 3) other = torch.Tensor(3, 2) torch_out = torch_x.view_as(other) ms_torch_x = pytorch.Tensor(2, 3) other = torch.Tensor(3, 2) ms_torch_out = ms_torch_x.view_as(other) assert np.allclose(torch_out.size(), ms_torch_out.size()) def test_logsumexp(): x = np.random.randn(3, 4, 5).astype(np.float32) torch_x = torch.tensor(x) torch_out = torch_x.logsumexp(dim=1, keepdim=False) ms_torch_x = pytorch.tensor(x) ms_torch_out = ms_torch_x.logsumexp(dim=1, keepdim=False) assert np.allclose(torch_out.size(), ms_torch_out.size()) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) x = np.random.randn(2, 3, 2).astype(np.int32) torch_x = torch.tensor(x) torch_out = torch_x.logsumexp(dim=2, keepdim=True) ms_torch_x = pytorch.tensor(x) ms_torch_out = ms_torch_x.logsumexp(dim=2, keepdim=True) assert np.allclose(torch_out.size(), ms_torch_out.size()) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy(), rtol=1e-04) def test_addmv(): M = np.random.randn(2).astype(np.float32) mat = np.random.randn(2, 3).astype(np.float32) vec = np.random.randn(3).astype(np.float32) torch_M = torch.tensor(M) torch_mat = torch.tensor(mat) torch_vec = torch.tensor(vec) torch_out = torch_M.addmv(torch_mat, torch_vec, beta=False, alpha=0.5) ms_torch_M = pytorch.tensor(M) ms_torch_mat = pytorch.tensor(mat) ms_torch_vec = pytorch.tensor(vec) ms_torch_out = ms_torch_M.addmv(ms_torch_mat, ms_torch_vec, beta=False, alpha=0.5) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_dot(): x = [2, 3] y = [2, 1] torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.dot(torch_y) ms_torch_x = pytorch.tensor(x) ms_torch_y = pytorch.tensor(y) ms_torch_out = ms_torch_x.dot(ms_torch_y) assert np.allclose(torch_out.numpy(), ms_torch_out.numpy()) def test_inverse(): A1 = np.random.randn(4, 4).astype(np.float64) torch_A1 = torch.tensor(A1) torch_out1 = torch_A1.inverse() ms_torch_A1 = pytorch.tensor(A1) ms_torch_out1 = ms_torch_A1.inverse() assert np.allclose(torch_out1.numpy(), ms_torch_out1.numpy()) A2 = np.random.randn(2, 3, 4, 4).astype(np.float32) torch_A2 = torch.tensor(A2) torch_out2 = torch_A2.inverse() ms_torch_A2 = pytorch.tensor(A2) ms_torch_out2 = ms_torch_A2.inverse() assert np.allclose(torch_out2.numpy(), ms_torch_out2.numpy(), rtol=1e-3, atol=1e-06) def test_asin(): a1 = np.random.randn(4).astype(np.int32) torch_a1 = torch.tensor(a1) torch_out1 = torch_a1.asin() ms_torch_a1 = pytorch.tensor(a1) ms_torch_out1 = ms_torch_a1.asin() assert np.allclose(torch_out1.numpy(), ms_torch_out1.numpy(), equal_nan=True, atol=1e-05) a2 = np.random.randn(3, 2).astype(np.float32) torch_a2 = torch.tensor(a2) torch_out2 = torch_a2.asin() ms_torch_a2 = pytorch.tensor(a2) ms_torch_out2 = ms_torch_a2.asin() assert np.allclose(torch_out2.numpy(), ms_torch_out2.numpy(), equal_nan=True) def test_atan(): a1 = np.random.randn(4).astype(np.uint8) torch_a1 = torch.tensor(a1) torch_out1 = torch_a1.atan() ms_torch_a1 = pytorch.tensor(a1) ms_torch_out1 = ms_torch_a1.atan() assert np.allclose(torch_out1.numpy(), ms_torch_out1.numpy()) a2 = np.random.randn(3, 2).astype(np.float32) torch_a2 = torch.tensor(a2) torch_out2 = torch_a2.atan() ms_torch_a2 = pytorch.tensor(a2) ms_torch_out2 = ms_torch_a2.atan() assert np.allclose(torch_out2.numpy(), ms_torch_out2.numpy(), equal_nan=True) def test_atan2(): x = np.random.randn(2).astype(np.int32) y = np.random.randn(2, 2).astype(np.int32) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out1 = torch_x.atan2(torch_y) ms_torch_x = pytorch.tensor(x) ms_torch_y = pytorch.tensor(y) ms_torch_out1 = ms_torch_x.atan2(ms_torch_y) assert np.allclose(torch_out1.numpy(), ms_torch_out1.numpy()) def test_count_nonzero(): x = [[0, 1, 2, 0, 9],[0, 1, 0, 0, 0]] torch_x = torch.tensor(x) torch_out = torch_x.count_nonzero() ms_x = pytorch.tensor(x) ms_out = ms_x.count_nonzero() assert np.allclose(torch_out.numpy(), ms_out.numpy()) torch_out = torch_x.count_nonzero(0) ms_out = ms_x.count_nonzero(0) assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_scatter_1(): ms_out = pytorch.full((2, 4), 2.).scatter_(1, pytorch.tensor([[2], [3]]), 1.23, reduce='add') torch_out = torch.full((2, 4), 2.).scatter_(1, torch.tensor([[2], [3]]), 1.23, reduce='add') assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_scatter_2(): ms_out = pytorch.full((2, 4), 2.).scatter_(1, pytorch.tensor([[2], [3]]), 1.23) torch_out = torch.full((2, 4), 2.).scatter_(1, torch.tensor([[2], [3]]), 1.23) assert np.allclose(ms_out.numpy(), torch_out.numpy()) # TODO: To support index.shape != src.shape ''' def test_scatter_3(): src = torch.arange(1, 11).reshape((2, 5)) index = torch.tensor([[0, 1, 2, 0]]) torch_out = torch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src) src = pytorch.arange(1, 11).reshape((2, 5)) index = pytorch.tensor([[0, 1, 2, 0]]) ms_out = pytorch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_scatter_4(): src = torch.arange(1, 11).reshape((2, 5)) index = torch.tensor([[0, 1, 2], [0, 1, 4]]) torch_out = torch.zeros(3, 5, dtype=src.dtype).scatter_(1, index, src) src = pytorch.arange(1, 11).reshape((2, 5)) index = pytorch.tensor([[0, 1, 2], [0, 1, 4]]) ms_out = pytorch.zeros(3, 5, dtype=src.dtype).scatter_(1, index, src) assert np.allclose(ms_out.numpy(), torch_out.numpy()) ''' def test_acosh(): np_array = np.random.rand(1, 1, 1, 1, 1, 2, 3, 2).astype(np.float64)*50+2 torch_tensor = torch.tensor(np_array) torch_out = torch.acosh(torch_tensor) ms_tensor = pytorch.tensor(np_array) ms_out = pytorch.acosh(ms_tensor) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_new_ones(): tensor = torch.tensor((), dtype=torch.int32) torch_out = tensor.new_ones((2, 3)) tensor = pytorch.tensor((), dtype=pytorch.int32) ms_out = tensor.new_ones((2, 3)) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_new_empty(): tensor = torch.ones(()) torch_out = tensor.new_empty((2, 3)) tensor = pytorch.ones(()) ms_out = tensor.new_empty((2, 3)) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_addcmul(): np_array = np.array([1, 1, 1]).astype(np.float64) np_x = np.array([[1], [2], [3]]).astype(np.float64) np_y = np.array([[1, 2, 3]]).astype(np.float64) torch_tensor = torch.tensor(np_array) torch_x = torch.tensor(np_x) torch_y = torch.tensor(np_y) torch_out = torch_tensor.addcmul(torch_x, torch_y, value=2.1) ms_tensor = pytorch.tensor(np_array) ms_x = pytorch.tensor(np_x) ms_y = pytorch.tensor(np_y) ms_out = ms_tensor.addcmul(ms_x, ms_y, value=2.1) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arccosh(): np_array = np.random.rand(1, 4, 5, 6).astype(np.float64)*5 np_array[0, 0, 0, 0] = 1 torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.arccosh() ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.arccosh() assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arcsin(): np_array = np.random.rand(1, 2, 3, 4).astype(np.float64)*5 torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.arcsin() ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.arcsin() assert np.allclose(ms_out.asnumpy(), torch_out.numpy(), equal_nan=True) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arctan(): np_array = np.random.rand(1, 2, 3, 4).astype(np.float64)*5-2 torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.arctan() ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.arctan() assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_arctan2(): np_array = np.random.rand(1, 4, 5, 6).astype(np.float64)*5-2 np_other = np.random.rand(1, 4, 5, 6).astype(np.float64)*5-2 torch_tensor = torch.tensor(np_array) torch_other = torch.tensor(np_other) torch_out = torch_tensor.arctan2(torch_other) ms_tensor = pytorch.tensor(np_array) ms_other = pytorch.tensor(np_other) ms_out = ms_tensor.arctan2(ms_other) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_bitwise_not(): np_array1 = np.arange(-5, 12).astype(np.uint8) np_array2 = np.arange(-10, 10).astype(np.int16) np_array3 = np.array([[[False, True, True, False, False]]]) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_tensor3 = torch.tensor(np_array3) torch_out1 = torch_tensor1.bitwise_not() torch_out2 = torch_tensor2.bitwise_not() torch_out3 = torch_tensor3.bitwise_not() ms_tensor1 = pytorch.tensor(np_array1) ms_tensor2 = pytorch.tensor(np_array2) ms_tensor3 = pytorch.tensor(np_array3) ms_out1 = ms_tensor1.bitwise_not() ms_out2 = ms_tensor2.bitwise_not() ms_out3 = ms_tensor3.bitwise_not() assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_bitwise(): np_array1 = np.arange(-5, 11).reshape(1, 2, 2, 4).astype(np.bool8) np_array2 = np.arange(-10, 6).reshape(2, 2, 4).astype(np.bool8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch_tensor1.bitwise_and(torch_tensor2) torch_out2 = torch_tensor1.bitwise_or(torch_tensor2) torch_out3 = torch_tensor1.bitwise_xor(torch_tensor2) ms_tensor1 = pytorch.tensor(np_array1) ms_tensor2 = pytorch.tensor(np_array2) ms_out1 = ms_tensor1.bitwise_and(ms_tensor2) ms_out2 = ms_tensor1.bitwise_or(ms_tensor2) ms_out3 = ms_tensor1.bitwise_xor(ms_tensor2) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_addbmm(): M_ = np.random.randn(3, 5) batch1_ = np.random.randn(10, 3, 4) batch2_ = np.random.randn(10, 4, 5) M = torch.tensor(M_) batch1 = torch.tensor(batch1_) batch2 = torch.tensor(batch2_) torch_output = M.addbmm(batch1, batch2, alpha=2, beta=3) M = pytorch.tensor(M_) batch1 = pytorch.tensor(batch1_) batch2 = pytorch.tensor(batch2_) ms_output = M.addbmm(batch1, batch2, alpha=2, beta=3) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_addr(): vec1 = torch.arange(1., 4.) vec2 = torch.arange(1., 3.) M = torch.zeros(3, 2) torch_output = M.addr(vec1, vec2, alpha=2, beta=3) vec1 = pytorch.arange(1., 4.) vec2 = pytorch.arange(1., 3.) M = pytorch.zeros(3, 2) ms_output = M.addr(vec1, vec2, alpha=2, beta=3) assert np.allclose(ms_output.numpy(), torch_output.numpy()) def test_all(): a = pytorch.rand(1, 2).bool() a1 = a.all() a = pytorch.arange(0, 3) b1 = a.all() a = torch.rand(1, 2).bool() a2 = a.all() a = torch.arange(0, 3) b2 = a.all() assert np.allclose(a1.numpy(), a2.numpy()) assert np.allclose(b1.numpy(), b2.numpy()) def test_allclose(): a1 = pytorch.tensor([10000., 1e-07]).allclose(pytorch.tensor([10000.1, 1e-08])) b1 = pytorch.tensor([10000., 1e-08]).allclose(pytorch.tensor([10000.1, 1e-09])) c1 = pytorch.tensor([1.0, float('nan')]).allclose(pytorch.tensor([1.0, float('nan')])) d1 = pytorch.tensor([1.0, float('nan')]).allclose(pytorch.tensor([1.0, float('nan')]), equal_nan=True) a2 = torch.tensor([10000., 1e-07]).allclose(torch.tensor([10000.1, 1e-08])) b2 = torch.tensor([10000., 1e-08]).allclose(torch.tensor([10000.1, 1e-09])) c2 = torch.tensor([1.0, float('nan')]).allclose(torch.tensor([1.0, float('nan')])) d2 = torch.tensor([1.0, float('nan')]).allclose(torch.tensor([1.0, float('nan')]), equal_nan=True) assert a1 == a2 assert b1 == b2 assert c1 == c2 assert d1 == d2 def test_isclose(): a1 = pytorch.tensor((1., 2, 3)).isclose(pytorch.tensor((1 + 1e-10, 3, 4))) b1 = pytorch.tensor((float('inf'), 4)).isclose(pytorch.tensor((float('inf'), 6)), rtol=.5) a2 = torch.tensor((1., 2, 3)).isclose(torch.tensor((1 + 1e-10, 3, 4))) b2 = torch.tensor((float('inf'), 4)).isclose(torch.tensor((float('inf'), 6)), rtol=.5) assert np.allclose(a1.numpy(), a2.numpy()) assert np.allclose(b1.numpy(), b2.numpy()) def test_addmm(): _M = np.random.randn(2, 3) _mat1 = np.random.randn(2, 3) _mat2 = np.random.randn(3, 3) M = pytorch.tensor(_M) mat1 = pytorch.tensor(_mat1) mat2 = pytorch.tensor(_mat2) a1 = M.addmm(mat1, mat2, alpha=2, beta=3) M = torch.tensor(_M) mat1 = torch.tensor(_mat1) mat2 = torch.tensor(_mat2) a2 = M.addmm(mat1, mat2, alpha=2, beta=3) assert np.allclose(a1.numpy(), a2.numpy()) def test_cholesky(): _data1 = np.random.randn(3, 3) _data2 = np.random.randn(3, 2, 2) a = torch.tensor(_data1) a = a @ a.mT + 1e-3 torch_out1 = a.cholesky() a = torch.tensor(_data2) a = a @ a.mT + 1e-03 torch_out2 = a.cholesky() a = pytorch.tensor(_data1) a = a @ a.mT + 1e-3 ms_out1 = a.cholesky() a = pytorch.tensor(_data2) a = a @ a.mT + 1e-03 ms_out2 = a.cholesky() assert np.allclose(torch_out1.numpy(), ms_out1.numpy()) assert np.allclose(torch_out2.numpy(), ms_out2.numpy()) def test_aminmax(): t = torch.arange(0, 10).view(2, 5) c1, d1 = t.aminmax(dim=0, keepdim=True) t = pytorch.arange(0, 10).view(2, 5) c2, d2 = t.aminmax(dim=0, keepdim=True) assert np.allclose(c1.numpy(), c2.numpy()) assert np.allclose(d1.numpy(), d2.numpy()) def test_any(): data1 = np.random.randn(1, 2) data2 = np.random.randn(4, 2) a = pytorch.tensor(data1).bool() a1 = a.any() a = pytorch.arange(0, 3) b1 = a.any() a = pytorch.tensor(data2) < 0 c1 = a.any(1) d1 = a.any(0) a = torch.tensor(data1).bool() a2 = a.any() a = torch.arange(0, 3) b2 = a.any() a = torch.tensor(data2) < 0 c2 = a.any(1) d2 = a.any(0) assert np.allclose(a1.numpy(), a2.numpy()) assert np.allclose(b1.numpy(), b2.numpy()) assert np.allclose(c1.numpy(), c2.numpy()) assert np.allclose(d1.numpy(), d2.numpy()) def test_bincount1(): np_array = np.random.randint(0, 8, (7,)).astype(np.uint8) weight_array = np.linspace(0, 2, len(np_array)) torch_tensor = torch.tensor(np_array) torch_weight = torch.tensor(weight_array) torch_out = torch_tensor.bincount(torch_weight) ms_tensor = pytorch.tensor(np_array) ms_weight = pytorch.tensor(weight_array) ms_out = ms_tensor.bincount(ms_weight) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_bincount2(): np_array = np.random.randint(0, 4, (8,)).astype(np.int64) torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.bincount(minlength=5) ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.bincount(minlength=5) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_bitwise_shift(): np_array1 = np.arange(-5, 11).reshape(2, 2, 4).astype(np.int8) np_array2 = np.array([1,2,7,1,2,5,1,2,-1,-3,-4,0,2,1,0,0]).reshape(2, 2, 4).astype(np.int8) torch_tensor1 = torch.tensor(np_array1) torch_tensor2 = torch.tensor(np_array2) torch_out1 = torch_tensor1.bitwise_left_shift(torch_tensor2) torch_out2 = torch_tensor1.bitwise_right_shift(3) ms_tensor1 = pytorch.tensor(np_array1) ms_tensor2 = pytorch.tensor(np_array2) ms_out1 = ms_tensor1.bitwise_left_shift(ms_tensor2) ms_out2 = ms_tensor1.bitwise_right_shift(3) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_cholesky_inverse(): _data = np.random.randn(3, 3) a = torch.tensor(_data) a = torch.mm(a, a.t()) + 1e-05 * torch.eye(3) u = torch.cholesky(a) torch_out = u.cholesky_inverse() a = pytorch.tensor(_data) a = pytorch.mm(a, a.t()) + 1e-05 * pytorch.eye(3) u = pytorch.cholesky(a) ms_out = u.cholesky_inverse() assert np.allclose(torch_out.numpy(), ms_out.numpy()) def test_clip(): np_array = np.random.rand(1, 2, 3, 4).astype(np.float32) min = 0.35 max = 0.65 torch_tensor = torch.tensor(np_array) torch_out1 = torch_tensor.clip(torch.tensor(min), torch.tensor(max)) torch_out2 = torch_tensor.clip(0.85, max) torch_out3 = torch_tensor.clip(None, max) ms_tensor = pytorch.tensor(np_array) ms_out1 = ms_tensor.clip(pytorch.tensor(min), pytorch.tensor(max)) ms_out2 = ms_tensor.clip(0.85, max) ms_out3 = ms_tensor.clip(None, max) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype assert np.allclose(ms_out3.asnumpy(), torch_out3.numpy()) assert ms_out3.asnumpy().dtype == torch_out3.numpy().dtype def test_copysign(): np_other = np.array([[0.0, -0.0, 1, 1.5, -0.5, -20]]).astype(np.half) np_array = np.array([0, -0, 0, 1, 1, 1]).astype(np.int32) for type_input in (np.int32, np.half, np.float64, np.float32): torch_tensor = torch.tensor(np_array.astype(type_input)) ms_tensor = pytorch.tensor(np_array.astype(type_input)) for type_other in (np.half, np.float64, np.int8, np.int32): torch_other = torch.tensor(np_other.astype(type_other)) torch_out1 = torch_tensor.copysign(torch_other) ms_other = pytorch.tensor(np_other.astype(type_other)) ms_out1 = ms_tensor.copysign(ms_other) assert np.allclose(ms_out1.asnumpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype torch_out2 = torch_tensor.copysign(-0) ms_out2 = ms_tensor.copysign(-0) assert np.allclose(ms_out2.asnumpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_cos(): x = torch.tensor([0.24, 0.83, 0.31, 0.09], dtype=torch.float32) t_r = x.cos() y = pytorch.tensor([0.24, 0.83, 0.31, 0.09], pytorch.float32) ms_r = y.cos() assert np.allclose(t_r.numpy(), ms_r.numpy()) def test_cosh(): np_array = np.random.rand(1, 1, 1, 1, 2, 3, 2).astype(np.float64) - 0.5 np_array = np_array * 20 torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.cosh() ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.cosh() assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_cummax(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = pytorch.tensor(np_1) ms_result_1, ms_result_2 = ms_tensor.cummax(dim=0) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch_tensor.cummax(dim=0) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cummin(): np_1 = np.random.randn(3, 4).astype(np.float32) ms_tensor = pytorch.tensor(np_1) ms_result_1, ms_result_2 = ms_tensor.cummin(dim=0) torch_tensor = torch.tensor(np_1) torch_result_1, torch_result_2 = torch_tensor.cummin(dim=0) assert np.allclose(ms_result_1.asnumpy(), torch_result_1.numpy()) assert ms_result_1.asnumpy().dtype == torch_result_1.numpy().dtype assert np.allclose(ms_result_2.asnumpy(), torch_result_2.numpy()) assert ms_result_2.asnumpy().dtype == torch_result_2.numpy().dtype def test_cumprod(): np_1 = np.random.randn(3, 4).astype(np.float64) ms_tensor = pytorch.tensor(np_1) ms_result = ms_tensor.cumprod(dim=0) torch_tensor = pytorch.tensor(np_1) torch_result = torch_tensor.cumprod(dim=0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_deg2rad(): np_array = np.array([[180, -180, 360, 30, 0, 57, 80]]) torch_tensor = torch.tensor(np_array) torch_out = torch_tensor.deg2rad() ms_tensor = pytorch.tensor(np_array) ms_out = ms_tensor.deg2rad() assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype assert ms_out.asnumpy().shape == torch_out.numpy().shape def test_diag(): np_1 = np.array([1, 2, 3]) ms_tensor_1 = pytorch.tensor(np_1) ms_result = ms_tensor_1.diag() torch_tensor_1 = torch.tensor(np_1) torch_result = torch_tensor_1.diag() assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagonal(): np_1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) ms_tensor_1 = pytorch.tensor(np_1) ms_result = ms_tensor_1.diagonal(0) torch_tensor_1 = torch.tensor(np_1) torch_result = torch_tensor_1.diagonal(0) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_diagflat(): np_1 = np.random.randn(8).astype(np.float32) ms_tensor = pytorch.tensor(np_1) ms_result = ms_tensor.diagflat(offset=-1) torch_tensor = torch.tensor(np_1) torch_result = torch_tensor.diagflat(offset=-1) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) assert ms_result.asnumpy().dtype == torch_result.numpy().dtype def test_iscomplex(): a = pytorch.tensor([1+1j, 2, 3]) assert a.is_complex() == True a = pytorch.tensor([1, 2, 3]) assert a.is_complex() == False def test_isinf(): data = [1, float('inf'), 2, float('-inf'), float('nan')] a = torch.tensor(data).isinf() b = pytorch.tensor(data).isinf() assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype def test_isneginf(): data = [-float('inf'), float('inf'), 1.2] a = torch.tensor(data).isneginf() b = pytorch.tensor(data).isneginf() assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype def test_isposinf(): data = [-float('inf'), float('inf'), 1.2] a = torch.tensor(data).isposinf() b = pytorch.tensor(data).isposinf() assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype def test_isreal(): data = [1, 1+1j, 2+0j] a = torch.tensor(data).isreal() b = pytorch.tensor(data).isreal() assert np.allclose(a.numpy(), b.numpy()) assert a.numpy().dtype == b.numpy().dtype if __name__ == '__main__': test_others_tensor() test_add_() test_mul() test_mul_() test_fill_() test_size_with_dim() test_size() test_zero_() test_normal_() # test_uniform_() test_clamp_() test_permute1() test_permute2() test_copy_() test_expand() test_sigmoid() test_sigmoid_() test_float() test_flip() test_sign() test_div() test_div2() test_max() test_numel() test_sum() test_split() test_numpy() test_ndimension() test_pow() test_repeat() test_repeat_interleave() test_reshape() test_reshape_as() test_arcsinh() test_arctanh() test_det() test_negative() test_abs() test_ndim() test_amax() test_amin() test_as_strided() test_bmm() test_clamp() test_dim() test_expand_as() test_item() test_log() test_log2() test_matmul() test_squeeze() test_stride() test_sub() test_is_floating_point() test_unbind() test_unsqueeze() # test_is_signed() test_transpose1() test_transpose2() test_floor() test_isfinite() test_isnan() test_expand() test_clone() test_detach() test_new_zeros() test_to() test_sort() test_msort() test_argsort() test_sqrt() test_rsqrt1() test_rsqrt2() test_rsqrt3() test_resize() test_resize_as() test_index_fill() test_data() test_new() test_cuda() test_t1() test_t2() test_mean1() test_mean2() test_prod1() test_index_select() test_std() test_exp() test_masked_fill() test_tolist() test_chunk() test_cumsum() test_absolute() test_acos() test_addcdiv() test_asinh() test_atanh() test_gather() test_fmod() test_argmin() test_argmax() test_baddbmm() test_topk() test_slice() test_maximum() test_minimum() test_multiply() test_neg() test_ravel() test_select() test_square() test_broadcast_to() test_divide() test_mm() test_view() test_logsumexp() test_addmv() test_dot() test_inverse() test_asin() test_atan() test_atan2() test_count_nonzero() test_scatter_1() test_scatter_2() test_addbmm() test_addr() test_cholesky() test_aminmax() test_any() test_bincount1() test_bincount2() test_bitwise_shift() test_cholesky_inverse() test_copysign() test_cos() test_cosh() test_cummax() test_cummin() test_cumprod() test_deg2rad() test_diag() test_diagflat() test_diagonal() test_isinf() test_isneginf() test_isposinf() test_isreal()msadapter/testing/ut/pytorch/tensor/test_tensor_compare.py000066400000000000000000000147021436713554100246130ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) #context.set_context(mode=ms.GRAPH_MODE) def test_ge(): x = [[1, 2], [3, 4]] y = [[1, 1], [4, 4]] torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.ge(torch_y) ms_x = ms_torch.tensor(x) ms_y = ms_torch.tensor(y) ms_out = ms_x.ge(ms_y) assert np.allclose(ms_out.numpy(), torch_out.numpy()) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_gt(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[4, 2]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_tensor_1.gt(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch_tensor_1.gt(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_le(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[4, 2]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_tensor_1.le(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch_tensor_1.le(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_lt(): np_1 = np.array([[1, 2],[3, 4]]).astype(np.int64) np_2 = np.array([[4, 2]]).astype(np.int64) ms_tensor_1 = ms_torch.tensor(np_1) ms_tensor_2 = ms_torch.tensor(np_2) ms_out = ms_tensor_1.lt(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1) torch_tensor_2 = torch.tensor(np_2) torch_out = torch_tensor_1.lt(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_eq(): x = np.random.randint(4, size=(2, 3, 4)) other = np.random.randint(4, size=(2, 3, 4)) torch_x = torch.tensor(x) torch_other = torch.tensor(other) ms_x = ms_torch.tensor(x) ms_other = ms_torch.tensor(other) torch_out1 = torch_x.eq(torch_other) ms_out1 = ms_x.eq(ms_other) torch_out2 = torch_x.eq(3) ms_out2 = ms_x.eq(3) assert np.allclose(ms_out1.numpy(), torch_out1.numpy()) assert ms_out1.asnumpy().dtype == torch_out1.numpy().dtype assert np.allclose(ms_out2.numpy(), torch_out2.numpy()) assert ms_out2.asnumpy().dtype == torch_out2.numpy().dtype def test_ne(): x = np.random.randn(2, 3, 4) y = np.random.randn(3, 4) ms_x = ms_torch.tensor(x) ms_y = ms_torch.tensor(y) ms_out = ms_x.ne(ms_y) torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.ne(torch_y) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_equal(): torch_x = torch.tensor([1, 2]) torch_out = torch_x.equal(torch.tensor([1, 2])) ms_x = ms_torch.tensor([1, 2]) ms_out = ms_x.equal(ms_torch.tensor([1, 2])) assert ms_out == torch_out assert type(ms_out) == type(torch_out) torch_x = torch.tensor([1., 1.]) torch_out = torch_x.equal(torch.tensor([1.])) ms_x = ms_torch.tensor([1., 1.]) ms_out = ms_x.equal(ms_torch.tensor([1.])) assert ms_out == torch_out assert type(ms_out) == type(torch_out) def test_greater_equal(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.float16) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.float16) ms_out = ms_tensor_1.ge(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.float16) torch_tensor_2 = torch.tensor(np_2, dtype=torch.float16) torch_out = torch_tensor_1.ge(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_greater(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.float64) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.float64) ms_out = ms_tensor_1.ge(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.float64) torch_tensor_2 = torch.tensor(np_2, dtype=torch.float64) torch_out = torch_tensor_1.ge(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_less_equal(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.double) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.double) ms_out = ms_tensor_1.less_equal(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.double) torch_tensor_2 = torch.tensor(np_2, dtype=torch.double) torch_out = torch_tensor_1.less_equal(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_less(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.long) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.long) ms_out = ms_tensor_1.ge(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.long) torch_tensor_2 = torch.tensor(np_2, dtype=torch.long) torch_out = torch_tensor_1.ge(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype def test_not_equal(): np_1 = np.array([1, 2, 3]) np_2 = np.array([1, 1, 4]) ms_tensor_1 = ms_torch.tensor(np_1, dtype=ms_torch.int32) ms_tensor_2 = ms_torch.tensor(np_2, dtype=ms_torch.int32) ms_out = ms_tensor_1.ge(ms_tensor_2) torch_tensor_1 = torch.tensor(np_1, dtype=torch.int32) torch_tensor_2 = torch.tensor(np_2, dtype=torch.int32) torch_out = torch_tensor_1.ge(torch_tensor_2) assert np.allclose(ms_out.asnumpy(), torch_out.numpy()) assert ms_out.asnumpy().dtype == torch_out.numpy().dtype if __name__ == '__main__': test_ge() test_gt() test_le() test_lt() test_eq() test_equal() test_greater_equal() test_greater() test_less_equal() test_not_equal() test_less() test_ne() msadapter/testing/ut/pytorch/tensor/test_type.py000066400000000000000000000036231436713554100225540ustar00rootroot00000000000000import numpy as np import mindspore as ms import torch from mindspore import context import ms_adapter.pytorch as ms_torch context.set_context(mode=ms.PYNATIVE_MODE) def test_tensor_type_1(): x = [[1, 2], [3, 4]] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch_x.type() ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_x.type() assert (ms_out in torch_out) def test_tensor_type_2(): x = [[1, 2], [3, 4]] torch_x = torch.tensor(x, dtype=torch.double) torch_out = torch_x.type(torch.int) ms_x = ms_torch.tensor(x, dtype=ms_torch.double) ms_out = ms_x.type(ms_torch.int) assert ms_out.numpy().dtype == torch_out.numpy().dtype def test_tensor_type_3(): x = [[1, 2], [3, 4]] torch_x = torch.tensor(x, dtype=torch.int64) torch_out = torch_x.type(torch.int64) ms_x = ms_torch.tensor(x, dtype=ms_torch.int64) ms_out = ms_x.type(ms_torch.int64) assert ms_out.numpy().dtype == torch_out.numpy().dtype assert id(ms_x) == id(ms_out) def test_tensor_type_4(): x = [[1, 2], [3, 4]] y = [1., 2.] torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out1 = torch_y.type() torch_out2 = torch_x.type(torch_out1) ms_x = ms_torch.tensor(x) ms_y = ms_torch.tensor(y) ms_out1 = ms_y.type() ms_out2 = ms_x.type(ms_out1) assert (ms_out1 in torch_out1) assert torch_out2.numpy().dtype == ms_out2.numpy().dtype def test_tensor_type_as(): x = [[1, 2], [3, 4]] y = [1., 2.] torch_x = torch.tensor(x) torch_y = torch.tensor(y) torch_out = torch_x.type_as(torch_y) ms_x = ms_torch.tensor(x) ms_y = ms_torch.tensor(y) ms_out = ms_x.type_as(ms_y) assert torch_out.numpy().dtype == ms_out.numpy().dtype if __name__ == '__main__': test_tensor_type_1() test_tensor_type_2() test_tensor_type_3() test_tensor_type_4() test_tensor_type_as() msadapter/testing/ut/pytorch/torch/000077500000000000000000000000001436713554100177635ustar00rootroot00000000000000msadapter/testing/ut/pytorch/torch/test_assert.py000066400000000000000000000034601436713554100227000ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_torch_assert1(): try: ms_torch._assert(1 == 2, 'assert fail') except: return 0 return 1 def test_torch_assert2(): try: ms_torch._assert(1 == 1, 'assert fail') except: return 1 return 0 def test_is_tensor(): x = [1, 2, -1, 2, 0, -3.5] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.is_tensor(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.is_tensor(ms_x) assert np.allclose(ms_out, torch_out) torch_out = torch.is_tensor(x) ms_out = ms_torch.is_tensor(x) assert np.allclose(ms_out, torch_out) def test_is_floating_point(): x = [1, 2, -1, 2, 0, -3.5] torch_x = torch.tensor(x, dtype=torch.float32) torch_out = torch.is_floating_point(torch_x) ms_x = ms_torch.tensor(x, dtype=ms_torch.float32) ms_out = ms_torch.is_floating_point(ms_x) assert np.allclose(ms_out, torch_out) def test_isfinite(): x = [1, float('inf'), 2, float('-inf'), float('nan')] torch_x = torch.tensor(x) torch_out = torch.isfinite(torch_x) ms_x = ms_torch.tensor(x) ms_out = ms_torch.isfinite(ms_x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) def test_isnan(): x = [1, float('nan'), 2] torch_x = torch.tensor(x) torch_out = torch.isnan(torch_x) ms_x = ms_torch.tensor(x) ms_out = ms_torch.isnan(ms_x) assert np.allclose(ms_out.numpy(), torch_out.numpy()) if __name__ == '__main__': test_torch_assert1() test_torch_assert2() test_is_tensor() test_is_floating_point() test_isfinite() test_isnan() msadapter/testing/ut/pytorch/torch/test_info.py000066400000000000000000000034771436713554100223420ustar00rootroot00000000000000import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import math from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def assert_finfo_attr(finfo, ms_finfo): assert finfo.bits == ms_finfo.bits assert abs(finfo.eps - ms_finfo.eps) <= 1e-6 assert abs(finfo.max - ms_finfo.max) <= 1e-6 assert abs(finfo.min - ms_finfo.min) <= 1e-6 assert abs(finfo.tiny - ms_finfo.tiny) <= 1e-6 # TODO: numpy vision >= 1.23 # assert abs(finfo.smallest_normal - ms_finfo.smallest_normal) <= 1e-6 assert abs(finfo.resolution - ms_finfo.resolution) <= 1e-6 def assert_iinfo_attr(iinfo, ms_iinfo): assert iinfo.bits == ms_iinfo.bits assert iinfo.min == ms_iinfo.min assert iinfo.max == ms_iinfo.max def test_finfo(): finfo = torch.finfo(torch.float16) ms_finfo = ms_torch.finfo(ms_torch.float16) assert_finfo_attr(finfo, ms_finfo) finfo = torch.finfo(torch.float32) ms_finfo = ms_torch.finfo(ms_torch.float32) assert_finfo_attr(finfo, ms_finfo) finfo = torch.finfo(torch.float64) ms_finfo = ms_torch.finfo(ms_torch.float64) assert_finfo_attr(finfo, ms_finfo) def test_iinfo(): iinfo = torch.iinfo(torch.uint8) ms_iinfo = ms_torch.iinfo(ms_torch.uint8) assert_iinfo_attr(iinfo, ms_iinfo) iinfo = torch.iinfo(torch.int8) ms_iinfo = ms_torch.iinfo(ms_torch.int8) assert_iinfo_attr(iinfo, ms_iinfo) iinfo = torch.iinfo(torch.int16) ms_iinfo = ms_torch.iinfo(ms_torch.int16) assert_iinfo_attr(iinfo, ms_iinfo) iinfo = torch.iinfo(torch.int32) ms_iinfo = ms_torch.iinfo(ms_torch.int32) assert_iinfo_attr(iinfo, ms_iinfo) iinfo = torch.iinfo(torch.int64) ms_iinfo = ms_torch.iinfo(ms_torch.int64) assert_iinfo_attr(iinfo, ms_iinfo) if __name__ == '__main__': test_finfo() test_iinfo()msadapter/testing/ut/pytorch/torch/test_multinomial.py000066400000000000000000000014121436713554100237240ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as torch from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) ms.set_seed(123) def test_multinomial1(): x = torch.tensor([0, 9, 4, 0]) out = torch.multinomial(x, 2) assert len(out) == 2 def test_multinomial2(): x = torch.tensor([0, 9, 4, 0]) out = torch.multinomial(x, 4, replacement=True) assert len(out) == 4 def test_multinomial3(): x = torch.tensor([0, 9, 4, 0]) output = torch.tensor([1, 1, 1, 1]) out = torch.multinomial(x, 4, replacement=True, out=output) assert len(out) == 4 assert id(output) == id(out) if __name__ == '__main__': test_multinomial1() test_multinomial2() test_multinomial3()msadapter/testing/ut/pytorch/torch/test_ones.py000066400000000000000000000011731436713554100223420ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_ones1(): ms_result = ms_torch.ones(1, 2, 3, dtype=ms_torch.float32) torch_result = torch.ones(1, 2, 3, dtype=torch.float32) assert np.allclose(ms_result.asnumpy(), torch_result.numpy()) def test_ones2(): ms_result = ms_torch.ones((1, 2, 3), dtype=ms_torch.float32) torch_result = torch.ones((1, 2, 3), dtype=torch.float32) assert np.allclose(ms_result.asnumpy(), torch_result.numpy())msadapter/testing/ut/pytorch/torch/test_randint.py000066400000000000000000000010761436713554100230370ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as torch from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) ms.set_seed(123) def test_randint1(): out = torch.randint(0, 10, (2, 2)) # TODO:out = torch.randint(10, (2, 2)) assert out.shape == (2, 2) def test_randint2(): output = torch.tensor([1, 1, 1, 1]) out = torch.randint(3, 10, (3, ), out=output) assert out.shape == (3, ) assert id(output) == id(out) if __name__ == '__main__': test_randint1() test_randint2()msadapter/testing/ut/pytorch/torch/test_randn.py000066400000000000000000000011241436713554100224740ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as ms_torch import numpy as np from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) def test_randn1(): try: ms_torch.randn(1, 2, 3) except: return 1 return 0 def test_randn2(): try: ms_torch.randn([1, 2, 3]) except: return 1 return 0 def test_randn3(): try: ms_torch.randn((1, 2, 3)) except: return 1 return 0 if __name__ == '__main__': test_randn1() test_randn2() test_randn3()msadapter/testing/ut/pytorch/torch/test_randperm.py000066400000000000000000000012241436713554100232030ustar00rootroot00000000000000#!/usr/bin/env python # -*- coding: utf-8 -*- import mindspore as ms import ms_adapter.pytorch as torch from mindspore import context context.set_context(mode=ms.PYNATIVE_MODE) ms.set_seed(123) def test_randperm1(): out = torch.randperm(4) assert len(out) == 4 def test_randperm2(): out = torch.randperm(4, dtype=torch.float32) assert len(out) == 4 assert out.dtype == torch.float32 def test_randperm3(): output = torch.tensor([1, 1, 1, 1]) out = torch.randperm(4, out=output) assert len(out) == 4 assert id(output) == id(out) if __name__ == '__main__': test_randperm1() test_randperm2() test_randperm3()msadapter/testing/ut/runtest.sh000066400000000000000000000006321436713554100172150ustar00rootroot00000000000000set -e if ! [ -x "$(command -v pytest)" ]; then pip install pytest fi TEST_DIR=$(realpath "$(dirname "$0")") PROJECT_DIR=$(realpath "$TEST_DIR/../../") if [ $# -eq 0 ] || [ $1 == "pytorch" ]; then echo "********** Run MSAdapter-pytorch ut." cd "$PROJECT_DIR" || exit UT_PATH="$TEST_DIR/pytorch/" pytest "$UT_PATH" echo "********** Test all MSAdapter-pytorch unit testcases success." fimsadapter/testing/ut/torchvision/000077500000000000000000000000001436713554100175235ustar00rootroot00000000000000msadapter/testing/ut/torchvision/assets/000077500000000000000000000000001436713554100210255ustar00rootroot00000000000000msadapter/testing/ut/torchvision/assets/encode_jpeg/000077500000000000000000000000001436713554100232675ustar00rootroot00000000000000msadapter/testing/ut/torchvision/assets/encode_jpeg/grace_hopper_517x606.jpg000066400000000000000000002200221436713554100274450ustar00rootroot00000000000000JFIF``^Photoshop 3.08BIM Resolution``8BIM FX Global Lighting Angle8BIMFX Global Altitude8BIM Print Flags 8BIM Copyright Flag8BIM'Japanese Print Flags 8BIMColor Halftone SettingsH/fflff/ff2Z5-8BIMColor Transfer Settingsp8BIMGuides@@8BIM URL overrides8BIMSlicesy^ Grace Hopper^8BIMICC Untagged Flag8BIMLayer ID Generator Base8BIM New Windows Thumbnail`p ~JFIFHHAdobed            p`"?   3!1AQa"q2B#$Rb34rC%Scs5&DTdE£t6UeuF'Vfv7GWgw5!1AQaq"2B#R3$brCScs4%&5DTdEU6teuFVfv'7GWgw ?=_+]?;'73O"v4Ӽkl. MyS/9,Ʋ인n}ݏUch}w{QobV:W^QǶưWӳ &Vt߉.V^NbYdc: ;kÛEt^?Yf鰾AγmFW1>ιn}97TVRm=ƪ<{}U$k=#(&d90D~cS)iWڃm5b}ou߷[JX}OdXg7zYeW=mk-o&nsmӇM~hk*mv9Nk~SWԺ<_PTl={'|)%euU1%"opE$d)_7]{Ko{=7=k?Ҧi-0.OO= ۝?V3HZFEV8sx:>.ߡo5v6NcTs>?G߷u1́=ӺFVmy| koNzWWM::wIvqǹ ]ݹQ`ތU{Gګf7)'+ᨐ9_v5sY[ D-{V~dmu,@6{Hճ3K3du",`UGSzNic0qh>ʴ^*ZdYg}5oO̳lm[ZOmZY]z^?c 2Ŏuau&Zc6z'}gՏ l4N{C/P4WUUcmi&'X.<x՞쬋2,׺^OawjEzjF}eլWk`sC\}r4lrrZmz9|HGzƵf_zlޱk1_eg}Qf9e-ϯXpzό`*}k]wݑ5ac?NZ׾Yc8Cvi.k+1 F"ҢKnkmNuz`6z^N۽hWnM3`.{;ֻ)׻ߓ1ub%r΋hXT zZ=VmH*٨^r72HҢ??Uǯ׺3+8Q]US^\k>1MK1ht~p^4zX+ǏL ȣk6J]Uo=m{qH3[Yxn{Ռ q>HmU_a%!&ZTXtc"[Y˦=%ՇFH1ߧ;]O `3wYNE95Y}mqTeno/غ5+~Mhw6l{}Tԛ1[U͹4<7 w4sGbG#3{v,OjخBm,S]svm|LdnkFq" 3F1P_+jO˪xe{]vk۽C9Ä_$L"e.!Dow^[r0k2I':Ϣdml.=RcLuֿF`}_R'+Y.cc]cm?;?駈H  kYb!) hã-|wdYz`to{N20Yk7j.1a` ߪO=fy%)nO (!*$jXze0qr5[UW0 1w>=;}Ḽͤ{i CGeuk=O[ٱeg_Z)ճŻ f?DX Wo+W3u]fEuu{!02zPY %;sk@O7\6Y)gol@cHn̑zw;3~ݕe6Dz|[@7!ZXRm\>7Snmy}+ÿ~ 9MŻ'oV]hgѸSIk_l<̪Q[4aq.bac1ZC%_q~[?]koIkYo*:9}o u&IoockѶ.{;g~4zSK˱SN8rwf9ڬw%[ 3@B/R[mW6^aXZ{Gu !eSdzvnZ|T* ;i1tL_V[_ZAk\ǀ<kpKn}P/7pk4~f}~V1 < bvԯ1KGG;}/yܮd[_m걍ѫĤ_Epcmss\;v3˛‘$U- ِ|uG.+˺VEX{=^^wF:xxA.~}c>Z%h~mM%< עq21`6?SvY[3}tPkmW:%O}Mhvƹ׵.!{^\9I$˜.`}e[쯵VS18{}/k&5&!4AXn>ipp~uzrrik݉f犝ncݺMgjwPb96^xVikm7gJ^ߣ/>/Mvy=7"YKc{U11+ͽ[ak7{U :$ ~p SoM_R/V˲@&C%ρU8MGɛldc*#ᢛnmsP ST4"UcxC#t!r]ߜRhӧG7ݳ)ęl>)A;O(̊=c>kDNH64!w.oh(EV'ݞϣ9?X:Kmq6\mc5?,E #(Z/rq ˌUK#;.WYnkZ^G*.qrE se9))'8BIM!Version compatibility infoUAdobe PhotoshopAdobe Photoshop 6.08BIM JPEG QualityAdobed              ^"!!1AQ"aq2BRbr#т3CSc$s4D%&Td!1AQaq"2BRrb#3c4႒$Ss%C5 ?hQ((\S D*Q?cF*e c7Hw P4ڥIJ" ̝J]Trm-/̄44RBR^8 O+GVzD\&ܺT sD#̘. *G1\,Ϙ*IRQIAǐ$!qW(tVef'obcY'ˊAc wR9H0 4S{HWA/,}p(D!9cBo4o/i;Ԯu{Mi~uÈkq9(QS>:\eτᤖB3WjʷH붭ܝ4'aIq)<d#qTYQDGKjzj1}S)Ļ׼=n]XSv1gE1"XmSWIUs,n{Xsa{DpS gS!@,25^ H!i-9!Uk)4se4.Ωޚ=;? i^E j@ۊʇ9'._#=/>Uz٧zW4s1ow]k^V}E` }J6_ouLTqi rDW4ҰQ<֥JU!&'ڋM_k-.eUB,vMNiR܀8RA7S- (iU(~W.rh\kN/&Vwj @U$$; H퇂]߈\5dJF e9$Tj )D҈B-$L'#(`@tT8b<LIC1 0:T> Zj!CQ-) " &'T!9 `ּ>bIY]Z,? @CIӥݚ.JÚje#2 IgPL{.C,׸CGjB#xa -AJ Հ2;Lyk3 !EBK#su#%2#b @c'Q# 5* B= g Q8C2QPՈ$8BpcHo.T'OtLBfG Wg)##lH|*d dAJBpMN gaT@G`\4Y/u** DI=@ DIa+!@` 8]9@ 3h yd%))HwNGñ쓊jq\jp (4U{*Kq(qDa )??¨Skx ,vgEB9D)1;C]bB| > 텸iz,i'N%S8oQ>P?aI)&@NKut@ry' 9:ײHW0_ycDTJIw*Ď#q ̀ ߈f57U)aZ`1 kf(]H J܂+,r1^Vi5$#&a% %QVM+|JA6%!>lJ8cÌd+K!<>`@!U /u"@BU!TX/őBqQx(qT MJ@9$I$j+ObBgfYbpB@EDU8)>c O,1S8 ZL*<" s%ǤzMe+ e:2?MWՕ}QN-HeL6F03cqwA>mv柅U\92O6Z⿻=/n=^jdc -ҡuS#1ݷnje W9.72XUМ$KS1jnhĵ{uDJ?;Yhm aYJ _v- dVeQ!Y P $H$Ghz,Sw! k\Hc]Zq5ʄ@dq? U:] l{t7Lwװ|ٌei=r8R<uj83FCIC؆;Bm=y]UGsoYߨ&E`=hOVӭuGq{\S$XuZl5qoCr)9Rãkt52S雪PWT3 v<\U$8Rs1í5YmmYM=A*vRw<ϼ]nW{M:mPWӭ|lm['UzQ.hQ b:խ3Uzh9f*$Q]Tj 9ܪA/f׹n ^܁V 1h/w/Q,xux[i_BSBG5Ti.ģIƉ :U;v}RQ VJF_k+\mwul/h:-j.vs@PriF>WGF9p쟉mpp N3. ,JE*.vO0z,uRL?+Û(2~jOqd7Š WxAw18&#v-n]mPTj+όJKfH*$<#\ZjfiѪ> TjjR yw(I!o4 Oj$AA Ctتr,LN -',W:CԘv d1"IA0d駔 D Y- 3# 48} 샊Ȯ/l1'$ R"H#Brjq43_PJ& J}HD4.D"$]ް j! 椒L; ~,24B(J eqE!.(2$0L`:Jf*uxwqRSD g@!}HdF8x $T$9{bC4j˜j Fgl\&=K!\152Rg`bBLpW/LVa`+"OtDH疑M:Ap% ZLa !BW"pL(504)MBjO1qr,pdWM3BqLc#fe?tuo|IvYWJ0t@av <{r]S*qS:2$`F}14jY1Vr ÄfSq+5ȯD$RDZjsR8!*FAd1K HDM3)A&@b!aˊDT&Ȳ;0H{Bpp^ɪ$5Xj^$㗌Usz8N8ot.QAe>;I~ҋm~aIԙb@̘AL",pX5@&M(pU 7J/oc"Ũq:$=1DVIR>1"H3!8$vΈtQj :8B~z>?J<z8w\ W`8!cnd}qtx\PC@(y.9q؅O>ʎ߬= P @B6e1GjLRPX.Si(T{U˯ɹ %c:o1Ƃ~(8IT8.y/JiI.{}XNӪa cW}aro6z:njWI}j'ƨzRۙ^ōc]]8 4r*=oʿXڶ[}\Thm'[}M? 8K}tȾ)GoMإ t'~Ws}yc[ߺb,.]Aim4=ͣR9i?CuQ1FA|Ꝿ^g7QTK}:o-3,v:f+nTkw:SiWbkgW/$FqcJCesMpoYٺ޵ mG Zh`5TyKOk *2sm鲱 41VsGziio4\Z\Ӵ+[]W`ʬ(AkKfjk7զM_>ѵn7oamQ+>AMik1ө56^ыTߒkpyzVt=o'~KSjn3^h}K ~594WCʟcՋWP-*;-ՕphԧhڏGO'Tz&*R-]?>"%IBtFyݕNF\4ΨC#CڸߴXhuMMJcM#6,}QW_2߳bISg3Hv=_RW;-*5_[n‹noVCj5BQOo;OaM . SMҊ.m]K[n1gN+QAİh}3ritWͨ56>jإԿuk.<⺔bZ4Fb F)9wcV{}_^R}a?FtR|MTq85)(iFDPdTE2&~ \9dI ph*J 8TM$a($D@^,)y `3n q9gQ$W8y,r(T!2 p2nHdrRyHT*sdh:Tt K # "L2A^?05 iiD01@H2B$< He @Be]JLK(D3_b$0!91)d"ド:B<`)tUfRFD%%*_ht1I8q&dCPa>)a(:b@O U(9 LrrqBƣFC^BȆ gaQRDd2E2=#B \/J)d ԛP.-'NI0E U?CI9BqD?f^i@Uz>(2ꌝuezd:wUgA/}zdмk41SoA2p(1vGm˩:qu56^ ^ۥ]#RKkK 1o|%T%%QbiqSnbzm7qoBե+|qQ w3PAL.qִ gHK湨J6eqNɞG"S wE J)_ya ]),p33R;'AKM-CJjȿA6>?J_%p_7@IYFs@`*}ͭi_^dSQGvìcWyQZXT1,{cs\dƵ'T ;qс,՚\u&)q?Im5/5Se)p qid@4T:(,)>VӶo՝Mnz5QE+.kKԙG#閞ԕy2_إS*n!'NbmAWW#XYmvvzvT(1jFPPa^i.97)*ס'B8iP[QRUԭmJs^!zMss2=>ۖim=EiOMZ4965 rgɜ*wѦ5j%1:k+jśmk_QڽjsVM76 QΘNP8<"D;-@y EBŗ\J1n% &tʹ{7m޶[ .N? TUSV../>VFަcsc3ҥsMz$TI>]1O7%\[RիH6֦zզYQKԭϧiEtOm-6S*8sXM#t7o6cUj'bS{gMnTjHݥS' odcצiSsԢ]ƸxT8ͽ)hdG)M0$v -@8=)LiMFkb 4z /<8GcoQvga\q̅2g(-lokFsGvt=+u[5oWרtD'ю da"T}ѕ 7HwC lHb;lm}aq?zi-ˆkVQX:rZAjƚ ) ^Fo!,~|.Zs N'|Ҷ]InT[pVծ²4?6~>F;Fd)XPj#o$M\Tj:zkwn46< 8yi1̩Xp`pFZGwYLzQ[6n7'iJEM:Z* %1Y]iBG;.[ruC*ӨuT[O~syթyڕEGSP׸w9ė=LW*]s;SnwQd/ߖi5zͅwYvˍ{MVZ vLJlӧ_qw={K2hK[uGg+3l:;jޙYU[ѫQZ׾J54uiUf3зbÄa kk3˵iOU P* eݵQn-U3gK\~7CA$cdp )P9@cCQS=1=O wR,ŵ=*z@3 ~HҶݮj:具ac\p'Nkj5J4!tCcf}gMm(fѧMB/ln-tn/A+YT9㫮RmJW$/qƍ*1Ywۮ5ksQwIZn<Nm.n۷;;sRe>`)s厢.3Ӥ|+nYI%oR&z|D:X1Zܚ[h6֝-/k|Mϣ]Zp{^0ʃMS+) uZTw?\lZQo :XA1-CM={0^v`?NK+Nu< ?FڝiU08UMybZj(HC-' %J0⊑w%@}Qk][Ώm+K:wׅ+*H~Y)qtthwtKh݀Ihھ2_fO=yn6uҁpƴJM5q: n5^+EzҎ mri$:EnkmwU=6]4P4yVLm`}CYyl.6Z/66g#=k[괮9 ;7 4Yv _H}g|M'eG.;{&WcR8a[=@6-'rUoާuj$ F* 6*.m:TQr[=HHEpc>sVmf2ىipͨI5']/*oDdۚ#8%S(j@Ǻ$I;pHQtLEdH@dD0%W,i.0N$A*&&eyżD+PeE;Ch { @*{xCH$\H!30I W;0DKq|! 0%؃+.*LiĪL! NV_ $@ JH&DG d>0#f\s)AI0B%I bA0̤i3 <Bf;s; E L=Ah D@ )SBFyBN; GfXPU4b%^PQ-I)4hVv4GH*G]Gһ7~ƜJX8\ZeNdnoL:{G4r7 6CUJKnNէYpwl=ջ hS\Vmv|)d/;u)6_8/- ݳVe@=1Q`;9UyZګWy{ܸO(1~V85:t 3zeR~GI~mS ţu#[OiғYY؄+=L>ӬX{pcSHȀP/nf\?e]J^M"c$c'hi[ލFƮN*\S0\J<=8=ɧ+x4:ǣl)6+Ӣt@ .Qu'ROUyr P4=:M ccLS ;ျWò)FX|N{w1uL^=&VZ;;޿IYޕMg鬪2ZVj7W؎3({'(FepCi^?U~͆8֨UU‡w.ڴVn8,.*6}NK[zoUNe٪+(v]-jA۽kpp֓inu(۲)z՞NH~m3{vREݸ\V!u\@/[˫mr62K~6Tے4ќ3I*`Σ,:ctmgVNJ|˛4Cy=Douv-u1^Ew䵾gu'T.wqӶmTh\tPUW\#yeW`H:}ƻPV.}G5ZiEw-ߋmE*Z ӧ:vi@u/Ujֱhzco/mw->շH5k6:Yy],˦mFClY|h7]vʔ6kjUrLÃojvkWIiR;WqݮAi\C)ɱ)mBMZ.5>`X^3ph5 M1W6p7wpmhZVcZm_IطһkzgԖomJ/'צ;Jl4W=A{m49)7AǭN3s?tljimEu.'Qݟ=v`6~6 Dϙ޲5і74q&JE\]|Rc`ڎm6ZC}G7^nwk?On5 F-iG1Ե8jQKgD[KӢʎm GӨ:tLlQ?tzꙥGhCTD;{Ҭ7K>/;ڶ-H9]hsjjvV39Pꯘ׫w[kxZJu<I;LJ*I)WHEOڲ1lw-Li}n5M-Zy-cjuQqDgne^wm+_NqS y\Z71菛m(7kToi~)\3N<ܕXشlyc}=wDs(ѿj65j LHn[VuVyz`\_0$!d+6oj]z.rn[mPHA3COt4V\8Ʒ[m7pujGҪ&oFIY4J>FjqRY5S(!92BHTĐ4O:P(b(+Q8aA$Des+-YBH 32 (@0+HB&I>DjB{f PTVRdBsШ%Tp )J@.rb, 8"c5A$#ҠLChID,̤ԉJcW AQH) UV 2pA"UV)(S+2e$Ye2[ieD fwt"HSy Cȟ Lʎl( ~n-hϏ ; "{tT R+]S>|sEQ["u=Ōcu/B Ln*?AՔ _{\ޕII3=<"fj- L3aܐ娠&cTI{b~ibm*MR6]B!ОC-.n1>R94ZDR$sz`kOHH;~Jk?U c|׫| #Al:u : DŽVzd꽆mڍRF-8v7)U'r<࠸*N]QKZOpoѸV֠M:AmR(IĂLv|"@'*54̢/,c˩P.\r9 1KsʧbꝖ[J{JU,~~Uheo?GEߺC:H_ϵvQ~^TkFQ* Mhr]$!e9EJMB5u{/h*3@$g%\cu_Y/wCuBHY#{1Xs\D5Le@TZYfI=33Q@|ǻL '!\OSl **g~PDI$IL9@!C$xWQKG0/d&^܁〆!gdqCdU!q>@i*`8s wLEJw Jv9PJj@c\PG@9l; 8*9TᒙB33vIpRQ_]1 ;$~o(CPa't6 PvA51>Hqu)HDMg!&#' Ġ.ab5WD \bF[#cꕬCj'mW;%>W x1k@Q=+FLU;cƍ5{%AA"O|[{=?}ci+xǘh.9Wج[^4_CgK? hid逺\tLizs7Μ{iZ5vO/.>>% E\1+ӵAES}ԛJڑR Y|[\9%">G?S N:uM71O+!1eOME*_חi_:QjpuL.+s$x$! sDqitL"(F$8 ˲F!RKFTEQ O,B 2IE8C6G9y䠔^}N 0(*03#`D s_d8Hv($9CI U"` T# y%A$\_悎@Tl&fFg $fLO咑3PfLN`e C!8e,!CUO )H DnW `Vjd~ h IqHp!dI$ELC'j ajt k)oq&"(X@P2pJFXa9y 2HF*F9MP PL @^PM|qH@bHL{]!H,;J)]bH*T&4ܞod+P18iAh$D  =H?8\,B({6."q:Vv ^ȧTÁ4I#Ke9OK$++3髇20XoI\3=JMG g iOOa>E|?f0HOcPA#vf^Ў ;Er'ƴys@Z#Z'Ǡ:{*Ao1@7j} ~bxh1:)5 v߃9~ ,"cbc?qnڦva0@)`!ԓ:' EJ92^0A*=+l*{8D-,+?(']H4i>f"8VK 9vGe??R`#Pt=f4́P? Rxv!B@Bx0 M&1 FU&!r13LJJF.* Jyw!gSOgGN i ]8Aq 0BS5PP9qOP&gb(g>8v@ bHp^#fB9s B)3sP'_0˜ "q8,FI^RUA%TϚ}l,9*$3.'%Bx "5h$&8`3W*k‡$(!@ns_K6M:5-[: Z-H%1,mN-'$9XntCޡ۩ cՐmwqy8#XUgS (;4upLz]6\J$EgmԳ~%6 Jx~+!qPX.yu2 WJDI~)~4钪Z>/icZtMW kRd+U+ q(HsT~_X_Omwf gjLRJf pHjtjL1E^h|NE0T_Ria}WT9=Z+ď.?Rh:%X[Pma͡Ծ.h)mU*7.iծ涓9_O+56{I@83IJ6}K➟N0DR4[QN6r;1woI,:6lB6;Fgh42l8tnEL) n]X daL:d@"c4-Wq\%)FKhZS3Äe2rj".|af#BQDq,sҹ CڄCS팆ah-=RC,vi@u˃@sZ%0KH9UJ kDPe@ƻH@qP|r ( Qm]52p3ZUI ÿHE .ZtSip*<(5{6M8. /2(0^tMI"UqHU.A Gt*F&9]\}P.kKCB]Ot JQ5: .{sdqXu)$Fm[7rsL (4 Υ^ zq Ǿ ܛYNUt0 v^Xts$m5:\Ij"Q.{M΁P(sfͯuIAȐ":nN߸ѭJ0nRituK#V$([L z %ujOBQ'}R$#]ִ)hiY8 Ն zu)V=4s U }bZUu:K jS?tIc˻ocdZ7V4bZtڈ1G_s9.{uA xs͎N+u}=y&[+f(dCPhI @B , @ & 2Ps Yqݚ3/i (EJP8BP,?'v~hS,p L5HYPJvA%Ia#pikٜ5ʊdf W,a'Gs ™ 8wCUDS$L:@p\ba55P @% T!8T ZH$-$i88|4R쐑"OsP:Derp ص& \I`PѨ|iu9=ycGunvۄ,io4(fw67[.4 _˝VbT9GtcHєgt6W_;ZEsS6ΗO<Q垧.MǪbAݽnUHΑ3hSR־^PtiYm:wZB\$(I6D P{]8*mFm:yInǦ' ^Q(9U<#.cG;u3W,`Tv,[olm ]ǶjOR!ȷc9ǩ Aֈ4K!T$Mc:;nG.z ><PCLXUĝSTR+V=x+Y8 یCp*ќxAwbԃFMHV2ьS>S8'dHR8V%2pNA$HDT\B㥮$ @e%v<"X :5ygz̨ͼR(i(D{Ur榍-JHE'G#}=7&k;h5OqҞ83+4QTQwFm?˳OԹgDW'z^mEJhQ}Bqu/3ГP.(t5 :ZREԭ0Rd|9BqQZ*قQmK ZNzeXi %ueFޱE[RNQEScfj!)nx}s3mD\@ib})rښCc( 60LdS7vSqkhh+y[]ʏ9 F.w SJwM 7mU5 ,RmvT#@'fpIbKt] ٩!#>εz\kPRB-f*3)޵AhD`8+'p/t s.eV`c T)\T(L6.* D4zLbOj@^ƙco퇻C[#]dup$jxF D5*`kK57vUM@L\FB'tl.AK4;Ȥ0H6Iy7 4O\Ɲ%QI[r}S'<@EO|;p7W4ˇ.MC!{իz,.DF.X%V9mID7rh}J.zy8q~Ҍc\$ikUéM%ħ`팖oZ@gUbϓ,^jԵML6?L@Fף\п9axO$qզ {t:f(C*5Ɠbsܻu uP=iXfG1+2LCXզT|%,S˯hݾ4F23^},GuQYW;nﶲֵ2Zni2U\nY6ڧSwwJF F)EZOv9kQժOq1:ڪ 2Vi#6ڛbҥRKYWIp\m 5"JL\ ˼w FgQӚ {p q%D`*4DCdVk*Vd$AȤ@Ihw*@ $)K0Ӌ$nөd%4NX,5!PL ZUEIӓ;s.0?ҹm4>ͥJT6 /67b!LRqI -'tzq7:CJ=T"Ȱ;rxF[:Ajb;eV "C/U_VpeZxPjj. 1ps:ֻzyG#ǻj[8{/85.N JNaNzG_N;;2g q8o^ 6N03QcFCrV@8en6Jpo:`m)[Z`Z0ƣwim5;z,. 3Ԯ=?Un:U `E"K옢v1 k8Rw[_q%k9iX B`@ėicEHvZGzIEdp%ySoV6 sN(3V\ ^؍"f;rۮ6kx#qhDӠZ(,w ;mzCƱ\dD1rngloM}66JfxrʕvVOTVRxw6 nJW5HWP0aF9wkw. 3!{'OjRut -ϸԎRt-9 k|Z+GhiS$ >̍~U, p-cyOyboK֬. ~m_TJJU451[[]\2 Zuuc*jD^Eu6*J ,[ɹLkK}:`Z4j6o]*05tІ\\]_iqM6f ) Ld< V7+nhc Jo3nrhPqmͺ.>E'.ʄUwlXXBP }G_j+ە U)h_Y4sfʅ*O`kߥ?-y\-ժ5[ ESv]QF:Fų&!Äl__ fpPic6]ožUwBғml.rmʤ\-a_ٷ)Vy\|tZ=?Erjڢ͗By&P|me]ƛiZAU8'S13Q W!q먺XSQ0!cb7pyJΦjaƘ_I}ɥպS6;ӖѳlzWlA}XUmBaa޺h.i<QOQR\Ɲ'թJƹƧ!-ilp"ޮ?"pb+V:]Wf 8ӎ6TYPILt}:2T^m"bvECFͰmKjc-pŏhD"2ܜjiҗ%U%ꥋiUmV*f24"DP%1Xe*-@fcI;1کqN,bɌ[J,)f{>nh%•)KNΥgH#Pr$ ڳXb93 C%*n4Uew]Um<4>' jTMQR-՚wtXoTs{[h* NNNqBTΣ5FtnmȾ(1Hc9X/W=Zy 齭/}TdkqΘ\ Qs&S:hYPOa`TmhrRJyZ$ݨ7{7vu6^{[TKE%*]QJk[4HebScWWg-0:zLm7^m5ϫV24/mOŶK7X‹Zt'?rǐ4"ٷ"ZLȮѷﭮ]wup[NZG3.mI&oAmyꝓvkԭJ[d=WνTirԮg/4Ei]L!)S(#qDn7;aZ:`>ܮp:Ӈ݌-6mCUBkF-BaԎҭj:45޽t xlM:$"LcMn_eѦ[e(Ԫ1jvTN]˷VmZu.TYH kPQdo/ߋpXGMGg(µCW6+ijI]LZe5UԪr}_6)ZqYғ@$ 1SnֿKyxvkg*sO^\Iٚ,C9軔|!e)}Oj8,qE=):9HqJrPKg:&>H~1?#-8\g3\xe 9H\5 dxN$Uahx'0eS (S'1`x:syAAJ@(L9p*sDw9P c&@ǃ @($CIr!)%NL IjA&GR7?`S4<)}0TҀ">jj?eLF+.uMA*:d`t,qr %kQU4Si* sdno|WqTBJ@ yMh0XR@3RLYj#h0!9f;'Zvu醖I4f$ӫq@ %JnoU(кˤ\λjqIx&q띧yNv`þ<ߺ擵S&F'ƫnwMkmOzZKՉǂݵ$?6*X ت^ l]1'cKD+DYqRuMS6sCQS7 "c&5YrtS&U:WːII*(>a!D a'{hYYr42^UܚjeESR ]alm>q,coںfiJj4%A#0֮)5л狌Y Ԯ^>0/|'p]ճjT$#&s۫^R%ZK\gO,hme_F\F'"{;qw^mAi5)Դݥu25Fˊr̔uD#W* rTVzu% .qR^O[EՀDfHZMCobcNkzT$5>PvA }'p6=56 !.ÌX.\q6ic{Z,#PuBfjKTߩqcPP, tmȕNVJ%i\QǥT@li"ǨXN#0Ke6_mL:Doh=^eN[Q]!^Y:Ӹ{ܠtP-]Ue6ӕrJ} u\7nPֹL"\1lR}zsK)g0ű TRSUz{n[mU=.q0#ҢԵFuC# gffgwnMK !ZVk;2n36SL6#4M *{sOiԤ4)̃X۰Vk4yq 2Pca_j.E ͨJ'ڊ9KƋ*Y `K "4jn R5) kʫ}2ם%BjFS*-M8#k>ҡ^մ Ei[*:N"M0 مΆTQ{ZpIYcCAyҨMI-\[L]=YQ *]yD7m\Z疺E Vޫ V-dLFޥ zL#oIvjl?NJ4C0L-%q' 4b@$8w$``xA {=QaHC0$sZ$,2Ki30'4$&\" x =3-N#LE& RJ@*}H pY{xqBw"fqih.HczP a B '`37H%!tO ܲ'*@!yJL ` $&g EYs^p:Te(O(`*:rB5f3H{ RS\<9pCB44O)\VZ #{[ Q! 2PAe#UIB':\Rk.9$̝M YP.bH, .pSjl=ҵ UJP$Bc L̬dTܣ,n8E 'Η ˰C)0u"#\d3V)X,`N!YD\t'4E9w|QB V HcT%2qX9ĶK0p(Yl"7hk1sTw}lʖUMd{`D%#DMѬ'VNv]ɮ{{q-P$ $%siD"nj$Ab£զ օPT /?,@]@#Q'OlsV~{ Xo4&NPoԡu$ZL3M="*Fã6=i?k9H*O8gVůno&KѶtuZFfShQ=(TESsF'Hؾޥf! xwt8ʤ!*q:e:Giib1EKҹiu[A}j5^I,9KR1}\v\| ÖdER TEkwҩQ鰆WsBIf5Q W[Ţj•0-5UMPO14=5,i,{Jk$p,v.s"L]Q4Ej!}@7(@e*4g4 :!٪OQiTVԹpcTv'q|i$4Q[ xrZEQk۪r~fPOR*5O84Ќ(b*VX.J_&+ qjfm+:c)%튯TeZMZ4̢j< @2J|MRIu5b[TO^I,*K(s;|}t7[-=oj8f9I*șba-iLdcb4$ĈsEfܪn-*n HPf]ǰJ;l+JUJs76&SVp{\48qVQddmCӪ5 AێWiujO :?i>HoJTި j_c. Ol )$(I\\PQ2S> ~#R7|E怨+S.`Ƣ5pRHq ? 'nIsQE]GlsB.2)hҎOHv1QIY%"kƚ,Ar梄"O>33%r{"wEE[HP"uGk[ !p#C I_Fx~NP)VGAVdU tģ rT/z>13/n6ڞ }T 4>tu\)uFvx;su픑hծr %qy\ 昑g]ӥXS= ٌZt=BʌF e9 4sUcpr/;W|ٰrm$;! UIÂwN=!A.|H@$CN ya_nb xS,3*JPiP}8BbL*2]8?!{bI,pdq0AOĘ `*J}fABtөF"BIdq; d|%%)݄i& %vC5L  \{7Lis& QP4T]VoAj5R;I=gN=K!KgR(:x8)碐w[;AD8Sf:VS".u:m .fD`GcKxք/#PѿfV G3ӤִvzVU+ݜ 0?N@MT{scKQ{T5us1cp{&h 0e[ kѭI$.e(VGv7}&&߲~̽uT9͸dij+Ч.u4FidyԖZ9`NgѩHԠ\zM[ ĥYi%L5d̙Fnxcv<11N*274y\i/wϢƶ.p{ZI'9F͕\*RcCۄjQ_uoI08ġqݧv(EZu1aX̻gx*ѮOf)Ohu Wڎ7M6(!É ۹29j29D@7z;D*.@#4GUO[[n,{]3H6nB3 ?d3- a/+m&Y ViCx1Ftv8̢ĵnKɫH2O;D28b1Юbfj\,GkKԮҨ8NJA7T%f#/u&ΨXE2;[@I3CՍٰ84IBu]q_M:kgޅ'D̷I*%Ip ^p^cUIz*(˲'!mPMϨƙ41}tc 81&g9?g~<׭ۤ^3ӁbA5@]T' Q%Rg9"()`H!e.Ք.S3L$Jc+ r~r20K#HzbARDa)2|10P&b?O)~w-#+@#KNܾ %( FWNeXӮaQ'5[iU2[: O`ZAIwq\C6'>0ͯy^^mJUe^#&Zʓ᎛^|߇Np8IA%-o4p 'dc.KZCfisf5yKfZD&6lKB[R!$D>-PL"jvVuGr'W*jROn0}03Yo1֠ ND;ͽ&Um" ш빲p]$9P'{s\N҃hk? [Uر,vRmzRkT{]P'Q%L:ۯ+vTe,{e~;  {[ڏA껈JM,Gny)\cf鵺"ctD;oꆅe*LWknvıtl"ikoMRRKީ)k1\ͷWc)/ *Em]qQUrIJ(Ʃ`j,i=P2P8P\~:-` X&ݴkm&>?wQuGR#w9ptΚnOKP3ހJhhQ;hsJWTpbJԖȺJ9cRzlI޾8q_ۚ=SFvy !"_k+@W:\{Bj0~4eJ*2i Ata[lu7|vfK~U(R{@Kj=⥍Gu78HIVQu[(`:,8Ef?@vJFm ij eP{:ӁS,D %+jM<UX;cvp3QF6WA-v |p嗄TQIF9(Eֱ4;]l}GsPkI16*{} Ӧ)Lp?-~U8`$>^-t#Sfi)K sDB|dҠa8̢zeQ[`zBoWI"kʸ*j[{[wUK]R %O,6tq=~eS|ݙr!#$f; ωBTQd s.GP$S_(P邝V5#AefX2j[Vv=pˋͻr}J4), s+J/.zx<|}"P0NP"'c>ԟKX\x ;J❘a\fm:jvFFdvΥlKl6H# LRBoTa.^pM2EnmK,YCZWJAg"BpLB`#&CH"|dte Bq%qHT&1e2eD`"q 10p5]_q8*+K.s>ǕztQ ORل$]2ף蹺i9*ovzKxQ-i4^'euJmi9Dhhųm=[(eg=(}D*H/?OjF%.4I8JC~h\ #qnT>ӧh"VXD+ͥx.k5C' zU\^$AŎܪX9rLUJEM`9FzB..+*^pYn ?`/(SUuJEgx WQ~uuJ9* ?{7 zTi70}y'*lA(⃵c"֮K>uԾqeR2& )hh҇.[ۍMqIFV~:u;mǽo[ EY9 JҏLL˞e y^y\QR꫶FOi3,QKnHּP*6d`eTyLZcauI3WVrV5*Me,M"m.XhAQ5##9H_21еg5c t&:TO4&ң]oYHs֗A=.>d:c6Ӣ,J8lnӻcZTiQS kEVZSdD:oҤ3 pmݞZ,{ڮ! !kl ;;F "J B?x%izR֫^PL1FGLjkINfvfI$Qc킛[J'KhKyy\Mi!*E7#mmpp($9(Ĺx4A]#yH1(ҦРW>¨,p-*˂Mig\jg#5]G 2<.}KWEh%qcEvWv8? ufiIiB;f*~ #~/{ڵ!ǒ ܳFʁ}W8 g6~_hŽN'7o[֐no/96xa!*+8$*(| pXR$!Giv@*8/gL/FSHA1([< Q0Pht&b0;p,BJĎ;9\&d% HզcTAPpQ?l&rQPLR}dH0Ih;g !9 _a`t ';HdS3iĨS+`P`Ry&0& BpLs˾$ x10_PnQ~ͿWQҧH_,qo[-Ʀmpw5^j$6j7W]J[T*UdS7dD®i"=|R~ʖܸ^)TW{-yi $a(!tT f2>x$L^8c9p4, @Tbbn M9Q ʆZVܮ:k`⑁&} b/mn{fʴ\/ȸg+R فvNEz73eRXJp2ĕN=+ۡ^ 8IY S8+MDqHd95Q=bXdb+&\3i9*^R~e6s1%7U(sH[e7rY*sH($Jݜm:f~MvT:I(f{XCBO3 j*\Y(U`]?P-xVg+}#2kkT:-dXJn`-ό_nuӅ̱mhs II)jN J~`tbѵsn*mG)c$Q5$*8:2/RZB L+k-5T#;i9uG̫ýlI,Uer ߪm 02hL}]uiMD466qI AѦz WĠ$T,!/dO-C.%"8D aA$\)S1 2[FȀ{231RTfZT]4noZn\Y46֦u7It{kfUV%e mG5lv8qra]+IJJ-I$n YOVCҼ@}+ݕ<ٝR,kLcnRo8a~,s)ɭ~OBztfƍ;]m^_pPGg"wT]SVP65j"XzANRDGLk_%TiXd&YogfᵃRJ\yU0p3 X,qj42<{TWxG`B12!J-8]u-  brEƣ; 1.EA- r4.̸ŕO \_x)Wih1sTkk/֚EyH8jeR)H8ԡ]#*FSPjE$ڢXuv{8Sm4(5^`)8{-G3~Ѥ4c~(ɭqOIk RquؤLnfAC/b-sY5yPF +7ԢFE9l`_%eB XBaoFR*5O0 ,_l:Ua']hW8j^}]P$ldﻍ=jշE6][5\\N1mBmʜ<=G"n Q4 '" bGePp2EIg ԙHbx H !@# +x F$C8_+a r܆ CT}}pr8 ICI\!p^2㌡ĖYl:cQ~Fd ȘJL]I} ~ q "IM9*U8C1' k)b%P).4) gJ̢](N%rc ,0H Y!Jh10,OdӞ->aCZgx ?zZ\bfIH{fw6%(e43IvLrl$@(2*beP%fdG'#ؑ:Ho"{"6؜F'@dhCLCF@F@"n%HR(ڵ!B+U+K&E| }TP5+1<h_;nT92m[uY0v`檒퍝Wͤ8E?/֘\p9ա 9N]kҪ@)đW~{=<;t*&p<(݋V-V[~-Eܛbk:F]I7CTqϴj}|;wCow6eFR&EӤZ=/ֶqBߚ-}WSNOuJtݮj*L`42nGq+pbJ7ԟ_*F/eo-JsO,SkVnjZ\),+MZ_z.xu oPLi@{mխx,saDQnY4U:Ι>\cei݅ͭűqhRrzku{oP5rĐ/O$%ZbpHM{b6ݷW Bө]A\YSDwWX4=Wi@ ?hr[CԽa%_W LQޛT5qa){<#?b`5zNDJ㉋?-5GiSc6!X̚O9ȩS]U+xMsJNֵqp0?xyٹtgO4Xk'M:]9\J$}af$(:Ң= 7$}]s^q$hIks~-yq27=ݾI}n|#ReIE˥>$[oU>g8t:lܮ cX>ċY]nI$lshhֽ>w3m% oޚgWD=c龕ܺ=95Wt b 6&:ZI퍅OIyRqSk-)5>J\Ϊ1&8/*C._icI9}c'ϫ7>B.@t#{dkFe=E6޿McU[B婀1]PӲz. {Lt= BG"w}֥:L5)*R;&O~T]. ~QS22~I].Ƕ o= k\Xv :F*vK5. ~ڠBzO7/m+kӖ"@Tri׃TJh h4$f0`#[9aU g7G+¸~c_OKG2 in8R}G+ͱi4,AO.X]X֤aPfti}ULӨCW"qG.hP `1Hƨge:SΠq:rC.(񸯷YQd 8F93USU1̓՘lk 2O|eؗB|SJ{J6 M0[+z sMK!;іգ?fۚl($/բ5rW?l>ƗMEj#;#fE$ac= AhRdŮKZ%a"gݾit1&QPuMEs`3YhH=լ!3ʚ_NTlnmMк_km m u@Kӊ>NyVq5co mڵmd}W2L! bC]ЉYb&qU0;0@"$suvKaWҨP9ZTi干u^ Z(6@G#:M~v8F HRc]5#/q92qw-Fkz5_1@t<; 4^:3rf<ȔRzf}EyHaz2A#&j(z@M7Ujmr_.'q_q:A2p1|߼qa6eÙp9D 1:q b]GuJ7 ꥊ$Y{!YC^ !3Î0SU| GT@DE^_Pn$comhl[g.w쫞iSRL)+K̒I?T4T\Hſkzֹ~G޹q4U0nW[ܯTntEiC=6H!xe8G2UW.!WDb'YI˽ p"Adx$"&ZA%\}(HȐe>1"S A&+i8S!sA%2X|sQ eg w#ZqM@K|_p~Pj>IjSh.Ln{=S͸/wLӨyA{b7"zړ%YX^zuhV\zCȹn5b@Wg,s݊:aSO5?ˬϹP _sAZ@u&O%%vOGCǫh+5@n9rEvp1q/h p\)S$ӧu}7!Ĩr PB4]Oi!yw՛P[\b5 5jzhV@{O\9 E@Ф^ڌ=HʉPMkh5H`qcHD3rNXKE.iUiESꊜ褌4[.L -USȻ<"J}F&̫`Q eQA&$O捭jF"g%#Fr5*}JMj=o6n[z4U77;VɨEn]pEڕ- -`Idf{qOk#PFAgg}NhOn-ܟxEmI-4{瘀VcdI!0I$iP$"EbR;2IAD O!9?!qi*( )4 d;5H;L$D1%2k+mS $jO¬ዮͪI'fi8dJ@%B"_ den^dwmu۷z~qN@.kO;1^ldt$YQԪ 5]2Zᖠ(m}}׍Έ)Jv_pźMmA|fSSfqvO#I^ͽ~]? ͨv7 p a15,SqŬ̄`˾ W!9!c&CBMSwihBsp$K0ES$Dp |L1 TL* '/a" ()b#@".YHNj;U;OU#F)0 TɒR'L}AqI;2 2`€KÄ(!$Ov}GPu"XqQ '4IQьp  )gfr,%\վQ!C-r_:0-N( :)7R~_^?A {Gsflm5j "5tcM.Ң-o -Hrʆ6{sHwsBT 9sb4*~.}"~FM|s?\1*hsaڂRuo B0T1qZ=d9D ;* dm9#-_"q:!OW)#=9A8KS()} $$r0GaO V^ BEABx ~*6.6@O()LFmhU%}SxLV4H^Qy#sX[,0Tҙs 2 V$ <6ć9ĐTeX)$s,:I* a_V{KRj #suZϲSqR׸)zkx$WtEi:ѭMuۛXiu;vdz^!oIB4+F΂+;{|[% 3Vr9Wdx,浥L?XP7v|}sVqڇ&3=Ct;7EnmxeG6 Xa?<@$)^# , 8jI]] 1._0 'ˀ JA`p8?Č&aP4s3YZM 3Es^Ԇ)f8$ GMp' OlIA҉I>(!HUT Hߋ44ҮlT\ KCf&IKQ\{#uVҴs5Rbpo4Y A81(T%jck7-?6֡ju9`n30әi#WyAhh(]˸TW.=-ε*EZ/kǎsˀTAa8("mJN-FJ]^uMtսe[×N ekʧ^ 5AT$aT6ϘUӅ=ޥGcQMj_i1mZ^ޣkRpi=æ3b)'B8DfG{IXP2 D2j+B$Dn9G LBG|P8xH bVped!@3y+h =F2"fGiH5t7:M<1'(}}A_gt?i}:,7&7u* iZ@ߵEnUY6vܕzQ.kΒx!Q7t4 .U9mW jI WƎT}zCnr#If UARzMMDSH̦ꦓDSAP5J\*I%%7M1RYQyIO4ZR#Cb 2 FŮjZ T$uaLfzCG4n7{;Ujz2)\HZ=Nv ɪdn&oN{Yޝ\kB:; ܆D41ϗw.=C-.A-2~㨲XFbr9Ukdc!XL}a%7Oe7U!{т8^5zVN3J48"}Ѹ{庻 s:jZ;%vTj HId ((H۔8- R؋DaUrN! ,PȑM0 0洹 y}58H &)030X$fk(bK;A"Q"ięC׵N dgYuI&a9qDXcsBH4 W yWbHI'\И,%9L*b\`s> /C%@ KP$)Y&CdFW$"dJpG)*@K,VB̩G*PVݹn;EQuܾ! cԯPw4kMA8Kl8.jW?6VnmrMCk3\Oo1ksKPEvkp7#ÐLϳ6[_P=QZi<+=ΤsNu3K}f{dي SnGqK*PM:nL?̚&a}UXNAFJ5) гuGwn-kL&cl*n.hSSsuOgLTՕ=M`OORy.t{Q4IѥQBI'ڈF8s{~%䨿O5ZgBEP3sZF7Hiy?R1u_ lUNR:nyU$tPRIWZƉI;@(8O(r!H$a1,aP$HL @8T#J*gAoqSЦh 3XLI0 @T8p)@!P Bb(ǸB#5TH_(5I{8q{!|FRP ,#'G a/0")RHgvN RI5| I!I0 (F !p?iOn'a*|JiT AHx&@U2WYz+w.߸(h -`D5'I |apkmH? -(8F)צLSxr Zj+4,2 PlfJ .7#³}B#8Pvv p ]Lǃkc )3*Ϛ,][woԚn/l4V .F kjm?mo\jS UxՠG6\k{bu̱(ҹe הI5d>FEGƝ r#&鵽yi%ir# a٭s^섑؛ô6~ӧaDeNm3Q›asZִ O{yYK=OCMmv!ރSQ[;|q:mtaVd,j FOގѽqkOE7a*6]toDf%(AmO6"3` Ls#?='Vґw9":UBxG>r ޺ۼ:kj I u9DN<29ω"% JC0B IDőLdOM$Vg/+I!RkK|#uCr_4jk ecocst)2SOI::d+IOnJ@gT* -bLJ*WReM,ڇHB J0dOf0*FpZ@SI& ⧼j%F"hOrCZJ)/GNDs_ՓBT$*\D3(`*]$ tRpl%PvAQP j`) H8w),0A, HݦXi-PC)@ paA!H=,e $x0J).(| ~\d 5'!,1# Й`L)\ L&#{o՚($<+:WW,{n7³jڗU\f?æ43/)&6}T72F@8|QﻮkZQ$ 4갫jRrs1ѯe3_Z M{ਪz]򁶻GT^tA?|v&/YS%(.V&A,!9sR2RJwa8i{9X(sI1T= i9= h$j}\SIV`A JjPOHPUGC ,!M& %HO-CF2 ġP p) *TRQ%(2J(>>? PD7awC$|8E.ЀyB8D3 )$";!`Z3 F KJ,ї2 [Bg!yI?ndq wy<~e?ƮwkJ( N!C&cQѫLT?Һ"1p233cW\93sr-LV8gS؛ ֡ʊPqxMdOn-Ā -̧UAI/&Xr.-@A 9@}*E P58y*d(J@$ʸLڧI%)Eq)H%JYHgr(b $B&4y–5H?L)AȄ$ Ǿf4BTlp.z@ CfOvp93q2qfD!D8\0@=p |!,O@ARg8k 턄L 81LR Lt/c* pUBL ,%N=d3Fx wUlT758'4%>l;MwH-xc`6{mG~ QGV%KRe;U.4mi]?Kӷs泬~zw:[OⷠmS%?W@HIUW8\sHF wZ@^>&^ąrd`bNvRd  f$i-C?dJIq'I(&*O) <09" NH)Cv(C&]O^EZm\4H-7jwgRVEj HaS"[Yi%3,A-xB1$!( b'Z }p5t[<8v _Bz\FA*pN IH%sd1<2RHq=0*>r&!F$`g #LЉ!4&D I>a;V!1=? s=G;&;! 42!ۜdh' t@#BOxsB;4]VGB.pTNqvGedBKO@1tЌQS>* $9S<@ r?t<RSN(\H=rYIBd`(j3E%ްfr"EmHkIð0 JXx=èC+0ҁqI4鈰jIc/n,ሊy\X 48嚒)vǺdMhꁪhw[~I~ԟ (~\~9?9\7?+5Կu|sg/[|Ӻu(?dd sZgkt+t|%vИGWH+w#ȩw%i6-xt?yKY)'焉)~p '' OzD )< @ D(SPs<ħ~p>%A A|`kKZeڰĂN3tȞ8y`rP\r2I!K rBBINBwg ,0`, %z L!0 JB TffNpH jLdџl\QNPNTR&<{M?npPҠ)2LR{a8iJxCtG09K'b4\p wb792Ì0Yj # BGIh iDf]8e!!1\CЂ45S 8Ȣȵ ]z7+u-?#sUT(n4Nxa]48R_Pj!!x.pU JJ$q@ECf0DjE˾$ c(GH=@!4v@ .fp$J$ftPAYO q@I MQx*R$(&S?c cf" #kU3ö%kKB9F}\PQh;g'yB)wLAVt`z\ڏiPg j-mY{@7{.wvӷ%>U4߷-[ml"-d41jHAWԜBg0qnwꞝ7WS#KK>caݵLoOmgeIVЦV7UbHTU9\xSWJɍ()RH5"I*xĭHPFk5(!rka0tfp" B D*THL&"6`*!^I8&0FTMSR=BH!x5Q|RLM0Jj*0K*TI#ѫ C)0Ԝϵ A?nhbväcC'HLq2 arO4FS_{8cꇗI 0f&8a CNS.9/AĢh)04ٓҥe UA2h=&9&(]aJJEI $4p RfnO*̪]D#!)j ij 0> (%9I^`}H7:iB< '3\ `",d$0@ROn)DF䖢9&a~Q8P\$%.la@AB*չM):v&ũi$gcSDPDq;wqpЂg&Q'<зՅrRV8iklBƮ2`"@j3DN8FFdIY&@qf>.)# /e U% <@Y@faU$Vxg HhSql\#n`v-evZ΅hYSiUt9thlVe@~+0ѾT]hyBS@`YGmW{ͭ+{Jf\Qj7?R@s):^Y\'ݸĊ :$qՈ <#/JƑ(@(=]rv¡K]<ԶeC&֥aNH<)PZ ޭjgԪSf5վQueXN@Z,HUPB(D ʈP% źY`wU ?^ y+չm[>ұhsX 6_0-hbIL&^k3qE^2BqL8$ *IT /tLbO這 eI^PHBwو`$"(MCBT*}qIH;&fS g$4l6kLhʬE7[ j1靯C0Qqib Eq>^ ]UzXYI;ClPR7M}11>P$Z`_i pe ey[=Nu w,JweT̔#VgoOBwp(PuB l65mD%|̹za*Y`ժqBC* aR5|Y%WB OS $'JMg?yF&hG8Ffjᙂ3sHPT(DN6QdU˜nraLNV5.=:VM[VӤh}OEYnu+}m4ROc[NuJoek9=hsaSWX4??SrS{jJ>3kN6O)Z9i:)zQKq5rQʟL*A[]ٰ= Ahm[KKF붗f28%HqQO=ӨvV*y`ӶmHz-5Rgu_XYӸoFn֊h kB+ %AQڷ鱔^KU?Whc6&F_n UUl=4xnMZj0>RVTk^R>>>3ioc(}UԮ{h<*U{> H1ϧ^eFҘp&aGHk1gt>.mjSZ:jK0NfZur7zfHXXF?ˣpս7;t,PP֗D7U21mǨ^8MzUpjOL2[JjqR9jKi[7Cr2"[AvU)*z抎-`ӧM/-Oͧ뇹 gv{L c!)Q/A z@2W5/&@}n*5¯-+mWj[u [A(I~%Wga`in! fp Ј3 0NHEv(%yLqp.P5$h$XkW$π].Vujӽ=[VY^61ŵW5KtyH/XJt;~avIC Ѩ"d8evޠ6꿩Cq]κc(UtӧT-g7jew iUoJB, ;#|uZ7i4fJDm~lZ:aKeG1tu vn)U~kiwZ+5*Q@SvT:[UƍqefUeM̾FGu1,rn+y{3qN>ncQ&jlRS 1 IqyєʖcZ6tV-:YbN]mJֆը᮵KN!˷;r G;闹7L%%i]ǣ9.0C>iñl '[v6J}o]RZVZO~7$<|-~Vˀ$ij ᎒(]Q#؛mK2Z9O9.gGo;f-o.-Qk)T-(O.h,_Pg˶em^5'8N))Ay:Z͞> À@I HCzI!pI(R{!Ht$r$q !$@E) a{Pn:<'a@׶9Yxè,G!b⋐G3C~k\rmV`gU}i|θw6~݁FE{ϵHjBK: ϿƻtAH'JL*h {sr_rʇiʌ @2Oaqv]E1(6Vh= &J  *%A%02 mdUO3"Z0J*7u+x̧C?KK_GSs*S"ʙt:{:F-d6OT TQqg\tLhHG-"pU%1>gm>;&4Bpvtlmc+Gr"xhN L?dn=+OryJW4)8ֺT. 剏3ѐewmj%Ԫ[t=K^525PgwS 꺣)[l(7 & ekVPnGv(ࠪp>,oXk9Cj63: m]ا1\.*Q{ ;a6+2!J r$D+p[2Ypg9`%i-3t' \ZS.86G^KȇQ-`x<"4T.DPs!8Cu@˂#niNV6}m>J=j98?lr rw+mߟlQ @{ ܛ}sRQMdtKfM]vmZjs3AL\VivOr.y:C],k~mF 7^DsIkRP[TM!~U{:b=PF2B!K˳a"pbe#-(0 2.i$CP7VH8b$(H%ą$ "K ң |-s&j0P5ȓ(N M$dtS =#Y)IQ g mjW< Y2R{+vmDVֽX|)0mvJ'1sk6Ќ&_#OtK Mt15#?ȵP6=s\yUmլ-7J7 [6卨UQKЌ'u][[e:UdF5+JϫRp&Vp$zziX) PB''ɔg(Q`}KLz8ܔa=y~X:or4pe{JjR}jFzcmTo7-m*TP!V~u^iUvmʩY,GF757TmVޥ7eJiR/Q[9c̜;.q}qmi9bn_-Ux[OVPor.❦)I\nW\c__1M 7kEQ{ڍe!ouR}6ӳ}[zYOI6--:w.˃u[U`mV7K鱼ߎ5}aRn @7ߙĴśf_u[ɗ[Ti+Q꽴X^Qro6T5}_6IڅKZ 6=2״{ Tƕ.Z\\{ <.RTkvT%-ұ5B}g[Ԓȷg(ʥ2j7XէY_ӦdN:̬G˪` 4,t ᡪu5uڹ-47x؟욌-.mbQwV՝QXj9Ñ^9R:u-.֧NG!RrH`\PzV.M\VE o+w%5 zեg+1uTDlЦSUPÁ1q&,Dr4Hp e2r>J'gd6@I +4$vS8aWSD R]t? &sABtDд 'G=B@SW/{fAk=d#CN֕TL)6xzmkS#Bdh%YF{xnUI7uYRP`>G%܎eYY =Ђa0h' 'LHJsN3'B4JG!xN% *'@k=KKC\a%-\ l  p\J$ (@L?P<;A\)|Ә_aґ,5BW̎A5 IúY8iP  ΄c_RI'TT[۩Ԥ(RRՅ_Sou]u?RVjtȥifk|ST AH8Z"rut.$aBQB엌(Ѵ޵ܱڀ2ȞUqƟ1p38E[78swqnQsV6&i선qsAlť%N q8qq@'A9|"K6GN2S޽]iP48cX+5\  ڍw;^ꧡIPpS0?.C M1!ꚏ'88. bOxlDŽ:|02<E ;j0CdzA {CI˛P;J'EU"D%)/2{8uF֒A3'K@2JxADpI8ĖBr<2s C D^c H2 c@s(pϺ pR# B&Y ;ĤO2l \Dy 9$LV\1p(0gj%̨^h(!؂c8ZZ&L!$50r0S邀9 dЩJ8)= %3&pI*d,a Y` JH&sRqt'FHq)-Nb0UEކ "NRgQ$0tP" B(?i+%D>TXn=KDY!e "IEbCH,5`.0 HG82$&ri89{; )!ԳJkrX8L!  H.S%!; RQ.R!A} GK ܀"dy H*B@0D9)NҰf'J#$P,MaR'] {'QLA {PK*8@ !5Ȩ=\Lʒ>q!JS)'=BRG;t 6xa.g +1`a $ (rQi )9`  (WkI(%r\r+5P3 ҆TäF-H$+I)@  i\w0 TPZA@N( ga̕F=p ! =(HRNH |8g dR8 5όjn@B{%}D2+@(B$EMWA(ӑ̑,@&NYa-E3"8L*+08':!$p(} 1>*10@G !8D0g%e @9pRUp<?{Nd!c9($C_0ԓ$3G@b=DH˼@ Nn3$$IOb( 1x/jB$ &D\!HG|Q8g*Ì)qsRL9A{HBS))!A%&vv\{< hXxc&1R B$`T2 R;;5%1 d$*jR g$).+v}+J03!MRJFP Ot0fτ\4I(U@rD 2,ı@3ό @J 0&\1 PP ĜsxR" .aGCl*nSlϓ\^1js]Q~۪%9|1U'n۹%ٯ:K˴QOdRQ}W%=_wT`_m{XvlAZJ:u ziN҅sBNZ MixZr4624_gޟB-osnMTyic4ki;G7*#Rx(-Ih bcim.ڷQBU _/2;i Gdu-N*QIUANtPPR_L0Ye4#f,>xA d f1!p #$Q\Z* tY@59 YCB aspBUDTpY U@>&H(B˼v H TB"A:A 9LLO5T~  Z ^Fw>ۿۦC>ߩhp)i'Xf1W=bMp$}`q\O| m<d(S8p2+Ma8c rs Fa<&$X * 39BB;2`a!d o(G (F]2^ SPL*" oi Pԋ0T2Eٔ'ːq0h $9#B $4] P|:s#a! (-^#5Opzs_vF-gbGcr+礉9SX}N֜+,W5 r LqC[Q rʎ-4xb-%|}s4ƚi#kM\Nlu.^tiB&1&TT2i,[r7 NY%WCu&XRlmݶ*eFiUܔ֛^:jZ'}ԽG:MϦ)гwFJaW6ӨO67}xwX[|%4pOl$χNEd0JC9 b$D4I!Z8hp 퇠 &L HXiN%?:k(msadapter/testing/ut/torchvision/assets/encode_jpeg/jpeg_write/000077500000000000000000000000001436713554100254265ustar00rootroot00000000000000msadapter/testing/ut/torchvision/assets/encode_jpeg/jpeg_write/grace_hopper_517x606_pil.jpg000066400000000000000000001651421436713554100324630ustar00rootroot00000000000000JFIFC    $.' ",#(7),01444'9=82<.342C  2!!22222222222222222222222222222222222222222222222222^" }!1AQa"q2#BR$3br %&'()*456789:CDEFGHIJSTUVWXYZcdefghijstuvwxyz w!1AQaq"2B #3Rbr $4%&'()*56789:CDEFGHIJSTUVWXYZcdefghijstuvwxyz ?<`dPLN0*Qc듌Uyr+ٛL\9$`c涥֮dz? o?ְb¶H8;-L1gM' GvZ e.܁׾1*,Yj`a?(q*6y?xrIɩȇ F8B-@}T,g9殟&58=AO0Yt͒DAy\}=']=b Q/'8Q=-"ZXICA^`_ WnOeP^u(_3y_G TpKqX[DD)xLams_Mxٕ|M8 [yFʲ ׽{ݏ}i{2E~nG_AzPXl-rm>?ʻki7=smQp\g 1G+|m&̸bpExy7NWfнn̊\ iRH@@IvJ6X2$ZR^/? <î۩Gpu\5;=ǃdr@@  ^CY6l|CYvS&FRSKK[\q"OνSmef$qxGPADbT z8``i!-G&ۓJy:ZBHÜHg\x#rzsړo#9$ x'KaTu89oݜg'G O( rqOl  F:g;{~=Rx'\Ƿ  FOLShF#ka 8 [9߽,y' 6Oi\c߅J! |p8N[FW OJt.aB ^p@d(EhȢ`aqNxU);Ԍb˷6:MTgמk=$2 NyZqv%^G'Qh Vs=R5,3ChZr:Y9\sI2Ts}I\ssS9@xҴbGu9CBQ2!H=«HD93a\vSVf_0Y,rsV*I''O|tySd1JI==X-XVDR 1?t\%HpM:2we{qm[y1Hx*35Zɬi~$N}q9VwPv]ebösL !}ZiЌdǟRwL=Ojj¬J8z ~ReFhJ̠pfҍ")+yvOEx݌{״x鼏:#|* xN0I+ȗ+};xO`iɝ$~z/jѢAzymk#ݻ$`Յ$`H`spkؖ3wZ3T>K&2>4^wzk,s,1rq]'ixjBFd8ו(on_ćcԚh /Ҿr.~ϩ)Ql5ơ &n5Epgʾg.52)bi D&p=eVJq8G4GiP{W %3,^0p{X4Vmt 7]1#]lX.g[ wlFp09[RX0{`XMolI7'u{gՙ=]Wviu$l wf,'`r $R2 r\MdE{S`uc=h<q(#{Qi昃8 0G39O4q4uv#J@{8F=Ҏ(67q=qrz aA{҆8=98apjhpqCPn'`Jz9R`@B)ÐNx$K HE H87y?LĐtQO9#sǷU>bj8$uUxq\5w譇ÀN;Sr[.c-9K; rzԫ; t~1A֨E==*lDD D$q#&#1xvǟ=ƱFzW<>$g3 +C&16ߋɉ[,3JI iq>35Kң8?*'Yg~]<.Tl԰ge8+OAzu>C{ =wc qs]Ưaa[bbb:q⳺zjZ_ ?j6:]43(p~/Mӊ 6x<`&K`F{ulx&ǣv=(8 JWHl=4R7S6pHbb7 6$'<`k_ҪHME92jΈB6nygBIxWQ^GLVc2?նr|׀<{WڶiO nY!Wwk2Bd~Z%x ucsѭVTFJgS=rzsWh)+^vZDݰV* մ&Hi1UU0 d$ d=Vcs]Gmhd zV%I9'RN{ioUbU qW 1xۑ\Tb ':d Z}TY9u>"O<ެL|ڪ1Va2#]9kaI/GB}j=le8 g8i719s^"A 'Bq {*OſT=2zeO~/?Cu j]wbf@+ǭ?SU|O5c{]>kbI]9+Gr ǜVks"<) u+c'A~>yQ"Ӏ?z2qNezņ1Y6Hr x5|%\YZJ2!\q#_|jg^`]#۽lE&Ә[NR[n {gVь$Q(UQhft2M""If +ur%Wз:Gґ̃e[o,˲0̬A<|s WuZ&M?(Xnzks,=hYs_ ԵiZ{liQ`c̾@ǵypG g<'ϊy* p9XSi'sGM[I٫in mU#АHo\x.ohT7X:3Ķ"œsҳnt;fo[&#Sp8$q~y#kZڠ!Kr;9'$käZ]5*i]jlwedhCloR= Qx"fi7qH&Lz;:BSMfo*(riVjM9 c/o xj-"{Dy%g'_xXK_?;ywιR#rν|*ϣ>jO}^k>𖪋Q,F zh2-GKmX珰q_|8WE2]խ܏1G/ǟz[R:NfSf>o`S!RT*}1MQ:潯 ybUQwʼv{ mn1O p1aPFFqT+'P›'8Jˁ9AWu%@83ڐed0䌚 s@$)GqFr6sA7v  g 8o=NA^ݸc'=sޛs`.2q(x)@nc< ' @ Ϩ)bpCFnOaqP1pv(wn$c 7|?Qќc'#dUԷW+Y"R5.yQv㧧}1¯}h OrGoq^?WZdt4qQ;KSS b<$ 'oJ4%(9QP~ =SŖҀѰlէ2%+2=GP{u 1'΁YPD,GpO;SSJK4ӳ$2@OSڽἋfʟ9k2xGzfV몍+Hz?s= g˯_ =ьc?׭WR 0'&xk$7i TӦ?<ҽ>GOW|DŽ$/4Do߂+B# uZx3PW5GI|¶$lnY&1ȖL|:߁ץw?"i.0̻N+#3o([ C׿kAF/Z"׉miko{#9)Hr>חxp+6vdN q[`$י>+Qk#9&IlOJHIo1Ry;{N]dZJ5w.'k sGRY rH`#:ΡSR>sY$:g7vҼMRdxH30a{ЧJJ*0I՞~ 2{Qwi $|M݁ؠG}?j/: t=U _&O\m%zay#%8ҷ25 a$xml!c4eZ;lck:+H[Uȸ =kf7VViז%3'+U%tj֜P1֖WU,no`Xf H|I -ֳsq|Ѳ @;NrOSm`6^ƄVU$b$U{KⴲYsEok.U8<VQ@OOʺ߉K:'im# Hһ(+'KƷ^KbȞceGw/a}k-wKGsv?}?܌0C'5o&:ͬkomYA9 9מkh3819B8IZ҄т"Y_-.PC+iފ3όdkP0<+z燶l\\~'lzuz,QfKK/.Mc>ҡέ/s.^3v=ýp8u-?N3{h2u>}GQ[]r9@y={,둚饈Qjp1$_&89ӑ޷)( y;'R mF][=ֹNUE*3OL81Yyv6KάAPW1|K𖥢jqD'ufuǸȮ^>iIx3z~,dQٺwXUK} i:ڜ]Ρ}&ZF;Oᯆt K44$\I T23-u=Jğ5,g˅a,X99 :v&>:@]:ioPDq >INwIS)}KGcy{})[`|<=zw> Ou}kًѶ0B7}\t``޻\c8u77ZOKO/-,u\HH$ @ quP\,dGSݻxb'c9`?OB?}3%ba-`Jt½'7e^|4@ ,3H4%c}:da!>g"NcxS 8kך0Qq%@2O\jڬqj3LPG=YL:(b*jRjJ=Kqh 1OZnzrִ8zWΔ\ef #H:3JD>J'n qL@Na7vi c$dq@ʑ#8$h$񻁂~An''(<Z33J^eӈ?ZP1q@$N@);t4y8N㑻 @97g`~U.~\ B9g1@O=zRz۞*@09=)0q昆q_$Ey\Hds5-t֡ 1sSEqּد? =]#r19$^` 1_,@h[ gsAsPsBQOLz;`MNLcD lc!c}:W|$f_JD$r+NzWb+$j*[ɲq#% 5?Cewčx:?gd?+9tT"sXE4r@ )B qȐܣMIe={V}62He^L33$ƜǕ9';Ś+ҽÚݨ6# drx'8U,d}zn-L!Mqբjshobk1N3\1^ki#&h^+n"țt9*}k*kOX ק|.~xvGy3\Km8+,$W|(bMO,0H;FvS~y~+'$ہXBK|Ğ·{*6yq\ϘdzWVЗky9M@998=(v&OVmȵ.ӏ]drHd Xdcg 6gҐw`g9h' zq(BN@89dp Vr2pMFz[=3KOrqKZ['3 B26<Urˁ8Ssw >bIlW5,2܂8)6H'q@n2~d@j'%Gazar«c 8ZU rOsU1`pHgr^q@\7e4帐 :On*IYJ6DsG)M?.x$I8&6e}:{UuNKD 9穦 \:+{ :j.KRebqʕxݍ D׼/O Mе՝ܡRdw c-k>GLnae jl:A\)Jjfu(4]6fѵkfvc}Ag#i#3Y?o7zZE[LT2;LNO[.[KR\.]W74^G$XLO>TY鞼cv{SZ+"NI#y4"I31{B:oLWRw؃5䑃WV %q_?sנ|.rBJ2t sElDx7|D[(|I|SЌ&-]f%1FG8)l㞕8G|-%N+ py:Qcs׸q sSH8$؏sl`9223]@aTO9!O(C?A^oҀ0zz*rOQ@J=Rrztl`‧}b~5a #Z\GGJ \rI2'#H4[AH k ;C0 v܃N;] 䏩cd5>z=Aǵ4dqC>OUB 䓌4ΙbxHtzM] ^ 9&N)rS32dc_YzZzrLr+UmS:tuV8#T):dң]:Fkh6$!xzdBc$0;}j8\1*۞v5'+yXNqVVN:`Vp{6HGLT(h e˂yڸ cs)jxF`@u\n84KRH[/w`"рz`SvSH"AP+ܪsN"3ǥ7r(\rO;0n#ڋ" kuv^b˵Pt:e]d1+F%s1爲P881-Ҩ@6 (x´l6 $ P;H%$ lRY4r7Ji,y;X@2*YW,:ƅG0Fȉ;?/PjztF ]Q^ix9쏙 GsF[ 䓀>}dfqF885:Jq˻9m' "09*:q^-$4:(Iop3DL& RW=*sRZ]J|BŬb c_Gf$|fAM]DvO[H'^ku=Rqx<p< p%N@w>c(䷭.oAPyv=q@#>1'8ךFx=ɧshN2xӊN9 3cI(Gtx#9q#i6}iNz{`0r 'ߖia_/Vr18 3<+L-$¹ϋ&[{Kg+$2W\LD* ~ωKSėc(վ:o.эM'Nk־V<+2^K8`_B$A~xri}O bh\8;>k>0x<&ipDs* ?ҹ -յ]!$qnٮWRTA㊄ӥz#Dv ea9!}kum ”9tqPO5/ĜdsT7.0FTpsj:s#'Zʋ[gkI1PI#SYrXV$ ܮWa>F\p?Z=ʅ3ޠnKjM'* aEa:\["o,zeq݇-9YwprA)*vW#VF07gEoSTsҬHys6 ,rsՒ$(23`I#>FDǽ=*8#w#O-%yzpO,Ԯ6ڗW9C횥5ɐ¹bI1h\6"7Pz5\n*w){`;ˢsV$ڕxw$uP1Sssad< 5k k2NR#-.c|}ǯZ/]FqiI&|I$w۞Z{'1Zok,D$A_P=`r@^X|]i}A:*+9~(5okxG{ǨG"H 9&|ylGc?qFkݎ! H`r:/kWS㶕#DC{ⅇwA,Dwj0K Ѻq'>˲|p {bJ6$ zu?V!u$$I)T?ӧ]N窓Z<'UzHQc3"VX#L 1c#9K0X(q-wN<\RT䌶;vv@bg#.Cr98j9(^@Rtgߚvs"QI!a#O6mn\c|\>hT.dMў8"!;P(}#÷+#Ao_Ed64Bd_H!Ԭ)C.A"n*tq{.#`@nFsYM{6#$-_)1>t+.WkIIN)z2ފw:=9N1Y1AwjKTz#{{ImQÓȨFF>1<NGU`pN9M]" qږ>uN1b.$ 0qw9}kZ&:Wku#t+GV%W`U0:kΝP)` 1]OI}4\d'{U*H"\'G5VAW^^߆tN=1Z̍[#W#ghB2ls5҄BC~^+A#nӢ%ˍsԫWhep\+9'jP=ԱlBنI/8rX7؜NdRp~Fύ5dxqn9W?:]>~$3ź@.uM=6 {dL+54ǧ4vȼ;&Yʥ8JBZlh$Hg85 !E ."27$  /An\#?rQ/9Ci&dm=@(rhӍ7sXWƖ ?h *|íxHڷ2(\BGAN-ɔM#$C*d}5>*FAve *')INg93mb'_ R=GOЯ쯢GxJ+ ?OJn34CTo_" %)۴"jN7GLi֛;,V)E'$j2qs uU)Tr=}h ^"1vHFRݛͬg3&PJ5jŔVap[MZeI4c\#ac %u2o L¥Y=Lvߊu{+L%؂qE|1E%ޠxcb3)C;֖\ڑyx!rB =p;ING Wٹ]Z댝Ge0J4-Y'iX[Ge(#j9ʗ'1hKܶ-ELOYD}ۗg SEjgKc'&'(da~w ?:U ʦ>&FV_¼GFfB"J=k^Sk aN Λr6OGFcxD;bjWVybNr3sG#1Y>FKRf<҆VFlpOT_ZCogH+'tukG{?Vv͊t_6~`+OTҥ B<3#}cYCo2\ȅB2#۶ v.*?*chhY6K̿@Ÿ$?.zgWU,Ds29ִ[N9Ty|(:U F;+ Ujcxzv&"na)sE|Ozp<PǮ3RC dwggqLRךbJǵ3V0OcUrl0$i鞝Ng`ѳG4A=d9zg:S8@8&#!AGl23)B{(Г@ rOZc ts99 qJOv dq@ O7ƽ⻻}C5k:2cw$;y’W ]n?i-K2Ф1܍jń Q@['n|LAC}XuRm#) N9'i* 2b@29IE>y=*6%nz M>Ⱥ> H\%VCf`FqVxZLǨэ?M.-rA":<ԧvz=(N>4 (`@8:u춳$Ipw]pZ S׮kZG@`O~k gwCz~E8mSGZT~^ $?HxG ׁ4&Uۻ{)뱧i踽63Z^GCjd|6GsXm>Xp wMIꊴlf%]#gq'}*ƛy$b(rֺk9eswuDqmq8 BδU]L$gK ̑c!溦R9;l v.|5`یpq{ lV1[#;"#1! [QRˡB=Q[osE5hI v2*I`DQ`Vum쬋#w(upzn%Q"USE$ UItf[}4+ch2 6>Q+ld+,2L  繭c'} %iFU/Vv8f bGżn' ׷ wL2sT]A$d՞i't)BT`ُIqnkR0ĐПJ f[*y^OYJrqJ2X?8ivM#!5qdj6gqLc1叐1(R9B{UPV[ሑ.cL`SL{@=]qP\0qCZDDÔd$I^svo bJ>zUxCyAe+O?2zdWnY'"3hm_@vL4: (ʀW$`\#8u=094~aAsGcqqH' (O O x8؞O_NmPX`"NxMcg =zҌ9\gg8Ȥ#?)$sހ:p9u9#{Ӈ # />9 7y {2q9I׸LTm ]PT^=@ۃ?(nr8yqք &0a]E֖[q?^S6/o,WiPA튗_ޒrCs|G8FRFJezX|-|`$>O I+yݶ{(Μ9ebZrJƤ,FKߟdqC(qÞ]DS1 BL!`F/r2Ogm$ƜSfb Bh{ ㎧ZzGw;;ƲIi2y(rY217C= Jڜ#8]~#7^ #bc$5d}qGrwXz`?ªXG!T$WCrFuVaIw\:7sssc>2}! 2`qJ PiH%4@旎rssx)T뎄bN8#v @">(shz(`c<J(Ï H<]܀v).{^Sӡ"G'OLgvɸ#J?t]Jg$A׌jFxڥ}s\б+x92yDmdӊ8 Bه:❀+.Ns#':+j%?(*Ls뚯'$0G'=*2n:YU3H`u$g=p7~MSӊnS%=8P`$@'Rq׷&pA>88W"D888vdd!  rVOi$`WxO wr|:WI}^k;0ig ә@+ӣE#0J#X$U,FEo&4=`$jϭ2 Cc)+T[ق3޺9=r$+?ҡڀ޸$3ЊΤU;TTOX=V,yٓi$FqJGcBzr:Z$z+ [`C.zc#~\!C 7RF{]kb@tu4鞜n>bwM='xŎO)%p9wrFx"nx#{ې9 = 0ҮwtQsWkQUNY?;1\9kgGN1{$$3$v >f/Fbqٸ'bn^;T?h Is4aNѝǎ=ѭ7o ЩҹF@ň}v^iz"\c՘ldy܆=~J)^pOv[J NOJ*JEi.s#'52>9JaҦlٹ99C]nIId7 F1O]K{_ݪ`z"@iG8H\RdsKcրs14s xxJqnqL@ QJFI܀@~1$ЛmJ&+"P6I ʴc?ξX}ƒ W3?+cjzXJ ?Η<%3U2bAq+wd{{c#\{r;{Wrݜ-dE-ǡd)jǙb40SNI]l۲ \l9힦 q׊HZ& Uvqߨ:W]Ћ[@|+Tr:tS㷨'4ۯ*1p=i;ӱzwsۢz~gMuk巴B1Ok%0 =BbOnQ9XSve x>A.Xr+aۜ};ӓG4:oҳ@1uw_Jxq\:ߣ;)HsIjă۸_sqG|€`#Z{V\[cK֡ w',XOMC# udX!݋Br\v5cu#FU?ֶ*o_SaqXBKФb-T*QRGX<񕶑53#Ʌ _A6Z=vmn$;;kpb csW`ḍi15Ө.7h.; 0sH)<)s]hhx2z]y0:o< T7>$ҕLOSTwQ%q$H;sNt)JO3QSgsSSגOHwdZ7_i|&H=1Ԍ Lp:n#?ȠC' P2*H=2yg=(rޙ#8L=ZC+0{U'*~9#ך_sӰ&sw@zQu2$$H< >c=2I9=84>8zBNAE01==h:xG+ۃNV`r#0 gzH)1Ӟ}M (1$)PNTB9X8#5?5c@ê5kWϩiveJqyaq݌ %ykc'3ntcr^*KD2"f`L<%]ZVSnk")+ b@ 3ZN̗uzd{nMh*'8?p)ѤgyKNJ}a9sy>lv2 O"!MJ5dp7t^fpC(אxvgp(${GVbv@ψVΪM"S>xI+HU˹=+/>bHs0`K@λN,,BAkMdt]oj7(zva(imr;n=[X%.!c\Rjffҏbϋ JN9ȨU걞_LKo-3&[7JB[q]ȓ0\ey|eU>k_kBeB0γF ǖx9 =j,vU !'#ط^8dbm38$UsK*?JiXyeAМ~ú{DRvpQIMxrJ(0OZ%m''m h,B))kdî g)E=jױﶙrO?ofbq\MO0ܭU: 8oh5o,@5 JWDng8?Tp6A6aF@@7Oz#/9l>McEsy2Kw y5]w 2( Kę`?K69'sћ[2^b9$UFwuR^#V$e{bʼn@:cڪȆM<y'+OWmO7UF})`n@'=E/=3G'7 (8$(N{JO>##F(pW>.pr9,_TM?]HWm_ŪIĎd%Ijں)_ţBJ8nHg=^% jUAT/>GPK3,|c ̈тkzslxX<le$sڥS~ '?JF֨vţsօˏyqfPaךќrJJ˝8ś㻩niy =C& 2Zv"QB Qi+f/en_iŌr+a1c&ִ L*_[tkJ9U_Ʈ^QzqeF 5j# c q38 GjqCahp}x=*ܖAgU9=]>Qq/#Rl5Zyѱ2uSF8&|o0@AT3whARG<,r=#5BOX[N \l+6pCH'#5%GL3rIҢ~Ki7+I+rpMmuaz%+Ԃy銐5 rRiw>`>DԣsSڨ.3Sm"OC~|mH=sZVrH >KE5#]zPE?VL-ŶDgm(2bXaUݖ҈hv|Ѱ9/ƪju,k,%;{i!]ZŻƒyd,*8UfYTw=>) ǭS.9'f۱ao]`Nw.qg}gP9ȭA !Eal.x 𮸮gR] YzboY$h(X88N)p~X,?qu}@#9F}w;NӐzޡCϧAS(8,8T\zzҖnqM|s猎(98VSONבPJ@>5jR@T:⒴2f[oW#={S,Ƃ^Zh@SWC 1eL;#q -="O1-ynJ99+EʆfTݵ;WV_6$|ɻ®ry"ef%?^߅V#0H)md;Ձ '>1Ь\1iJS:(Gʤ;f%4 V-dy͠$H)edgN7K{;W (t`| zTe áԩqq%דƜya4iq"(T2 F<ӟj|* g׏YEH*K\׵]mA(T SZ]=%?1GLVl3=T\$|rFp3j98PqM1dSG$x8XOKhۧVfMPy$Vu؉DenG$T4nfE3"Ub#@,y?8#jQQˆS:èx!8 sB-#J&VQ‘Ge%*2UP\!#Ug]SCѬsgaD`~a3]hR*]1Rjhb/x x[FҩRILG 9e+C2I ]9W)CAy$r{pNG8<Ӱ z 8GZO:X1qގݎ(>g P43sINϰcB99H m,1鞹[nL2{4NSbH8 x=<12O gv;bu,s'\ڕC6pN>nQϧRM)Q('PFA I'nH^y#Id +1(sT)qkN}j-1/rCqc.mQ&kxl$H+$q9z|KB@\dҫخӝu?,jkq?1,y>f L<1qi.$^҂ܯk`x̱F@&0@8I#5/&ll6"#aNT䁭jy/Zi6[L۱4/=HkQᨴڼ-ΰ Bq996W˨]Ak<,(B@^~<)k GKo 5У"|V`HI~?ZgI|(}FO\5oص{5YWT<ȯ6OSn[ O//🽎JW%8=p3+,T&AַWv}{u$WK1;4z] 7`0:Tkvh,a " OBQxPղpka+*r^%Y.|Πz~5֖PC8zԒy`jx|A q2kOq4އߢ/'!ʯe={Չ5ay0xaߥy'UUXNzQsÔ^[8=qҚ÷/Mnze’x#);q\4!l@V\O 1#o|6z,jzϸ?tKjSI #xǿjI,xw?y]m,Kr֍󓹥"kH[L3!E}R|&q_jmTgn 㧵qۂpmzsPWj49 xi-${E8dw_s^kqou${ׅ09 [7n+9'[rI>4#M59$3JG'Sҵ;=^;97r7Fܒ:\<I4ֶrL3H23<4#v O@2pFq; HdbUAQԇ<='Ar`G8K9ך`n=(;6zR\8hg zz(߃@PI rC029=\`=c3Mq*BP0xjE8ӭF'ֺ? V3'An򝪽IlD{Kך<Yh -IyԈk#kvGM6b"pszS2IGO+՟7\e[I?J0x_ [8\9jNݎX k@kZC ~f+X<NÑj-GW[L;\9]63ZZ a`9 bϲ=[Fr<~A\d>6k OG|<tW=jڽٸ_vUp+ۖwO0+uh4lIa#?^;,坎X$^O=zPמ8T)э2"gysh.#v&/Lr~[W|hHGK=A P7X?WKK Nn3үܼ3Dr}>u`I+#T sY df6& 0`1۵yR,=x;1E3F?RK`TZ߼P#a+>=>l+dcjL|04Ě,,_2x=9rXr'oS"!s3\>ae88V4sԤǝ;{`[036)m<WZjI2+ ջ{7Ԑڷ 9AZV在^1NcZt.6ĒH V̰**3G4vIP}0*9e1@+Oݝ񊊲ͼv*'WPn+@85cZ҆6,w+9{gͻp>S+n_=c!\4*~^}-no~ 8g~|~f;.vyʗ|@wzA8A94z7qzhizr/,~es}yTߌ,|B#`wauQgpzA5G*zt3N3GX?xi6A95u z1b@ypG@{=qE.Z)d嘞=1Ex?*Ϸ40*Ha8xSyOH==ޤ''@ 0tYTt^ !°F|yZmoDUcּ"/$yJoW.kTncv'W}])~:WuwZ[UM"kshvqE4xɁqҪ'=4t˦IԢ3f2`gk%$9##'Ҕq}ю(^ȨdQ\ Њy\9#ص#kC˝jJz_S1Yģ+ >< zB)y|9 3Bi5\cTn~"M-h[!Z璿I4I$#~MKSǯթ>ex]FW^>8gց8U/'|{V #P$H `Òp2GNfdl $V@\ *#=령v=9j+0,`Z&Ge@>Ry3(#8s]QIh[GߠPMh;0:c"[#R1]xAG?Z ; _3jKψ1>k 8:"UIv<#<HF1 9F܁nt#Qdz1I”1ϡA9CgN9L\9<@K,/t7[pQһ]';y-^:c\0i `p>NZpؙS֧i>@$)PpG*h33^ 4R!ggtd=+}n5%8g\1IKCxvwn8'.H'{sX?ko"~+qY NҚqkF4#zK H~Qܷځ/BG~LƤg>{pih^gi:Ƃ  C=A h8r=f9cҁW-::HyI4pܑ3ۚV8Ờ(i#ފNH ϶>x^V@CFJz:czֻO4qw2ǯ'3<+*VEi}3^ǽyg#ֽR.vA-uCEGh@፽=T ֥7VpI-ĥ/(@E89mc ǚJjZ_1,Y`@:cx*'5xℚ+C<_~8ޮTSO1*\ЗWs a:'q,sjOFprA 4Z~) ϯ@zӿ=0q47o4g8TsßgH$ӰC ӨB}q(A>SmʞW=h-@'8RE)y7CFX {R 9 FGǩsn=zJ[ B+N={ Hܜ8\IgYFvȤr5.6NXkvE9f\cHq<6qVf̖!F7őOZt-Yu 41@kܮml zHXr1|q_x/œiwA3,pb)uGl JÊ378ǭ_.;Q8 $3\;6llb]~F9k^#2 +hș?{+@il!rێs_ˁEycg#$K@rIZ#p .T㊹'b6T8L⥴8HDf\vNO]!#9 /Gq9Zm'9h$w_L9P+qdQss5V< ~ |WgQokl:NWpV1:ʽ^*`8_[zT6劌Wn}淊p:wCK\15}"fUR#uFr޴f;Bys|9}jUm9{SBdI$N}k&51AF![W^57=z7Ӝ'u8}:0Tn}+V` /p?A]OᕲmޢLcobxO JjԝsIKtӓ+$ќ{}ם$Ҍ9Q?tA=Mtnۃ|Ҟx<4pNJ0u=~PsiHfb'8zRaTtM8*A~#@w3G `g.ь'ԓCdqzPdHz )~;P6=q4ߨ=zQm<ހ8ڜ=׾)'9sȠTdpvN:R9'tdI1)$@_%H#8)[q$x7$1uz4{e{e `y Z#Y+w =b?ԨcmP1,|Yk4yh'AGZ!Vʅgߚs{jp#|k|.[q 9 ^gޤ1#v~c? [TJ>T\0$rX 9 RJe,$`>FD SssZtLorcp6@mثGf8$blp{q[ZU d$p=dfdj&*w+l$ש>!lI5)2H91/><Ŏ(}VpH:{ר*(:Wm6Wg݉]bқ*1U087N<bz5֎ ?'3^ZnB|w|| cVNJuCRԛ˞d1 :} ($8&@敗ڜF >i03s2 8G~f: w&YpNFs׿JzH⛓Fq*I`y=pJ@ =F=1coRy)̮@=A<[؞1N&SpGZcTg#/C'={9HrIٷ8N?5u) 3]s8g|"^ ye+5ܼg6 y?f64Vu,%Fq ?NfnKX%6m'h$4p1O (jv \:x=yZA%rGrz#w)S۽)3L_nfO8Nz'8Rz ?v! 3:u3?J(`cfwq<zGJOQqG9^`NQu&Hsq7`؞4^89^iiFq89  :f`?i3ރmH>O'H&_j~,)@ f3ׁ\99^F0*hNbn+ ƛبϳm)X1lAQ\w5OOazL\Խ-X5xmiIH Y j)ԕ9sD*g; I1h\''8z.s~u$7N@Q/?|gdH|ЂqyR=u(Hek)€˚KdBp[pg"OG@\$2+ J.ZჵϧU5Q,e;8[tO0F!,9'x'T6nW>HnӐ>)h~cxYƢk(:Ԑp1ڽ_7Ѽk2}G>N$Ў?+z8nW#e8)iSK]g!^g> 3QoՏ^!ϥ|ăZ񉵁[2jsǯ&Cy:~?p=t 3נCFi1$󊔀T&x84>#O8 1 W=pg2p8&=G&#'ӽAMg#"ᘱ'4O:M(4r\ @x#s aGv\w8ɤ'kps~ ƌ8*2FIucAz^HZNNW$98'AR@z4a'*%݈du=tpH'HI߭ ;9y?{j\5A8叮 5}H9!NŒ`R3c :[W3Q?z^Al~UVWjvשJ w~#UtRXu r W >x߸pP0; Q]L:8GKVKfٵ8銷}UG|fA;l}5C QcU,yk޼sCxm7*7 g85e$( k59ctҕ4B ,pRHr *;IԊ敗0s›œ<Aa$wROlqA=jƒ#g=x玔L7<c?M 3cJ@.Fzd21(g=N~c@8F(n4Og#npGCHG /8{Jqazҗ{:I\QҎ `($d)HX`vpvׁJ{@xqx_NiS F>+Y,G̠y^.?#_6|mPpQA1$S[~4gPy u qzn!_>0ǡ6ƊSL-Gnqt99IƒBQp\[qճWlګ/^G8+o>)xw6!# av7YVU"EB~x1Xwp<Ӝ^hs2E  jޓZ\d+OH4n9^03EsO_,qy''4zCHCV=OaLǷ>iFq v鎄T33҅$`_j:cHpxC'hQQJ~9Au鞕C"nz`Js4 s<چvh{8gXm3؛ fC*>X 52+".۫qkZmRKqvps@3s\ L Fx搎2~z+Hy\qIo^AcJrs809֘$# <8H q4 9$=~>!?7^SY )@ wM9 H Q0:`փ7@89Fxހw@gsАqтs ϡ$8NG\u\RzJP9 LӎXd4@3FF뜎dý)L7N):`S{c2hSAK-.$popHt 4{%5`W gz\_|sө ǒqҝ6;/M>1^G4?O@8<7;9X#=!39>Ɣs8=i:gFv6'h 00:ԊCyҡ9oSR) N Ұ ~R3N;A xI@88p>xZ̗%xsĀ)e?wk'<#?0=`i0b{˗'#xI:G6UʸI48e @9Gp@=7oR;FsK3ҔǧӊFrH'קjvsߏjbNԘ$Ƕ(9\8;dL[aGCKFpG{sAH:M ͻ\ 1&|:c"TʘH9c='ͅ^3RrI:EwϷ4; xz(,7çn(b hj<3#0`s@t FOϴ;@y}1FϐmĊGthea2JzdpzHU'> 2Q@=gtss67`ژdӺ}C@ުFGhomR-Yl١lrE{qҒy4zI'q 5d8{YO'ke+|ڼ1+{hb^_s^G"M}v[9?j֒ɨI^y<8g@s8;CژAl9 G$F2r@u s~v鞂OCgGNu<eInG)~0 *r}FO4I˞O$ xcF1Hq댞Ԁ8G"q@ 9$񎸤83 ?3OS7o90=NR`zTn(Coc :zb :*8}iq@#l[(5HJ26$8A@2ݲx҅88OAzJp OzCg׽s8B@rc8^>9 zP!6sMz98}@jN}C(>'h'^3֨c c{ːISކ!˓dޕ=N2r 9A m${E cnNxB|h^}rÞG?hQ=qHbQӸN8r(P[d i pO͜h1qN@ڸ#&@䌎CЎ9@ Ӂnrך^YQ t'40r{uIӃ׮A'߮)~ڤb-799y)Aװ0G Ec򢀹+TsXv8*yH^UT_9γNY:_BL-L_92]Xu9/D#ziۺ1#x<I4y9ҺNAOsi?8f9h\cgN4 g$*`瞹ɠcp󃎔qand|9EyidVU$1LT<|O#zhY;E\d80GW'&>Q0=)@Ad}'&~3fF$ۼΧ5I 8p˽ Nnwu(LJUK^"p@b`8tTvCm~>\aPrɶ3Ӛw=P ׬]?Λdq+F1dᗋ 7#\FxwSҵ "mY-e*+c#Ţ1ړy|qz1 Or9'&8s)[$1O|@Vs4H\cQ9H2:Pp|rqqP pN1bMi#CHYYWG<_L{#GLWR둃W|= j>nlѤ$ _%i{V#IWr9|g5/]Ij*ZH< Myo|N҆P{X<1 gp!s@ޭˮjz>:@0큁[0e1e/ ?\m tyvTih!$:{PS3?0u@>FcшQM$gisJ@''NGM?:LcI1Ǧ{Ԁ0A,88悤9&^$)1䞀Pz38?L8@ pǁIܞ>}17iw'+Zo 60-;|O^_Wai8ZW[F&U#sJ>U=ܞ Ո'o93J<1Tg['}hE}J3cn52g?Ƌ k:nݘK &6_5H0N0GCϵ;S Vэ? x{A?ʼo$6 9EZ78.AF)9=TlsJvcns O8"9ӽ+8 Nڐ<1#h#ԋD9y 4ڞ)"ʎI rN?;-^8<1ĒW_$7Al.pp]sup)2Mۺz{`q= 68֑BgE TqJ1{\I8SAo-ĢP'KV8EVDg9^x㊎B|b+Ծ|$7Z#v!Nh4#Qt'%sH99^E9#koFWlv'Z<3|Mմimg8~ ãGc޳m­ʓIhqG%7&T;;JUo$rUät'3t#ZC;{‘i7o+4I>!l rڷt+{9lڳ18$&*6O; Uqj|wx= j ''ʵt xz}6M[XKI mlp}WU*ʤQӋɫ{Ydyp =7⽫ Ɠ}+k%sq/(H~Z5;;uuhc rUU1kF5R;;R ֍);sJ=G8{Q^z[  S$v#=1IԨs@1ך0-=i@L8=qMO}dvL~tP`J^vi-T37X=^GCykVR??ƼpuǛHm ۨIc\Bp ڽ jѱvS'{sH#ސ隳!8993F @'|v0/Vۜ8RC1U 9$s}EK+ux#Pp,T@q܌w\8~;Rgg;ÖKKLq[prkMk3z][`u)w_\"I1bPp3浓Ǟ [o&ˉy'cf$֓k N1ЌIŽx$jD u)W2܎E `֜y4r おUb'zdJd)C֔9'؃Q8-dg=5ꅾg n A?i>Ll~xSpc 5+'Um8u=7],ڒ]qrIGdtR=qKừN!kl d˴z;s^YqtWr\>K18 t–3L29E5er+J{Ü9'p8}iXoq*p1/M.! $gIRnGJ7'# M 1i Խ!۞s>ׇ5}& ֹg,U!~yL0'HTInGTv8#Aau~K[-`O@O8 qL'47O)i[,xӞ*VSFPYP+➣}P}>[{;9|Fـ N r:Hԣ{[7 ]GacTaӓ]\Dx/$lxiO Jtla6 evs\dRɦ2ȊAUqU`HPF5,I'5EumkBxwr  ՘UE=VQb$Upe:$WV,6.NFWJ55a㊤z3,`RKONF>^%PV}ow7Rʉ'$s׊t|'ʶoMR:kCyns٧gzZoA1N9h'42xZSajo0h=8'91_$qໆ8EIEPDVO#ך! pr}I˼y *O^y6=<3Ɠ (shT=OGhN1ސ4ǵ0SYRҍ.2GqZPnG2B*1:;xN%C/nRyx^!${⌈@A\瑓\b @r8.A{g;r&~3 VPZ-yPe #II-b?/§{W<7OvʼnI9(i45&S{Lt++QVsrI1i $؞R|KS]P[@6_:[Ϡ)rmIYwwfLj渖贲]hsӷZuD~~ b):N{uңñұҲs_QCOrdGRnJ6ðvh7qcsdoY-B`aœʲA} y:=9 砦*NU%7v񎆜qG9B#&O#QM L96܎T1"ۉN\qA<_;Ӹ ,B{%x㞤Իbc֩l d}y>ϗv30s4'1S @O8ʜTX8=;Җ = p3;A> qq\Q'4=~o!cF:犌8:zr2g,ۥU@z8')SsJss@$w㩥$z `黷@h==r)9vO\}h'91K߆ϠIzҳnI$2z1@Ö\:vpe@c=24 0A/=hpW\t8w 8郜 r@`w&4dO2!ޘ dxo͐qsRpt4rrAƀ0wgFZC=8q1@ (Vp?:V's֓=iqLF zZOfsFI}4]<Ѐr׏"l mxN7@& N2'ށ f;vIǭ0w Nrd^@sy{c'"'ғsۥ/'IߊOggzw91SӞi8rN7\: Ё@SINh㌟RԘ}Ϯ;ێjش<ּjGBo#sMǏj~ 77#w<($vFўHx4&?=z[;Ҷ[)$BO˜ `"8$`HpqSsd}ӿJLr?4ʹg' :xP`zp~4Lq|tݎtRޙq'~q\@$r:䃚 @iCJx@aOQ##8ݒB^higws8ryG9#11Sa^sHFqAyROQ0lcޔL8M֗09(dGО=3Pڃ0ݨ O^\ zPAzqtִ$HK \4 I QS:GAIѯ.jDBdvxϵP;g5ːRXʓlJJ;!:@; ;1;sV?O,[m["u&&rF3s{qi5\pjR0+TsG^թi-j xǿe-ZM]u8ABj])O"/9*O\P0q=FO##dQp y#ѻ=zў?@%GN3q9mLfgPA?!2O^q9w d~T7ӓ@ z֕(,FpE8="ZSdzox=┶'^s4G'pҼٝ,ʼnPS Z7H#3Criq9O >ؤ&OLA<_L␹i R11@!='#%8q:i0w1@ IfGAR^02 7$cA4 }*HUPQ߁`AHO ޘ:0AsI'${<#=Gך^t8y8@$g_OqKs;{Rnlh=N?4׎2iFqzS rLAl'Ӛ g~4u N0{㨠' 'wd=Ȧ18<;ҮAG_H`>q1PzztPK('jwxZF8l4bys )˒G^hpK`N:PBLdzR` qA#(ȮWopF *X\}խ28 :dAzn'F{5NWy&mr˅bO<=ppIjƻ>?6g$V^G#AFUc\y{k2ꖒMcEa Pd$6JGNaO[\e=lBpnmssq\C]Q6̖.Nn;JܓTyڼ/8`y3;ٚ=>w[v,1fBHϩF'v{ץHr 9L3ׁMM; 8frOLwtN[<@Bu<ޟ@#㩧g֓ (#ߵ{ 4msadapter/testing/ut/torchvision/assets/masks.tiff000066400000000000000000012603441436713554100230270ustar00rootroot00000000000000II*(1 ]{"shape": [7, 224, 224]}tifffile.py ]](] ]](L^ M]](^ ]](_ ]](>` ]](msadapter/testing/ut/torchvision/common_utils.py000066400000000000000000000153751436713554100226200ustar00rootroot00000000000000import contextlib import functools import os import random import shutil import tempfile import numpy as np import ms_adapter.pytorch as torch from ms_adapter.pytorch.tensor import cast_to_adapter_tensor from PIL import Image from ms_adapter.torchvision import io import __main__ # noqa: 401 IN_CIRCLE_CI = os.getenv("CIRCLECI", False) == "true" IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1" CUDA_NOT_AVAILABLE_MSG = "CUDA device not available" CIRCLECI_GPU_NO_CUDA_MSG = "We're in a CircleCI GPU machine, and this test doesn't need cuda." @contextlib.contextmanager def get_tmp_dir(src=None, **kwargs): tmp_dir = tempfile.mkdtemp(**kwargs) if src is not None: os.rmdir(tmp_dir) shutil.copytree(src, tmp_dir) try: yield tmp_dir finally: shutil.rmtree(tmp_dir) # def set_rng_seed(seed): # torch.manual_seed(seed) # random.seed(seed) class MapNestedTensorObjectImpl: def __init__(self, tensor_map_fn): self.tensor_map_fn = tensor_map_fn def __call__(self, object): if isinstance(object, torch.Tensor): return self.tensor_map_fn(object) elif isinstance(object, dict): mapped_dict = {} for key, value in object.items(): mapped_dict[self(key)] = self(value) return mapped_dict elif isinstance(object, (list, tuple)): mapped_iter = [] for iter in object: mapped_iter.append(self(iter)) return mapped_iter if not isinstance(object, tuple) else tuple(mapped_iter) else: return object def map_nested_tensor_object(object, tensor_map_fn): impl = MapNestedTensorObjectImpl(tensor_map_fn) return impl(object) def is_iterable(obj): try: iter(obj) return True except TypeError: return False # @contextlib.contextmanager # def freeze_rng_state(): # rng_state = torch.get_rng_state() # if torch.cuda.is_available(): # cuda_rng_state = torch.cuda.get_rng_state() # yield # if torch.cuda.is_available(): # torch.cuda.set_rng_state(cuda_rng_state) # torch.set_rng_state(rng_state) def cycle_over(objs): for idx, obj1 in enumerate(objs): for obj2 in objs[:idx] + objs[idx + 1 :]: yield obj1, obj2 def int_dtypes(): return (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64) def float_dtypes(): return (torch.float32, torch.float64) @contextlib.contextmanager def disable_console_output(): with contextlib.ExitStack() as stack, open(os.devnull, "w") as devnull: stack.enter_context(contextlib.redirect_stdout(devnull)) stack.enter_context(contextlib.redirect_stderr(devnull)) yield def cpu_and_gpu(): import pytest # noqa # return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda)) return ["cpu"] def needs_cuda(test_func): import pytest # noqa return pytest.mark.needs_cuda(test_func) def _create_data(height=3, width=3, channels=3, device="cpu"): # TODO: When all relevant tests are ported to pytest, turn this into a module-level fixture tensor = torch.randint(0, 256, (channels, height, width), dtype=torch.uint8, device=device) # tensor = torch.randint(0, 256, (channels, height, width), dtype=torch.float32, device=device) data = tensor.permute(1, 2, 0).contiguous().cpu().numpy() mode = "RGB" if channels == 1: mode = "L" data = data[..., 0] pil_img = Image.fromarray(data, mode=mode) return tensor, pil_img def _create_data_batch(height=3, width=3, channels=3, num_samples=4, device="cpu"): # TODO: When all relevant tests are ported to pytest, turn this into a module-level fixture batch_tensor = torch.randint(0, 256, (num_samples, channels, height, width), dtype=torch.uint8, device=device) return batch_tensor # assert_equal = functools.partial(torch.testing.assert_close, rtol=0, atol=0) def get_list_of_videos(tmpdir, num_videos=5, sizes=None, fps=None): names = [] for i in range(num_videos): if sizes is None: size = 5 * (i + 1) else: size = sizes[i] if fps is None: f = 5 else: f = fps[i] data = torch.randint(0, 256, (size, 300, 400, 3), dtype=torch.uint8) name = os.path.join(tmpdir, f"{i}.mp4") names.append(name) io.write_video(name, data, fps=f) return names def _assert_equal_tensor_to_pil(tensor, pil_image, msg=None): np_pil_image = np.array(pil_image) if np_pil_image.ndim == 2: np_pil_image = np_pil_image[:, :, None] pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1))) if msg is None: msg = f"tensor:\n{tensor} \ndid not equal PIL tensor:\n{pil_tensor}" # assert_equal(tensor.cpu(), pil_tensor, msg=msg) assert np.allclose(tensor.numpy(), pil_tensor.numpy()) def _assert_approx_equal_tensor_to_pil( tensor, pil_image, tol=1e-5, msg=None, agg_method="mean", allowed_percentage_diff=None ): # TODO: we could just merge this into _assert_equal_tensor_to_pil np_pil_image = np.array(pil_image) if np_pil_image.ndim == 2: np_pil_image = np_pil_image[:, :, None] pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1))).to(tensor) if allowed_percentage_diff is not None: # Assert that less than a given %age of pixels are different assert (tensor != pil_tensor).to(torch.float).mean() <= allowed_percentage_diff # error value can be mean absolute error, max abs error # Convert to float to avoid underflow when computing absolute difference tensor = tensor.to(torch.float) pil_tensor = pil_tensor.to(torch.float) err = getattr(torch, agg_method)(torch.abs(tensor - pil_tensor)).item() assert err < tol, f"{err} vs {tol}" def _test_fn_on_batch(batch_tensors, fn, scripted_fn_atol=1e-8, **fn_kwargs): transformed_batch = fn(batch_tensors, **fn_kwargs) for i in range(len(batch_tensors)): img_tensor = cast_to_adapter_tensor(batch_tensors[i, ...]) transformed_img = fn(img_tensor, **fn_kwargs) # torch.testing.assert_close(transformed_img, transformed_batch[i, ...], rtol=0, atol=1e-6) assert np.allclose(transformed_img.numpy(), transformed_batch.numpy()[i, ...], rtol=0, atol=1e-6) if scripted_fn_atol >= 0: # scripted_fn = torch.jit.script(fn) scripted_fn = fn # scriptable function test s_transformed_batch = scripted_fn(batch_tensors, **fn_kwargs) # torch.testing.assert_close(transformed_batch, s_transformed_batch, rtol=1e-5, atol=scripted_fn_atol) assert np.allclose(transformed_batch.numpy(), s_transformed_batch.numpy(), rtol=1e-5, atol=scripted_fn_atol) msadapter/testing/ut/torchvision/conftest.py000066400000000000000000000005321436713554100217220ustar00rootroot00000000000000import random import numpy as np import pytest def pytest_configure(config): # register an additional marker (see pytest_collection_modifyitems) config.addinivalue_line("markers", "needs_cuda: mark for tests that rely on a CUDA device") config.addinivalue_line("markers", "dont_collect: mark for tests that should not be collected")msadapter/testing/ut/torchvision/datasets_utils.py000066400000000000000000001106541436713554100231340ustar00rootroot00000000000000import contextlib import functools import importlib import inspect import itertools import os import pathlib import random import shutil import string import struct import tarfile import unittest import unittest.mock import zipfile from collections import defaultdict from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union import numpy as np import PIL import PIL.Image import pytest import ms_adapter.pytorch as torch from ms_adapter.torchvision.datasets import VisionDataset from ms_adapter.torchvision.io import write_video from common_utils import disable_console_output, get_tmp_dir from ms_adapter.torchvision.transforms.functional import get_dimensions __all__ = [ "UsageError", "lazy_importer", "test_all_configs", "DatasetTestCase", "ImageDatasetTestCase", "VideoDatasetTestCase", "create_image_or_video_tensor", "create_image_file", "create_image_folder", "create_video_file", "create_video_folder", "make_tar", "make_zip", "create_random_string", ] class UsageError(Exception): """Should be raised in case an error happens in the setup rather than the test.""" class LazyImporter: r"""Lazy importer for additional dependencies. Some datasets require additional packages that are no direct dependencies of torchvision. Instances of this class provide modules listed in MODULES as attributes. They are only imported when accessed. """ MODULES = ( "av", "lmdb", "pycocotools", "requests", "scipy.io", "scipy.sparse", "h5py", ) def __init__(self): modules = defaultdict(list) for module in self.MODULES: module, *submodules = module.split(".", 1) if submodules: modules[module].append(submodules[0]) else: # This introduces the module so that it is known when we later iterate over the dictionary. modules.__missing__(module) for module, submodules in modules.items(): # We need the quirky 'module=module' and submodules=submodules arguments to the lambda since otherwise the # lookup for these would happen at runtime rather than at definition. Thus, without it, every property # would try to import the last item in 'modules' setattr( type(self), module, property(lambda self, module=module, submodules=submodules: LazyImporter._import(module, submodules)), ) @staticmethod def _import(package, subpackages): try: module = importlib.import_module(package) except ImportError as error: raise UsageError( f"Failed to import module '{package}'. " f"This probably means that the current test case needs '{package}' installed, " f"but it is not a dependency of torchvision. " f"You need to install it manually, for example 'pip install {package}'." ) from error for name in subpackages: importlib.import_module(f".{name}", package=package) return module lazy_importer = LazyImporter() def requires_lazy_imports(*modules): def outer_wrapper(fn): @functools.wraps(fn) def inner_wrapper(*args, **kwargs): for module in modules: getattr(lazy_importer, module.replace(".", "_")) return fn(*args, **kwargs) return inner_wrapper return outer_wrapper def test_all_configs(test): """Decorator to run test against all configurations. Add this as decorator to an arbitrary test to run it against all configurations. This includes :attr:`DatasetTestCase.DEFAULT_CONFIG` and :attr:`DatasetTestCase.ADDITIONAL_CONFIGS`. The current configuration is provided as the first parameter for the test: .. code-block:: @test_all_configs() def test_foo(self, config): pass .. note:: This will try to remove duplicate configurations. During this process it will not not preserve a potential ordering of the configurations or an inner ordering of a configuration. """ def maybe_remove_duplicates(configs): try: return [dict(config_) for config_ in {tuple(sorted(config.items())) for config in configs}] except TypeError: # A TypeError will be raised if a value of any config is not hashable, e.g. a list. In that case duplicate # removal would be a lot more elaborate and we simply bail out. return configs @functools.wraps(test) def wrapper(self): configs = [] if self.DEFAULT_CONFIG is not None: configs.append(self.DEFAULT_CONFIG) if self.ADDITIONAL_CONFIGS is not None: configs.extend(self.ADDITIONAL_CONFIGS) if not configs: configs = [self._KWARG_DEFAULTS.copy()] else: configs = maybe_remove_duplicates(configs) for config in configs: with self.subTest(**config): test(self, config) return wrapper def combinations_grid(**kwargs): """Creates a grid of input combinations. Each element in the returned sequence is a dictionary containing one possible combination as values. Example: >>> combinations_grid(foo=("bar", "baz"), spam=("eggs", "ham")) [ {'foo': 'bar', 'spam': 'eggs'}, {'foo': 'bar', 'spam': 'ham'}, {'foo': 'baz', 'spam': 'eggs'}, {'foo': 'baz', 'spam': 'ham'} ] """ return [dict(zip(kwargs.keys(), values)) for values in itertools.product(*kwargs.values())] class DatasetTestCase(unittest.TestCase): """Abstract base class for all dataset testcases. You have to overwrite the following class attributes: - DATASET_CLASS (torchvision.datasets.VisionDataset): Class of dataset to be tested. - FEATURE_TYPES (Sequence[Any]): Types of the elements returned by index access of the dataset. Instead of providing these manually, you can instead subclass ``ImageDatasetTestCase`` or ``VideoDatasetTestCase```to get a reasonable default, that should work for most cases. Each entry of the sequence may be a tuple, to indicate multiple possible values. Optionally, you can overwrite the following class attributes: - DEFAULT_CONFIG (Dict[str, Any]): Config that will be used by default. If omitted, this defaults to all keyword arguments of the dataset minus ``transform``, ``target_transform``, ``transforms``, and ``download``. Overwrite this if you want to use a default value for a parameter for which the dataset does not provide one. - ADDITIONAL_CONFIGS (Sequence[Dict[str, Any]]): Additional configs that should be tested. Each dictionary can contain an arbitrary combination of dataset parameters that are **not** ``transform``, ``target_transform``, ``transforms``, or ``download``. - REQUIRED_PACKAGES (Iterable[str]): Additional dependencies to use the dataset. If these packages are not available, the tests are skipped. Additionally, you need to overwrite the ``inject_fake_data()`` method that provides the data that the tests rely on. The fake data should resemble the original data as close as necessary, while containing only few examples. During the creation of the dataset check-, download-, and extract-functions from ``torchvision.datasets.utils`` are disabled. Without further configuration, the testcase will test if 1. the dataset raises a :class:`FileNotFoundError` or a :class:`RuntimeError` if the data files are not found or corrupted, 2. the dataset inherits from `torchvision.datasets.VisionDataset`, 3. the dataset can be turned into a string, 4. the feature types of a returned example matches ``FEATURE_TYPES``, 5. the number of examples matches the injected fake data, and 6. the dataset calls ``transform``, ``target_transform``, or ``transforms`` if available when accessing data. Case 3. to 6. are tested against all configurations in ``CONFIGS``. To add dataset-specific tests, create a new method that takes no arguments with ``test_`` as a name prefix: .. code-block:: def test_foo(self): pass If you want to run the test against all configs, add the ``@test_all_configs`` decorator to the definition and accept a single argument: .. code-block:: @test_all_configs def test_bar(self, config): pass Within the test you can use the ``create_dataset()`` method that yields the dataset as well as additional information provided by the ``ìnject_fake_data()`` method: .. code-block:: def test_baz(self): with self.create_dataset() as (dataset, info): pass """ DATASET_CLASS = None FEATURE_TYPES = None DEFAULT_CONFIG = None ADDITIONAL_CONFIGS = None REQUIRED_PACKAGES = None # These keyword arguments are checked by test_transforms in case they are available in DATASET_CLASS. _TRANSFORM_KWARGS = { "transform", "target_transform", "transforms", } # These keyword arguments get a 'special' treatment and should not be set in DEFAULT_CONFIG or ADDITIONAL_CONFIGS. _SPECIAL_KWARGS = { *_TRANSFORM_KWARGS, "download", } # These fields are populated during setupClass() within _populate_private_class_attributes() # This will be a dictionary containing all keyword arguments with their respective default values extracted from # the dataset constructor. _KWARG_DEFAULTS = None # This will be a set of all _SPECIAL_KWARGS that the dataset constructor takes. _HAS_SPECIAL_KWARG = None # These functions are disabled during dataset creation in create_dataset(). _CHECK_FUNCTIONS = { "check_md5", "check_integrity", } _DOWNLOAD_EXTRACT_FUNCTIONS = { "download_url", "download_file_from_google_drive", "extract_archive", "download_and_extract_archive", } def dataset_args(self, tmpdir: str, config: Dict[str, Any]) -> Sequence[Any]: """Define positional arguments passed to the dataset. .. note:: The default behavior is only valid if the dataset to be tested has ``root`` as the only required parameter. Otherwise you need to overwrite this method. Args: tmpdir (str): Path to a temporary directory. For most cases this acts as root directory for the dataset to be created and in turn also for the fake data injected here. config (Dict[str, Any]): Configuration that will be passed to the dataset constructor. It provides at least fields for all dataset parameters with default values. Returns: (Tuple[str]): ``tmpdir`` which corresponds to ``root`` for most datasets. """ return (tmpdir,) def inject_fake_data(self, tmpdir: str, config: Dict[str, Any]) -> Union[int, Dict[str, Any]]: """Inject fake data for dataset into a temporary directory. During the creation of the dataset the download and extract logic is disabled. Thus, the fake data injected here needs to resemble the raw data, i.e. the state of the dataset directly after the files are downloaded and potentially extracted. Args: tmpdir (str): Path to a temporary directory. For most cases this acts as root directory for the dataset to be created and in turn also for the fake data injected here. config (Dict[str, Any]): Configuration that will be passed to the dataset constructor. It provides at least fields for all dataset parameters with default values. Needs to return one of the following: 1. (int): Number of examples in the dataset to be created, or 2. (Dict[str, Any]): Additional information about the injected fake data. Must contain the field ``"num_examples"`` that corresponds to the number of examples in the dataset to be created. """ raise NotImplementedError("You need to provide fake data in order for the tests to run.") @contextlib.contextmanager def create_dataset( self, config: Optional[Dict[str, Any]] = None, inject_fake_data: bool = True, patch_checks: Optional[bool] = None, **kwargs: Any, ): r"""Create the dataset in a temporary directory. The configuration passed to the dataset is populated to contain at least all parameters with default values. For this the following order of precedence is used: 1. Parameters in :attr:`kwargs`. 2. Configuration in :attr:`config`. 3. Configuration in :attr:`~DatasetTestCase.DEFAULT_CONFIG`. 4. Default parameters of the dataset. Args: config (Optional[Dict[str, Any]]): Configuration that will be used to create the dataset. inject_fake_data (bool): If ``True`` (default) inject the fake data with :meth:`.inject_fake_data` before creating the dataset. patch_checks (Optional[bool]): If ``True`` disable integrity check logic while creating the dataset. If omitted defaults to the same value as ``inject_fake_data``. **kwargs (Any): Additional parameters passed to the dataset. These parameters take precedence in case they overlap with ``config``. Yields: dataset (torchvision.dataset.VisionDataset): Dataset. info (Dict[str, Any]): Additional information about the injected fake data. See :meth:`.inject_fake_data` for details. """ if patch_checks is None: patch_checks = inject_fake_data special_kwargs, other_kwargs = self._split_kwargs(kwargs) complete_config = self._KWARG_DEFAULTS.copy() if self.DEFAULT_CONFIG: complete_config.update(self.DEFAULT_CONFIG) if config: complete_config.update(config) if other_kwargs: complete_config.update(other_kwargs) if "download" in self._HAS_SPECIAL_KWARG and special_kwargs.get("download", False): # override download param to False param if its default is truthy special_kwargs["download"] = False patchers = self._patch_download_extract() if patch_checks: patchers.update(self._patch_checks()) with get_tmp_dir() as tmpdir: args = self.dataset_args(tmpdir, complete_config) info = self._inject_fake_data(tmpdir, complete_config) if inject_fake_data else None with self._maybe_apply_patches(patchers), disable_console_output(): dataset = self.DATASET_CLASS(*args, **complete_config, **special_kwargs) yield dataset, info @classmethod def setUpClass(cls): cls._verify_required_public_class_attributes() cls._populate_private_class_attributes() cls._process_optional_public_class_attributes() super().setUpClass() @classmethod def _verify_required_public_class_attributes(cls): if cls.DATASET_CLASS is None: raise UsageError( "The class attribute 'DATASET_CLASS' needs to be overwritten. " "It should contain the class of the dataset to be tested." ) if cls.FEATURE_TYPES is None: raise UsageError( "The class attribute 'FEATURE_TYPES' needs to be overwritten. " "It should contain a sequence of types that the dataset returns when accessed by index." ) @classmethod def _populate_private_class_attributes(cls): defaults = [] for cls_ in cls.DATASET_CLASS.__mro__: if cls_ is VisionDataset: break argspec = inspect.getfullargspec(cls_.__init__) if not argspec.defaults: continue defaults.append( { kwarg: default for kwarg, default in zip(argspec.args[-len(argspec.defaults) :], argspec.defaults) if not kwarg.startswith("_") } ) if not argspec.varkw: break kwarg_defaults = dict() for config in reversed(defaults): kwarg_defaults.update(config) has_special_kwargs = set() for name in cls._SPECIAL_KWARGS: if name not in kwarg_defaults: continue del kwarg_defaults[name] has_special_kwargs.add(name) cls._KWARG_DEFAULTS = kwarg_defaults cls._HAS_SPECIAL_KWARG = has_special_kwargs @classmethod def _process_optional_public_class_attributes(cls): def check_config(config, name): special_kwargs = tuple(f"'{name}'" for name in cls._SPECIAL_KWARGS if name in config) if special_kwargs: raise UsageError( f"{name} contains a value for the parameter(s) {', '.join(special_kwargs)}. " f"These are handled separately by the test case and should not be set here. " f"If you need to test some custom behavior regarding these parameters, " f"you need to write a custom test (*not* test case), e.g. test_custom_transform()." ) if cls.DEFAULT_CONFIG is not None: check_config(cls.DEFAULT_CONFIG, "DEFAULT_CONFIG") if cls.ADDITIONAL_CONFIGS is not None: for idx, config in enumerate(cls.ADDITIONAL_CONFIGS): check_config(config, f"CONFIGS[{idx}]") if cls.REQUIRED_PACKAGES: missing_pkgs = [] for pkg in cls.REQUIRED_PACKAGES: try: importlib.import_module(pkg) except ImportError: missing_pkgs.append(f"'{pkg}'") if missing_pkgs: raise unittest.SkipTest( f"The package(s) {', '.join(missing_pkgs)} are required to load the dataset " f"'{cls.DATASET_CLASS.__name__}', but are not installed." ) def _split_kwargs(self, kwargs): special_kwargs = kwargs.copy() other_kwargs = {key: special_kwargs.pop(key) for key in set(special_kwargs.keys()) - self._SPECIAL_KWARGS} return special_kwargs, other_kwargs def _inject_fake_data(self, tmpdir, config): info = self.inject_fake_data(tmpdir, config) if info is None: raise UsageError( "The method 'inject_fake_data' needs to return at least an integer indicating the number of " "examples for the current configuration." ) elif isinstance(info, int): info = dict(num_examples=info) elif not isinstance(info, dict): raise UsageError( f"The additional information returned by the method 'inject_fake_data' must be either an " f"integer indicating the number of examples for the current configuration or a dictionary with " f"the same content. Got {type(info)} instead." ) elif "num_examples" not in info: raise UsageError( "The information dictionary returned by the method 'inject_fake_data' must contain a " "'num_examples' field that holds the number of examples for the current configuration." ) return info def _patch_download_extract(self): module = inspect.getmodule(self.DATASET_CLASS).__name__ return {unittest.mock.patch(f"{module}.{function}") for function in self._DOWNLOAD_EXTRACT_FUNCTIONS} def _patch_checks(self): module = inspect.getmodule(self.DATASET_CLASS).__name__ return {unittest.mock.patch(f"{module}.{function}", return_value=True) for function in self._CHECK_FUNCTIONS} @contextlib.contextmanager def _maybe_apply_patches(self, patchers): with contextlib.ExitStack() as stack: mocks = {} for patcher in patchers: with contextlib.suppress(AttributeError): mocks[patcher.target] = stack.enter_context(patcher) yield mocks def test_not_found_or_corrupted(self): with pytest.raises((FileNotFoundError, RuntimeError)): with self.create_dataset(inject_fake_data=False): pass def test_smoke(self): with self.create_dataset() as (dataset, _): assert isinstance(dataset, VisionDataset) @test_all_configs def test_str_smoke(self, config): with self.create_dataset(config) as (dataset, _): assert isinstance(str(dataset), str) @test_all_configs def test_feature_types(self, config): with self.create_dataset(config) as (dataset, _): example = dataset[0] if len(self.FEATURE_TYPES) > 1: actual = len(example) expected = len(self.FEATURE_TYPES) assert ( actual == expected ), "The number of the returned features does not match the the number of elements in FEATURE_TYPES: " f"{actual} != {expected}" else: example = (example,) for idx, (feature, expected_feature_type) in enumerate(zip(example, self.FEATURE_TYPES)): with self.subTest(idx=idx): assert isinstance(feature, expected_feature_type) @test_all_configs def test_num_examples(self, config): with self.create_dataset(config) as (dataset, info): assert len(dataset) == info["num_examples"] @test_all_configs def test_transforms(self, config): mock = unittest.mock.Mock(wraps=lambda *args: args[0] if len(args) == 1 else args) for kwarg in self._TRANSFORM_KWARGS: if kwarg not in self._HAS_SPECIAL_KWARG: continue mock.reset_mock() with self.subTest(kwarg=kwarg): with self.create_dataset(config, **{kwarg: mock}) as (dataset, _): dataset[0] mock.assert_called() class ImageDatasetTestCase(DatasetTestCase): """Abstract base class for image dataset testcases. - Overwrites the FEATURE_TYPES class attribute to expect a :class:`PIL.Image.Image` and an integer label. """ FEATURE_TYPES = (PIL.Image.Image, int) @contextlib.contextmanager def create_dataset( self, config: Optional[Dict[str, Any]] = None, inject_fake_data: bool = True, patch_checks: Optional[bool] = None, **kwargs: Any, ): with super().create_dataset( config=config, inject_fake_data=inject_fake_data, patch_checks=patch_checks, **kwargs, ) as (dataset, info): # PIL.Image.open() only loads the image meta data upfront and keeps the file open until the first access # to the pixel data occurs. Trying to delete such a file results in an PermissionError on Windows. Thus, we # force-load opened images. # This problem only occurs during testing since some tests, e.g. DatasetTestCase.test_feature_types open an # image, but never use the underlying data. During normal operation it is reasonable to assume that the # user wants to work with the image he just opened rather than deleting the underlying file. with self._force_load_images(): yield dataset, info @contextlib.contextmanager def _force_load_images(self): open = PIL.Image.open def new(fp, *args, **kwargs): image = open(fp, *args, **kwargs) if isinstance(fp, (str, pathlib.Path)): image.load() return image with unittest.mock.patch("PIL.Image.open", new=new): yield class VideoDatasetTestCase(DatasetTestCase): """Abstract base class for video dataset testcases. - Overwrites the 'FEATURE_TYPES' class attribute to expect two :class:`torch.Tensor` s for the video and audio as well as an integer label. - Overwrites the 'REQUIRED_PACKAGES' class attribute to require PyAV (``av``). - Adds the 'DEFAULT_FRAMES_PER_CLIP' class attribute. If no 'frames_per_clip' is provided by 'inject_fake_data()' and it is the last parameter without a default value in the dataset constructor, the value of the 'DEFAULT_FRAMES_PER_CLIP' class attribute is appended to the output. """ FEATURE_TYPES = (torch.Tensor, torch.Tensor, int) REQUIRED_PACKAGES = ("av",) DEFAULT_FRAMES_PER_CLIP = 1 def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.dataset_args = self._set_default_frames_per_clip(self.dataset_args) def _set_default_frames_per_clip(self, inject_fake_data): argspec = inspect.getfullargspec(self.DATASET_CLASS.__init__) args_without_default = argspec.args[1 : (-len(argspec.defaults) if argspec.defaults else None)] frames_per_clip_last = args_without_default[-1] == "frames_per_clip" @functools.wraps(inject_fake_data) def wrapper(tmpdir, config): args = inject_fake_data(tmpdir, config) if frames_per_clip_last and len(args) == len(args_without_default) - 1: args = (*args, self.DEFAULT_FRAMES_PER_CLIP) return args return wrapper def create_image_or_video_tensor(size: Sequence[int]) -> torch.Tensor: r"""Create a random uint8 tensor. Args: size (Sequence[int]): Size of the tensor. """ return torch.randint(0, 256, size, dtype=torch.uint8) def create_image_file( root: Union[pathlib.Path, str], name: Union[pathlib.Path, str], size: Union[Sequence[int], int] = 10, **kwargs: Any ) -> pathlib.Path: """Create an image file from random data. Args: root (Union[str, pathlib.Path]): Root directory the image file will be placed in. name (Union[str, pathlib.Path]): Name of the image file. size (Union[Sequence[int], int]): Size of the image that represents the ``(num_channels, height, width)``. If scalar, the value is used for the height and width. If not provided, three channels are assumed. kwargs (Any): Additional parameters passed to :meth:`PIL.Image.Image.save`. Returns: pathlib.Path: Path to the created image file. """ if isinstance(size, int): size = (size, size) if len(size) == 2: size = (3, *size) if len(size) != 3: raise UsageError( f"The 'size' argument should either be an int or a sequence of length 2 or 3. Got {len(size)} instead" ) image = create_image_or_video_tensor(size) file = pathlib.Path(root) / name # torch (num_channels x height x width) -> PIL (width x height x num_channels) image = image.permute(2, 1, 0) # For grayscale images PIL doesn't use a channel dimension if image.shape[2] == 1: image = torch.squeeze(image, 2) PIL.Image.fromarray(image.numpy()).save(file, **kwargs) return file def create_image_folder( root: Union[pathlib.Path, str], name: Union[pathlib.Path, str], file_name_fn: Callable[[int], str], num_examples: int, size: Optional[Union[Sequence[int], int, Callable[[int], Union[Sequence[int], int]]]] = None, **kwargs: Any, ) -> List[pathlib.Path]: """Create a folder of random images. Args: root (Union[str, pathlib.Path]): Root directory the image folder will be placed in. name (Union[str, pathlib.Path]): Name of the image folder. file_name_fn (Callable[[int], str]): Should return a file name if called with the file index. num_examples (int): Number of images to create. size (Optional[Union[Sequence[int], int, Callable[[int], Union[Sequence[int], int]]]]): Size of the images. If callable, will be called with the index of the corresponding file. If omitted, a random height and width between 3 and 10 pixels is selected on a per-image basis. kwargs (Any): Additional parameters passed to :func:`create_image_file`. Returns: List[pathlib.Path]: Paths to all created image files. .. seealso:: - :func:`create_image_file` """ if size is None: def size(idx: int) -> Tuple[int, int, int]: num_channels = 3 height, width = torch.randint(3, 11, size=(2,), dtype=torch.int32).tolist() return (num_channels, height, width) root = pathlib.Path(root) / name os.makedirs(root, exist_ok=True) return [ create_image_file(root, file_name_fn(idx), size=size(idx) if callable(size) else size, **kwargs) for idx in range(num_examples) ] def shape_test_for_stereo( left: PIL.Image.Image, right: PIL.Image.Image, disparity: Optional[np.ndarray] = None, valid_mask: Optional[np.ndarray] = None, ): left_dims = get_dimensions(left) right_dims = get_dimensions(right) c, h, w = left_dims # check that left and right are the same size assert left_dims == right_dims assert c == 3 # check that the disparity has the same spatial dimensions # as the input if disparity is not None: assert disparity.ndim == 3 assert disparity.shape == (1, h, w) if valid_mask is not None: # check that valid mask is the same size as the disparity _, dh, dw = disparity.shape mh, mw = valid_mask.shape assert dh == mh assert dw == mw @requires_lazy_imports("av") def create_video_file( root: Union[pathlib.Path, str], name: Union[pathlib.Path, str], size: Union[Sequence[int], int] = (1, 3, 10, 10), fps: float = 25, **kwargs: Any, ) -> pathlib.Path: """Create an video file from random data. Args: root (Union[str, pathlib.Path]): Root directory the video file will be placed in. name (Union[str, pathlib.Path]): Name of the video file. size (Union[Sequence[int], int]): Size of the video that represents the ``(num_frames, num_channels, height, width)``. If scalar, the value is used for the height and width. If not provided, ``num_frames=1`` and ``num_channels=3`` are assumed. fps (float): Frame rate in frames per second. kwargs (Any): Additional parameters passed to :func:`torchvision.io.write_video`. Returns: pathlib.Path: Path to the created image file. Raises: UsageError: If PyAV is not available. """ if isinstance(size, int): size = (size, size) if len(size) == 2: size = (3, *size) if len(size) == 3: size = (1, *size) if len(size) != 4: raise UsageError( f"The 'size' argument should either be an int or a sequence of length 2, 3, or 4. Got {len(size)} instead" ) video = create_image_or_video_tensor(size) file = pathlib.Path(root) / name write_video(str(file), video.permute(0, 2, 3, 1), fps, **kwargs) return file @requires_lazy_imports("av") def create_video_folder( root: Union[str, pathlib.Path], name: Union[str, pathlib.Path], file_name_fn: Callable[[int], str], num_examples: int, size: Optional[Union[Sequence[int], int, Callable[[int], Union[Sequence[int], int]]]] = None, fps=25, **kwargs, ) -> List[pathlib.Path]: """Create a folder of random videos. Args: root (Union[str, pathlib.Path]): Root directory the video folder will be placed in. name (Union[str, pathlib.Path]): Name of the video folder. file_name_fn (Callable[[int], str]): Should return a file name if called with the file index. num_examples (int): Number of videos to create. size (Optional[Union[Sequence[int], int, Callable[[int], Union[Sequence[int], int]]]]): Size of the videos. If callable, will be called with the index of the corresponding file. If omitted, a random even height and width between 4 and 10 pixels is selected on a per-video basis. fps (float): Frame rate in frames per second. kwargs (Any): Additional parameters passed to :func:`create_video_file`. Returns: List[pathlib.Path]: Paths to all created video files. Raises: UsageError: If PyAV is not available. .. seealso:: - :func:`create_video_file` """ if size is None: def size(idx): num_frames = 1 num_channels = 3 # The 'libx264' video codec, which is the default of torchvision.io.write_video, requires the height and # width of the video to be divisible by 2. height, width = (torch.randint(2, 6, size=(2,), dtype=torch.int32) * 2).tolist() return (num_frames, num_channels, height, width) root = pathlib.Path(root) / name os.makedirs(root, exist_ok=True) return [ create_video_file(root, file_name_fn(idx), size=size(idx) if callable(size) else size, **kwargs) for idx in range(num_examples) ] def _split_files_or_dirs(root, *files_or_dirs): files = set() dirs = set() for file_or_dir in files_or_dirs: path = pathlib.Path(file_or_dir) if not path.is_absolute(): path = root / path if path.is_file(): files.add(path) else: dirs.add(path) for sub_file_or_dir in path.glob("**/*"): if sub_file_or_dir.is_file(): files.add(sub_file_or_dir) else: dirs.add(sub_file_or_dir) if root in dirs: dirs.remove(root) return files, dirs def _make_archive(root, name, *files_or_dirs, opener, adder, remove=True): archive = pathlib.Path(root) / name if not files_or_dirs: # We need to invoke `Path.with_suffix("")`, since call only applies to the last suffix if multiple suffixes are # present. For example, `pathlib.Path("foo.tar.gz").with_suffix("")` results in `foo.tar`. file_or_dir = archive for _ in range(len(archive.suffixes)): file_or_dir = file_or_dir.with_suffix("") if file_or_dir.exists(): files_or_dirs = (file_or_dir,) else: raise ValueError("No file or dir provided.") files, dirs = _split_files_or_dirs(root, *files_or_dirs) with opener(archive) as fh: for file in sorted(files): adder(fh, file, file.relative_to(root)) if remove: for file in files: os.remove(file) for dir in dirs: shutil.rmtree(dir, ignore_errors=True) return archive def make_tar(root, name, *files_or_dirs, remove=True, compression=None): # TODO: detect compression from name return _make_archive( root, name, *files_or_dirs, opener=lambda archive: tarfile.open(archive, f"w:{compression}" if compression else "w"), adder=lambda fh, file, relative_file: fh.add(file, arcname=relative_file), remove=remove, ) def make_zip(root, name, *files_or_dirs, remove=True): return _make_archive( root, name, *files_or_dirs, opener=lambda archive: zipfile.ZipFile(archive, "w"), adder=lambda fh, file, relative_file: fh.write(file, arcname=relative_file), remove=remove, ) def create_random_string(length: int, *digits: str) -> str: """Create a random string. Args: length (int): Number of characters in the generated string. *characters (str): Characters to sample from. If omitted defaults to :attr:`string.ascii_lowercase`. """ if not digits: digits = string.ascii_lowercase else: digits = "".join(itertools.chain(*digits)) return "".join(random.choice(digits) for _ in range(length)) def make_fake_pfm_file(h, w, file_name): values = list(range(3 * h * w)) # Note: we pack everything in little endian: -1.0, and "<" content = f"PF \n{w} {h} \n-1.0\n".encode() + struct.pack("<" + "f" * len(values), *values) with open(file_name, "wb") as f: f.write(content) def make_fake_flo_file(h, w, file_name): """Creates a fake flow file in .flo format.""" # Everything needs to be in little Endian according to # https://vision.middlebury.edu/flow/code/flow-code/README.txt values = list(range(2 * h * w)) content = ( struct.pack("<4c", *(c.encode() for c in "PIEH")) + struct.pack("") for image in (None, *images): self._add_image(fh, image, num_captions_per_image) fh.write("") def _add_image(self, fh, image, num_captions_per_image): fh.write("") self._add_image_header(fh, image) fh.write("
    ") self._add_image_captions(fh, num_captions_per_image) fh.write("
") def _add_image_header(self, fh, image=None): if image: url = f"http://www.flickr.com/photos/user/{image.name.split('_')[0]}/" data = f'{url}' else: data = "Image Not Found" fh.write(f"{data}") def _add_image_captions(self, fh, num_captions_per_image): for caption in self._create_captions(num_captions_per_image): fh.write(f"
  • {caption}") def _create_captions(self, num_captions_per_image): return [str(idx) for idx in range(num_captions_per_image)] def test_captions(self): with self.create_dataset() as (dataset, info): _, captions = dataset[0] assert len(captions) == len(info["captions"]) assert all([a == b for a, b in zip(captions, info["captions"])]) class Flickr30kTestCase(Flickr8kTestCase): DATASET_CLASS = datasets.Flickr30k FEATURE_TYPES = (PIL.Image.Image, list) _ANNOTATIONS_FILE = "captions.token" def _image_file_name(self, idx): return f"{idx}.jpg" def _create_annotations_file(self, root, name, images, num_captions_per_image): with open(root / name, "w") as fh: for image, (idx, caption) in itertools.product( images, enumerate(self._create_captions(num_captions_per_image)) ): fh.write(f"{image.name}#{idx}\t{caption}\n") class MNISTTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.MNIST ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(train=(True, False)) _MAGIC_DTYPES = { torch.uint8: 8, torch.int8: 9, torch.int16: 11, torch.int32: 12, torch.float32: 13, torch.float64: 14, } _IMAGES_SIZE = (28, 28) _IMAGES_DTYPE = torch.uint8 _LABELS_SIZE = () _LABELS_DTYPE = torch.uint8 def inject_fake_data(self, tmpdir, config): raw_dir = pathlib.Path(tmpdir) / self.DATASET_CLASS.__name__ / "raw" os.makedirs(raw_dir, exist_ok=True) num_images = self._num_images(config) self._create_binary_file( raw_dir, self._images_file(config), (num_images, *self._IMAGES_SIZE), self._IMAGES_DTYPE ) self._create_binary_file( raw_dir, self._labels_file(config), (num_images, *self._LABELS_SIZE), self._LABELS_DTYPE ) return num_images def _num_images(self, config): return 2 if config["train"] else 1 def _images_file(self, config): return f"{self._prefix(config)}-images-idx3-ubyte" def _labels_file(self, config): return f"{self._prefix(config)}-labels-idx1-ubyte" def _prefix(self, config): return "train" if config["train"] else "t10k" def _create_binary_file(self, root, filename, size, dtype): with open(pathlib.Path(root) / filename, "wb") as fh: for meta in (self._magic(dtype, len(size)), *size): fh.write(self._encode(meta)) # If ever an MNIST variant is added that uses floating point data, this should be adapted. data = torch.randint(0, torch.iinfo(dtype).max + 1, size, dtype=dtype) fh.write(data.numpy().tobytes()) def _magic(self, dtype, dims): return self._MAGIC_DTYPES[dtype] * 256 + dims def _encode(self, v): return torch.tensor(v, dtype=torch.int32).numpy().tobytes()[::-1] class FashionMNISTTestCase(MNISTTestCase): DATASET_CLASS = datasets.FashionMNIST class KMNISTTestCase(MNISTTestCase): DATASET_CLASS = datasets.KMNIST class EMNISTTestCase(MNISTTestCase): DATASET_CLASS = datasets.EMNIST DEFAULT_CONFIG = dict(split="byclass") ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( split=("byclass", "bymerge", "balanced", "letters", "digits", "mnist"), train=(True, False) ) def _prefix(self, config): return f"emnist-{config['split']}-{'train' if config['train'] else 'test'}" class QMNISTTestCase(MNISTTestCase): DATASET_CLASS = datasets.QMNIST ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(what=("train", "test", "test10k", "nist")) _LABELS_SIZE = (8,) _LABELS_DTYPE = torch.int32 def _num_images(self, config): if config["what"] == "nist": return 3 elif config["what"] == "train": return 2 elif config["what"] == "test50k": # The split 'test50k' is defined as the last 50k images beginning at index 10000. Thus, we need to create # more than 10000 images for the dataset to not be empty. Since this takes significantly longer than the # creation of all other splits, this is excluded from the 'ADDITIONAL_CONFIGS' and is tested only once in # 'test_num_examples_test50k'. return 10001 else: return 1 def _labels_file(self, config): return f"{self._prefix(config)}-labels-idx2-int" def _prefix(self, config): if config["what"] == "nist": return "xnist" if config["what"] is None: what = "train" if config["train"] else "test" elif config["what"].startswith("test"): what = "test" else: what = config["what"] return f"qmnist-{what}" def test_num_examples_test50k(self): with self.create_dataset(what="test50k") as (dataset, info): # Since the split 'test50k' selects all images beginning from the index 10000, we subtract the number of # created examples by this. assert len(dataset) == info["num_examples"] - 10000 class DatasetFolderTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.DatasetFolder # The dataset has no fixed return type since it is defined by the loader parameter. For testing, we use a loader # that simply returns the path as type 'str' instead of loading anything. See the 'dataset_args()' method. FEATURE_TYPES = (str, int) _IMAGE_EXTENSIONS = ("jpg", "png") _VIDEO_EXTENSIONS = ("avi", "mp4") _EXTENSIONS = (*_IMAGE_EXTENSIONS, *_VIDEO_EXTENSIONS) # DatasetFolder has two mutually exclusive parameters: 'extensions' and 'is_valid_file'. One of both is required. # We only iterate over different 'extensions' here and handle the tests for 'is_valid_file' in the # 'test_is_valid_file()' method. DEFAULT_CONFIG = dict(extensions=_EXTENSIONS) ADDITIONAL_CONFIGS = ( *datasets_utils.combinations_grid(extensions=[(ext,) for ext in _IMAGE_EXTENSIONS]), dict(extensions=_IMAGE_EXTENSIONS), *datasets_utils.combinations_grid(extensions=[(ext,) for ext in _VIDEO_EXTENSIONS]), dict(extensions=_VIDEO_EXTENSIONS), ) def dataset_args(self, tmpdir, config): return tmpdir, lambda x: x def inject_fake_data(self, tmpdir, config): extensions = config["extensions"] or self._is_valid_file_to_extensions(config["is_valid_file"]) num_examples_total = 0 classes = [] for ext, cls in zip(self._EXTENSIONS, string.ascii_letters): if ext not in extensions: continue create_example_folder = ( datasets_utils.create_image_folder if ext in self._IMAGE_EXTENSIONS else datasets_utils.create_video_folder ) num_examples = torch.randint(1, 3, size=()).item() create_example_folder(tmpdir, cls, lambda idx: self._file_name_fn(cls, ext, idx), num_examples) num_examples_total += num_examples classes.append(cls) return dict(num_examples=num_examples_total, classes=classes) def _file_name_fn(self, cls, ext, idx): return f"{cls}_{idx}.{ext}" def _is_valid_file_to_extensions(self, is_valid_file): return {ext for ext in self._EXTENSIONS if is_valid_file(f"foo.{ext}")} @datasets_utils.test_all_configs def test_is_valid_file(self, config): extensions = config.pop("extensions") # We need to explicitly pass extensions=None here or otherwise it would be filled by the value from the # DEFAULT_CONFIG. with self.create_dataset( config, extensions=None, is_valid_file=lambda file: pathlib.Path(file).suffix[1:] in extensions ) as (dataset, info): assert len(dataset) == info["num_examples"] @datasets_utils.test_all_configs def test_classes(self, config): with self.create_dataset(config) as (dataset, info): assert len(dataset.classes) == len(info["classes"]) assert all([a == b for a, b in zip(dataset.classes, info["classes"])]) class ImageFolderTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.ImageFolder def inject_fake_data(self, tmpdir, config): num_examples_total = 0 classes = ("a", "b") for cls in classes: num_examples = torch.randint(1, 3, size=()).item() num_examples_total += num_examples datasets_utils.create_image_folder(tmpdir, cls, lambda idx: f"{cls}_{idx}.png", num_examples) return dict(num_examples=num_examples_total, classes=classes) @datasets_utils.test_all_configs def test_classes(self, config): with self.create_dataset(config) as (dataset, info): assert len(dataset.classes) == len(info["classes"]) assert all([a == b for a, b in zip(dataset.classes, info["classes"])]) class KittiTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.Kitti FEATURE_TYPES = (PIL.Image.Image, (list, type(None))) # test split returns None as target ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(train=(True, False)) def inject_fake_data(self, tmpdir, config): kitti_dir = os.path.join(tmpdir, "Kitti", "raw") os.makedirs(kitti_dir) split_to_num_examples = { True: 1, False: 2, } # We need to create all folders(training and testing). for is_training in (True, False): num_examples = split_to_num_examples[is_training] datasets_utils.create_image_folder( root=kitti_dir, name=os.path.join("training" if is_training else "testing", "image_2"), file_name_fn=lambda image_idx: f"{image_idx:06d}.png", num_examples=num_examples, ) if is_training: for image_idx in range(num_examples): target_file_dir = os.path.join(kitti_dir, "training", "label_2") os.makedirs(target_file_dir) target_file_name = os.path.join(target_file_dir, f"{image_idx:06d}.txt") target_contents = "Pedestrian 0.00 0 -0.20 712.40 143.00 810.73 307.92 1.89 0.48 1.20 1.84 1.47 8.41 0.01\n" # noqa with open(target_file_name, "w") as target_file: target_file.write(target_contents) return split_to_num_examples[config["train"]] class SvhnTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.SVHN REQUIRED_PACKAGES = ("scipy",) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test", "extra")) def inject_fake_data(self, tmpdir, config): import scipy.io as sio split = config["split"] num_examples = { "train": 2, "test": 3, "extra": 4, }.get(split) file = f"{split}_32x32.mat" images = np.zeros((32, 32, 3, num_examples), dtype=np.uint8) targets = np.zeros((num_examples,), dtype=np.uint8) sio.savemat(os.path.join(tmpdir, file), {"X": images, "y": targets}) return num_examples class Places365TestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.Places365 ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( split=("train-standard", "train-challenge", "val"), small=(False, True), ) _CATEGORIES = "categories_places365.txt" # {split: file} _FILE_LISTS = { "train-standard": "places365_train_standard.txt", "train-challenge": "places365_train_challenge.txt", "val": "places365_val.txt", } # {(split, small): folder_name} _IMAGES = { ("train-standard", False): "data_large_standard", ("train-challenge", False): "data_large_challenge", ("val", False): "val_large", ("train-standard", True): "data_256_standard", ("train-challenge", True): "data_256_challenge", ("val", True): "val_256", } # (class, idx) _CATEGORIES_CONTENT = ( ("/a/airfield", 0), ("/a/apartment_building/outdoor", 8), ("/b/badlands", 30), ) # (file, idx) _FILE_LIST_CONTENT = ( ("Places365_val_00000001.png", 0), *((f"{category}/Places365_train_00000001.png", idx) for category, idx in _CATEGORIES_CONTENT), ) @staticmethod def _make_txt(root, name, seq): file = os.path.join(root, name) with open(file, "w") as fh: for text, idx in seq: fh.write(f"{text} {idx}\n") @staticmethod def _make_categories_txt(root, name): Places365TestCase._make_txt(root, name, Places365TestCase._CATEGORIES_CONTENT) @staticmethod def _make_file_list_txt(root, name): Places365TestCase._make_txt(root, name, Places365TestCase._FILE_LIST_CONTENT) @staticmethod def _make_image(file_name, size): os.makedirs(os.path.dirname(file_name), exist_ok=True) PIL.Image.fromarray(np.zeros((*size, 3), dtype=np.uint8)).save(file_name) @staticmethod def _make_devkit_archive(root, split): Places365TestCase._make_categories_txt(root, Places365TestCase._CATEGORIES) Places365TestCase._make_file_list_txt(root, Places365TestCase._FILE_LISTS[split]) @staticmethod def _make_images_archive(root, split, small): folder_name = Places365TestCase._IMAGES[(split, small)] image_size = (256, 256) if small else (512, random.randint(512, 1024)) files, idcs = zip(*Places365TestCase._FILE_LIST_CONTENT) images = [f.lstrip("/").replace("/", os.sep) for f in files] for image in images: Places365TestCase._make_image(os.path.join(root, folder_name, image), image_size) return [(os.path.join(root, folder_name, image), idx) for image, idx in zip(images, idcs)] def inject_fake_data(self, tmpdir, config): self._make_devkit_archive(tmpdir, config["split"]) return len(self._make_images_archive(tmpdir, config["split"], config["small"])) def test_classes(self): classes = list(map(lambda x: x[0], self._CATEGORIES_CONTENT)) with self.create_dataset() as (dataset, _): assert dataset.classes == classes def test_class_to_idx(self): class_to_idx = dict(self._CATEGORIES_CONTENT) with self.create_dataset() as (dataset, _): assert dataset.class_to_idx == class_to_idx def test_images_download_preexisting(self): with pytest.raises(RuntimeError): with self.create_dataset({"download": True}): pass class INaturalistTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.INaturalist FEATURE_TYPES = (PIL.Image.Image, (int, tuple)) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( target_type=("kingdom", "full", "genus", ["kingdom", "phylum", "class", "order", "family", "genus", "full"]), version=("2021_train",), ) def inject_fake_data(self, tmpdir, config): categories = [ "00000_Akingdom_0phylum_Aclass_Aorder_Afamily_Agenus_Aspecies", "00001_Akingdom_1phylum_Aclass_Border_Afamily_Bgenus_Aspecies", "00002_Akingdom_2phylum_Cclass_Corder_Cfamily_Cgenus_Cspecies", ] num_images_per_category = 3 for category in categories: datasets_utils.create_image_folder( root=os.path.join(tmpdir, config["version"]), name=category, file_name_fn=lambda idx: f"image_{idx + 1:04d}.jpg", num_examples=num_images_per_category, ) return num_images_per_category * len(categories) def test_targets(self): target_types = ["kingdom", "phylum", "class", "order", "family", "genus", "full"] with self.create_dataset(target_type=target_types, version="2021_valid") as (dataset, _): items = [d[1] for d in dataset] for i, item in enumerate(items): assert dataset.category_name("kingdom", item[0]) == "Akingdom" assert dataset.category_name("phylum", item[1]) == f"{i // 3}phylum" assert item[6] == i // 3 class LFWPeopleTestCase(datasets_utils.DatasetTestCase): DATASET_CLASS = datasets.LFWPeople FEATURE_TYPES = (PIL.Image.Image, int) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( split=("10fold", "train", "test"), image_set=("original", "funneled", "deepfunneled") ) _IMAGES_DIR = {"original": "lfw", "funneled": "lfw_funneled", "deepfunneled": "lfw-deepfunneled"} _file_id = {"10fold": "", "train": "DevTrain", "test": "DevTest"} def inject_fake_data(self, tmpdir, config): tmpdir = pathlib.Path(tmpdir) / "lfw-py" os.makedirs(tmpdir, exist_ok=True) return dict( num_examples=self._create_images_dir(tmpdir, self._IMAGES_DIR[config["image_set"]], config["split"]), split=config["split"], ) def _create_images_dir(self, root, idir, split): idir = os.path.join(root, idir) os.makedirs(idir, exist_ok=True) n, flines = (10, ["10\n"]) if split == "10fold" else (1, []) num_examples = 0 names = [] for _ in range(n): num_people = random.randint(2, 5) flines.append(f"{num_people}\n") for i in range(num_people): name = self._create_random_id() no = random.randint(1, 10) flines.append(f"{name}\t{no}\n") names.append(f"{name}\t{no}\n") datasets_utils.create_image_folder(idir, name, lambda n: f"{name}_{n+1:04d}.jpg", no, 250) num_examples += no with open(pathlib.Path(root) / f"people{self._file_id[split]}.txt", "w") as f: f.writelines(flines) with open(pathlib.Path(root) / "lfw-names.txt", "w") as f: f.writelines(sorted(names)) return num_examples def _create_random_id(self): part1 = datasets_utils.create_random_string(random.randint(5, 7)) part2 = datasets_utils.create_random_string(random.randint(4, 7)) return f"{part1}_{part2}" class LFWPairsTestCase(LFWPeopleTestCase): DATASET_CLASS = datasets.LFWPairs FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, int) def _create_images_dir(self, root, idir, split): idir = os.path.join(root, idir) os.makedirs(idir, exist_ok=True) num_pairs = 7 # effectively 7*2*n = 14*n n, self.flines = (10, [f"10\t{num_pairs}"]) if split == "10fold" else (1, [str(num_pairs)]) for _ in range(n): self._inject_pairs(idir, num_pairs, True) self._inject_pairs(idir, num_pairs, False) with open(pathlib.Path(root) / f"pairs{self._file_id[split]}.txt", "w") as f: f.writelines(self.flines) return num_pairs * 2 * n def _inject_pairs(self, root, num_pairs, same): for i in range(num_pairs): name1 = self._create_random_id() name2 = name1 if same else self._create_random_id() no1, no2 = random.randint(1, 100), random.randint(1, 100) if same: self.flines.append(f"\n{name1}\t{no1}\t{no2}") else: self.flines.append(f"\n{name1}\t{no1}\t{name2}\t{no2}") datasets_utils.create_image_folder(root, name1, lambda _: f"{name1}_{no1:04d}.jpg", 1, 250) datasets_utils.create_image_folder(root, name2, lambda _: f"{name2}_{no2:04d}.jpg", 1, 250) class SintelTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.Sintel ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test"), pass_name=("clean", "final", "both")) FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None))) FLOW_H, FLOW_W = 3, 4 def inject_fake_data(self, tmpdir, config): root = pathlib.Path(tmpdir) / "Sintel" num_images_per_scene = 3 if config["split"] == "train" else 4 num_scenes = 2 for split_dir in ("training", "test"): for pass_name in ("clean", "final"): image_root = root / split_dir / pass_name for scene_id in range(num_scenes): scene_dir = image_root / f"scene_{scene_id}" datasets_utils.create_image_folder( image_root, name=str(scene_dir), file_name_fn=lambda image_idx: f"frame_000{image_idx}.png", num_examples=num_images_per_scene, ) flow_root = root / "training" / "flow" for scene_id in range(num_scenes): scene_dir = flow_root / f"scene_{scene_id}" os.makedirs(scene_dir) for i in range(num_images_per_scene - 1): file_name = str(scene_dir / f"frame_000{i}.flo") datasets_utils.make_fake_flo_file(h=self.FLOW_H, w=self.FLOW_W, file_name=file_name) # with e.g. num_images_per_scene = 3, for a single scene with have 3 images # which are frame_0000, frame_0001 and frame_0002 # They will be consecutively paired as (frame_0000, frame_0001), (frame_0001, frame_0002), # that is 3 - 1 = 2 examples. Hence the formula below num_passes = 2 if config["pass_name"] == "both" else 1 num_examples = (num_images_per_scene - 1) * num_scenes * num_passes return num_examples def test_flow(self): # Make sure flow exists for train split, and make sure there are as many flow values as (pairs of) images h, w = self.FLOW_H, self.FLOW_W expected_flow = np.arange(2 * h * w).reshape(h, w, 2).transpose(2, 0, 1) with self.create_dataset(split="train") as (dataset, _): assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list) for _, _, flow in dataset: assert flow.shape == (2, h, w) np.testing.assert_allclose(flow, expected_flow) # Make sure flow is always None for test split with self.create_dataset(split="test") as (dataset, _): assert dataset._image_list and not dataset._flow_list for _, _, flow in dataset: assert flow is None def test_bad_input(self): with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"): with self.create_dataset(split="bad"): pass with pytest.raises(ValueError, match="Unknown value 'bad' for argument pass_name"): with self.create_dataset(pass_name="bad"): pass class KittiFlowTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.KittiFlow ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test")) FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)), (np.ndarray, type(None))) def inject_fake_data(self, tmpdir, config): root = pathlib.Path(tmpdir) / "KittiFlow" num_examples = 2 if config["split"] == "train" else 3 for split_dir in ("training", "testing"): datasets_utils.create_image_folder( root / split_dir, name="image_2", file_name_fn=lambda image_idx: f"{image_idx}_10.png", num_examples=num_examples, ) datasets_utils.create_image_folder( root / split_dir, name="image_2", file_name_fn=lambda image_idx: f"{image_idx}_11.png", num_examples=num_examples, ) # For kitti the ground truth flows are encoded as 16-bits pngs. # create_image_folder() will actually create 8-bits pngs, but it doesn't # matter much: the flow reader will still be able to read the files, it # will just be garbage flow value - but we don't care about that here. datasets_utils.create_image_folder( root / "training", name="flow_occ", file_name_fn=lambda image_idx: f"{image_idx}_10.png", num_examples=num_examples, ) return num_examples def test_flow_and_valid(self): # Make sure flow exists for train split, and make sure there are as many flow values as (pairs of) images # Also assert flow and valid are of the expected shape with self.create_dataset(split="train") as (dataset, _): assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list) for _, _, flow, valid in dataset: two, h, w = flow.shape assert two == 2 assert valid.shape == (h, w) # Make sure flow and valid are always None for test split with self.create_dataset(split="test") as (dataset, _): assert dataset._image_list and not dataset._flow_list for _, _, flow, valid in dataset: assert flow is None assert valid is None def test_bad_input(self): with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"): with self.create_dataset(split="bad"): pass class FlyingChairsTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.FlyingChairs ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "val")) FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None))) FLOW_H, FLOW_W = 3, 4 def _make_split_file(self, root, num_examples): # We create a fake split file here, but users are asked to download the real one from the authors website split_ids = [1] * num_examples["train"] + [2] * num_examples["val"] random.shuffle(split_ids) with open(str(root / "FlyingChairs_train_val.txt"), "w+") as split_file: for split_id in split_ids: split_file.write(f"{split_id}\n") def inject_fake_data(self, tmpdir, config): root = pathlib.Path(tmpdir) / "FlyingChairs" num_examples = {"train": 5, "val": 3} num_examples_total = sum(num_examples.values()) datasets_utils.create_image_folder( # img1 root, name="data", file_name_fn=lambda image_idx: f"00{image_idx}_img1.ppm", num_examples=num_examples_total, ) datasets_utils.create_image_folder( # img2 root, name="data", file_name_fn=lambda image_idx: f"00{image_idx}_img2.ppm", num_examples=num_examples_total, ) for i in range(num_examples_total): file_name = str(root / "data" / f"00{i}_flow.flo") datasets_utils.make_fake_flo_file(h=self.FLOW_H, w=self.FLOW_W, file_name=file_name) self._make_split_file(root, num_examples) return num_examples[config["split"]] @datasets_utils.test_all_configs def test_flow(self, config): # Make sure flow always exists, and make sure there are as many flow values as (pairs of) images # Also make sure the flow is properly decoded h, w = self.FLOW_H, self.FLOW_W expected_flow = np.arange(2 * h * w).reshape(h, w, 2).transpose(2, 0, 1) with self.create_dataset(config=config) as (dataset, _): assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list) for _, _, flow in dataset: assert flow.shape == (2, h, w) np.testing.assert_allclose(flow, expected_flow) class FlyingThings3DTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.FlyingThings3D ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( split=("train", "test"), pass_name=("clean", "final", "both"), camera=("left", "right", "both") ) FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None))) FLOW_H, FLOW_W = 3, 4 def inject_fake_data(self, tmpdir, config): root = pathlib.Path(tmpdir) / "FlyingThings3D" num_images_per_camera = 3 if config["split"] == "train" else 4 passes = ("frames_cleanpass", "frames_finalpass") splits = ("TRAIN", "TEST") letters = ("A", "B", "C") subfolders = ("0000", "0001") cameras = ("left", "right") for pass_name, split, letter, subfolder, camera in itertools.product( passes, splits, letters, subfolders, cameras ): current_folder = root / pass_name / split / letter / subfolder datasets_utils.create_image_folder( current_folder, name=camera, file_name_fn=lambda image_idx: f"00{image_idx}.png", num_examples=num_images_per_camera, ) directions = ("into_future", "into_past") for split, letter, subfolder, direction, camera in itertools.product( splits, letters, subfolders, directions, cameras ): current_folder = root / "optical_flow" / split / letter / subfolder / direction / camera os.makedirs(str(current_folder), exist_ok=True) for i in range(num_images_per_camera): datasets_utils.make_fake_pfm_file(self.FLOW_H, self.FLOW_W, file_name=str(current_folder / f"{i}.pfm")) num_cameras = 2 if config["camera"] == "both" else 1 num_passes = 2 if config["pass_name"] == "both" else 1 num_examples = ( (num_images_per_camera - 1) * num_cameras * len(subfolders) * len(letters) * len(splits) * num_passes ) return num_examples @datasets_utils.test_all_configs def test_flow(self, config): h, w = self.FLOW_H, self.FLOW_W expected_flow = np.arange(3 * h * w).reshape(h, w, 3).transpose(2, 0, 1) expected_flow = np.flip(expected_flow, axis=1) expected_flow = expected_flow[:2, :, :] with self.create_dataset(config=config) as (dataset, _): assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list) for _, _, flow in dataset: assert flow.shape == (2, self.FLOW_H, self.FLOW_W) np.testing.assert_allclose(flow, expected_flow) def test_bad_input(self): with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"): with self.create_dataset(split="bad"): pass with pytest.raises(ValueError, match="Unknown value 'bad' for argument pass_name"): with self.create_dataset(pass_name="bad"): pass with pytest.raises(ValueError, match="Unknown value 'bad' for argument camera"): with self.create_dataset(camera="bad"): pass class HD1KTestCase(KittiFlowTestCase): DATASET_CLASS = datasets.HD1K def inject_fake_data(self, tmpdir, config): root = pathlib.Path(tmpdir) / "hd1k" num_sequences = 4 if config["split"] == "train" else 3 num_examples_per_train_sequence = 3 for seq_idx in range(num_sequences): # Training data datasets_utils.create_image_folder( root / "hd1k_input", name="image_2", file_name_fn=lambda image_idx: f"{seq_idx:06d}_{image_idx}.png", num_examples=num_examples_per_train_sequence, ) datasets_utils.create_image_folder( root / "hd1k_flow_gt", name="flow_occ", file_name_fn=lambda image_idx: f"{seq_idx:06d}_{image_idx}.png", num_examples=num_examples_per_train_sequence, ) # Test data datasets_utils.create_image_folder( root / "hd1k_challenge", name="image_2", file_name_fn=lambda _: f"{seq_idx:06d}_10.png", num_examples=1, ) datasets_utils.create_image_folder( root / "hd1k_challenge", name="image_2", file_name_fn=lambda _: f"{seq_idx:06d}_11.png", num_examples=1, ) num_examples_per_sequence = num_examples_per_train_sequence if config["split"] == "train" else 2 return num_sequences * (num_examples_per_sequence - 1) class EuroSATTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.EuroSAT FEATURE_TYPES = (PIL.Image.Image, int) def inject_fake_data(self, tmpdir, config): data_folder = os.path.join(tmpdir, "eurosat", "2750") os.makedirs(data_folder) num_examples_per_class = 3 classes = ("AnnualCrop", "Forest") for cls in classes: datasets_utils.create_image_folder( root=data_folder, name=cls, file_name_fn=lambda idx: f"{cls}_{idx}.jpg", num_examples=num_examples_per_class, ) return len(classes) * num_examples_per_class class Food101TestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.Food101 FEATURE_TYPES = (PIL.Image.Image, int) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test")) def inject_fake_data(self, tmpdir: str, config): root_folder = pathlib.Path(tmpdir) / "food-101" image_folder = root_folder / "images" meta_folder = root_folder / "meta" image_folder.mkdir(parents=True) meta_folder.mkdir() num_images_per_class = 5 metadata = {} n_samples_per_class = 3 if config["split"] == "train" else 2 sampled_classes = ("apple_pie", "crab_cakes", "gyoza") for cls in sampled_classes: im_fnames = datasets_utils.create_image_folder( image_folder, cls, file_name_fn=lambda idx: f"{idx}.jpg", num_examples=num_images_per_class, ) metadata[cls] = [ "/".join(fname.relative_to(image_folder).with_suffix("").parts) for fname in random.choices(im_fnames, k=n_samples_per_class) ] with open(meta_folder / f"{config['split']}.json", "w") as file: file.write(json.dumps(metadata)) return len(sampled_classes * n_samples_per_class) class FGVCAircraftTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.FGVCAircraft ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( split=("train", "val", "trainval", "test"), annotation_level=("variant", "family", "manufacturer") ) def inject_fake_data(self, tmpdir: str, config): split = config["split"] annotation_level = config["annotation_level"] annotation_level_to_file = { "variant": "variants.txt", "family": "families.txt", "manufacturer": "manufacturers.txt", } root_folder = pathlib.Path(tmpdir) / "fgvc-aircraft-2013b" data_folder = root_folder / "data" classes = ["707-320", "Hawk T1", "Tornado"] num_images_per_class = 5 datasets_utils.create_image_folder( data_folder, "images", file_name_fn=lambda idx: f"{idx}.jpg", num_examples=num_images_per_class * len(classes), ) annotation_file = data_folder / annotation_level_to_file[annotation_level] with open(annotation_file, "w") as file: file.write("\n".join(classes)) num_samples_per_class = 4 if split == "trainval" else 2 images_classes = [] for i in range(len(classes)): images_classes.extend( [ f"{idx} {classes[i]}" for idx in random.sample( range(i * num_images_per_class, (i + 1) * num_images_per_class), num_samples_per_class ) ] ) images_annotation_file = data_folder / f"images_{annotation_level}_{split}.txt" with open(images_annotation_file, "w") as file: file.write("\n".join(images_classes)) return len(classes * num_samples_per_class) class SUN397TestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.SUN397 def inject_fake_data(self, tmpdir: str, config): data_dir = pathlib.Path(tmpdir) / "SUN397" data_dir.mkdir() num_images_per_class = 5 sampled_classes = ("abbey", "airplane_cabin", "airport_terminal") im_paths = [] for cls in sampled_classes: image_folder = data_dir / cls[0] im_paths.extend( datasets_utils.create_image_folder( image_folder, image_folder / cls, file_name_fn=lambda idx: f"sun_{idx}.jpg", num_examples=num_images_per_class, ) ) with open(data_dir / "ClassName.txt", "w") as file: file.writelines("\n".join(f"/{cls[0]}/{cls}" for cls in sampled_classes)) num_samples = len(im_paths) return num_samples class DTDTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.DTD FEATURE_TYPES = (PIL.Image.Image, int) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( split=("train", "test", "val"), # There is no need to test the whole matrix here, since each fold is treated exactly the same partition=(1, 5, 10), ) def inject_fake_data(self, tmpdir: str, config): data_folder = pathlib.Path(tmpdir) / "dtd" / "dtd" num_images_per_class = 3 image_folder = data_folder / "images" image_files = [] for cls in ("banded", "marbled", "zigzagged"): image_files.extend( datasets_utils.create_image_folder( image_folder, cls, file_name_fn=lambda idx: f"{cls}_{idx:04d}.jpg", num_examples=num_images_per_class, ) ) meta_folder = data_folder / "labels" meta_folder.mkdir() image_ids = [str(path.relative_to(path.parents[1])).replace(os.sep, "/") for path in image_files] image_ids_in_config = random.choices(image_ids, k=len(image_files) // 2) with open(meta_folder / f"{config['split']}{config['partition']}.txt", "w") as file: file.write("\n".join(image_ids_in_config) + "\n") return len(image_ids_in_config) class FER2013TestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.FER2013 ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test")) FEATURE_TYPES = (PIL.Image.Image, (int, type(None))) def inject_fake_data(self, tmpdir, config): base_folder = os.path.join(tmpdir, "fer2013") os.makedirs(base_folder) num_samples = 5 with open(os.path.join(base_folder, f"{config['split']}.csv"), "w", newline="") as file: writer = csv.DictWriter( file, fieldnames=("emotion", "pixels") if config["split"] == "train" else ("pixels",), quoting=csv.QUOTE_NONNUMERIC, quotechar='"', ) writer.writeheader() for _ in range(num_samples): row = dict( pixels=" ".join( str(pixel) for pixel in datasets_utils.create_image_or_video_tensor((48, 48)).view(-1).tolist() ) ) if config["split"] == "train": row["emotion"] = str(int(torch.randint(0, 7, ()))) writer.writerow(row) return num_samples class GTSRBTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.GTSRB FEATURE_TYPES = (PIL.Image.Image, int) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test")) def inject_fake_data(self, tmpdir: str, config): root_folder = os.path.join(tmpdir, "gtsrb") os.makedirs(root_folder, exist_ok=True) # Train data train_folder = os.path.join(root_folder, "GTSRB", "Training") os.makedirs(train_folder, exist_ok=True) num_examples = 3 if config["split"] == "train" else 4 classes = ("00000", "00042", "00012") for class_idx in classes: datasets_utils.create_image_folder( train_folder, name=class_idx, file_name_fn=lambda image_idx: f"{class_idx}_{image_idx:05d}.ppm", num_examples=num_examples, ) total_number_of_examples = num_examples * len(classes) # Test data test_folder = os.path.join(root_folder, "GTSRB", "Final_Test", "Images") os.makedirs(test_folder, exist_ok=True) with open(os.path.join(root_folder, "GT-final_test.csv"), "w") as csv_file: csv_file.write("Filename;Width;Height;Roi.X1;Roi.Y1;Roi.X2;Roi.Y2;ClassId\n") for _ in range(total_number_of_examples): image_file = datasets_utils.create_random_string(5, string.digits) + ".ppm" datasets_utils.create_image_file(test_folder, image_file) row = [ image_file, torch.randint(1, 100, size=()).item(), torch.randint(1, 100, size=()).item(), torch.randint(1, 100, size=()).item(), torch.randint(1, 100, size=()).item(), torch.randint(1, 100, size=()).item(), torch.randint(1, 100, size=()).item(), torch.randint(0, 43, size=()).item(), ] csv_file.write(";".join(map(str, row)) + "\n") return total_number_of_examples class CLEVRClassificationTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.CLEVRClassification FEATURE_TYPES = (PIL.Image.Image, (int, type(None))) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "val", "test")) def inject_fake_data(self, tmpdir, config): data_folder = pathlib.Path(tmpdir) / "clevr" / "CLEVR_v1.0" images_folder = data_folder / "images" image_files = datasets_utils.create_image_folder( images_folder, config["split"], lambda idx: f"CLEVR_{config['split']}_{idx:06d}.png", num_examples=5 ) scenes_folder = data_folder / "scenes" scenes_folder.mkdir() if config["split"] != "test": with open(scenes_folder / f"CLEVR_{config['split']}_scenes.json", "w") as file: json.dump( dict( info=dict(), scenes=[ dict(image_filename=image_file.name, objects=[dict()] * int(torch.randint(0, 10, ()))) for image_file in image_files ], ), file, ) return len(image_files) class OxfordIIITPetTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.OxfordIIITPet FEATURE_TYPES = (PIL.Image.Image, (int, PIL.Image.Image, tuple, type(None))) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( split=("trainval", "test"), target_types=("category", "segmentation", ["category", "segmentation"], []), ) def inject_fake_data(self, tmpdir, config): base_folder = os.path.join(tmpdir, "oxford-iiit-pet") classification_anns_meta = ( dict(cls="Abyssinian", label=0, species="cat"), dict(cls="Keeshond", label=18, species="dog"), dict(cls="Yorkshire Terrier", label=37, species="dog"), ) split_and_classification_anns = [ self._meta_to_split_and_classification_ann(meta, idx) for meta, idx in itertools.product(classification_anns_meta, (1, 2, 10)) ] image_ids, *_ = zip(*split_and_classification_anns) image_files = datasets_utils.create_image_folder( base_folder, "images", file_name_fn=lambda idx: f"{image_ids[idx]}.jpg", num_examples=len(image_ids) ) anns_folder = os.path.join(base_folder, "annotations") os.makedirs(anns_folder) split_and_classification_anns_in_split = random.choices(split_and_classification_anns, k=len(image_ids) // 2) with open(os.path.join(anns_folder, f"{config['split']}.txt"), "w", newline="") as file: writer = csv.writer(file, delimiter=" ") for split_and_classification_ann in split_and_classification_anns_in_split: writer.writerow(split_and_classification_ann) segmentation_files = datasets_utils.create_image_folder( anns_folder, "trimaps", file_name_fn=lambda idx: f"{image_ids[idx]}.png", num_examples=len(image_ids) ) # The dataset has some rogue files for path in image_files[:2]: path.with_suffix(".mat").touch() for path in segmentation_files: path.with_name(f".{path.name}").touch() return len(split_and_classification_anns_in_split) def _meta_to_split_and_classification_ann(self, meta, idx): image_id = "_".join( [ *[(str.title if meta["species"] == "cat" else str.lower)(part) for part in meta["cls"].split()], str(idx), ] ) class_id = str(meta["label"] + 1) species = "1" if meta["species"] == "cat" else "2" breed_id = "-1" return (image_id, class_id, species, breed_id) class StanfordCarsTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.StanfordCars REQUIRED_PACKAGES = ("scipy",) ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test")) def inject_fake_data(self, tmpdir, config): import scipy.io as io from numpy.core.records import fromarrays num_examples = {"train": 5, "test": 7}[config["split"]] num_classes = 3 base_folder = pathlib.Path(tmpdir) / "stanford_cars" devkit = base_folder / "devkit" devkit.mkdir(parents=True) if config["split"] == "train": images_folder_name = "cars_train" annotations_mat_path = devkit / "cars_train_annos.mat" else: images_folder_name = "cars_test" annotations_mat_path = base_folder / "cars_test_annos_withlabels.mat" datasets_utils.create_image_folder( root=base_folder, name=images_folder_name, file_name_fn=lambda image_index: f"{image_index:5d}.jpg", num_examples=num_examples, ) classes = np.random.randint(1, num_classes + 1, num_examples, dtype=np.uint8) fnames = [f"{i:5d}.jpg" for i in range(num_examples)] rec_array = fromarrays( [classes, fnames], names=["class", "fname"], ) io.savemat(annotations_mat_path, {"annotations": rec_array}) random_class_names = ["random_name"] * num_classes io.savemat(devkit / "cars_meta.mat", {"class_names": random_class_names}) return num_examples class Country211TestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.Country211 ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "valid", "test")) def inject_fake_data(self, tmpdir: str, config): split_folder = pathlib.Path(tmpdir) / "country211" / config["split"] split_folder.mkdir(parents=True, exist_ok=True) num_examples = { "train": 3, "valid": 4, "test": 5, }[config["split"]] classes = ("AD", "BS", "GR") for cls in classes: datasets_utils.create_image_folder( split_folder, name=cls, file_name_fn=lambda idx: f"{idx}.jpg", num_examples=num_examples, ) return num_examples * len(classes) class Flowers102TestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.Flowers102 ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "val", "test")) REQUIRED_PACKAGES = ("scipy",) def inject_fake_data(self, tmpdir: str, config): base_folder = pathlib.Path(tmpdir) / "flowers-102" num_classes = 3 num_images_per_split = dict(train=5, val=4, test=3) num_images_total = sum(num_images_per_split.values()) datasets_utils.create_image_folder( base_folder, "jpg", file_name_fn=lambda idx: f"image_{idx + 1:05d}.jpg", num_examples=num_images_total, ) label_dict = dict( labels=np.random.randint(1, num_classes + 1, size=(1, num_images_total), dtype=np.uint8), ) datasets_utils.lazy_importer.scipy.io.savemat(str(base_folder / "imagelabels.mat"), label_dict) setid_mat = np.arange(1, num_images_total + 1, dtype=np.uint16) np.random.shuffle(setid_mat) setid_dict = dict( trnid=setid_mat[: num_images_per_split["train"]].reshape(1, -1), valid=setid_mat[num_images_per_split["train"] : -num_images_per_split["test"]].reshape(1, -1), tstid=setid_mat[-num_images_per_split["test"] :].reshape(1, -1), ) datasets_utils.lazy_importer.scipy.io.savemat(str(base_folder / "setid.mat"), setid_dict) return num_images_per_split[config["split"]] class PCAMTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.PCAM ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "val", "test")) REQUIRED_PACKAGES = ("h5py",) def inject_fake_data(self, tmpdir: str, config): base_folder = pathlib.Path(tmpdir) / "pcam" base_folder.mkdir() num_images = {"train": 2, "test": 3, "val": 4}[config["split"]] images_file = datasets.PCAM._FILES[config["split"]]["images"][0] with datasets_utils.lazy_importer.h5py.File(str(base_folder / images_file), "w") as f: f["x"] = np.random.randint(0, 256, size=(num_images, 10, 10, 3), dtype=np.uint8) targets_file = datasets.PCAM._FILES[config["split"]]["targets"][0] with datasets_utils.lazy_importer.h5py.File(str(base_folder / targets_file), "w") as f: f["y"] = np.random.randint(0, 2, size=(num_images, 1, 1, 1), dtype=np.uint8) return num_images class RenderedSST2TestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.RenderedSST2 ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "val", "test")) SPLIT_TO_FOLDER = {"train": "train", "val": "valid", "test": "test"} def inject_fake_data(self, tmpdir: str, config): root_folder = pathlib.Path(tmpdir) / "rendered-sst2" image_folder = root_folder / self.SPLIT_TO_FOLDER[config["split"]] num_images_per_class = {"train": 5, "test": 6, "val": 7} sampled_classes = ["positive", "negative"] for cls in sampled_classes: datasets_utils.create_image_folder( image_folder, cls, file_name_fn=lambda idx: f"{idx}.png", num_examples=num_images_per_class[config["split"]], ) return len(sampled_classes) * num_images_per_class[config["split"]] if __name__ == "__main__": unittest.main() msadapter/testing/ut/torchvision/test_datasets_samplers.py000066400000000000000000000077001436713554100246560ustar00rootroot00000000000000import numpy as np import pytest import ms_adapter.pytorch as torch from common_utils import get_list_of_videos from ms_adapter.torchvision import io from ms_adapter.torchvision.datasets.samplers import DistributedSampler, RandomClipSampler, UniformClipSampler from ms_adapter.torchvision.datasets.video_utils import VideoClips @pytest.mark.skipif(not io.video._av_available(), reason="this test requires av") class TestDatasetsSamplers: def test_random_clip_sampler(self, tmpdir): video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[25, 25, 25]) video_clips = VideoClips(video_list, 5, 5) sampler = RandomClipSampler(video_clips, 3) assert len(sampler) == 3 * 3 indices = torch.tensor(list(iter(sampler))) videos = torch.div(indices, 5, rounding_mode="floor").numpy() v_idxs, count = np.unique(videos, return_counts=True) assert np.allclose(v_idxs, torch.tensor([0, 1, 2]).numpy()) assert np.allclose(count, torch.tensor([3, 3, 3]).numpy()) def test_random_clip_sampler_unequal(self, tmpdir): video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[10, 25, 25]) video_clips = VideoClips(video_list, 5, 5) sampler = RandomClipSampler(video_clips, 3) assert len(sampler) == 2 + 3 + 3 indices = list(iter(sampler)) assert 0 in indices assert 1 in indices # remove elements of the first video, to simplify testing indices.remove(0) indices.remove(1) indices = torch.tensor(indices) - 2 videos = torch.div(indices, 5, rounding_mode="floor").numpy() v_idxs, count = np.unique(videos, return_counts=True) assert np.allclose(v_idxs, torch.tensor([0, 1]).numpy()) assert np.allclose(count, torch.tensor([3, 3]).numpy()) def test_uniform_clip_sampler(self, tmpdir): video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[25, 25, 25]) video_clips = VideoClips(video_list, 5, 5) sampler = UniformClipSampler(video_clips, 3) assert len(sampler) == 3 * 3 indices = torch.tensor(list(iter(sampler))) videos = torch.div(indices, 5, rounding_mode="floor").numpy() v_idxs, count = np.unique(videos, return_counts=True) assert np.allclose(v_idxs, torch.tensor([0, 1, 2]).numpy()) assert np.allclose(count, torch.tensor([3, 3, 3]).numpy()) assert np.allclose(indices.numpy(), torch.tensor([0, 2, 4, 5, 7, 9, 10, 12, 14]).numpy()) def test_uniform_clip_sampler_insufficient_clips(self, tmpdir): video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[10, 25, 25]) video_clips = VideoClips(video_list, 5, 5) sampler = UniformClipSampler(video_clips, 3) assert len(sampler) == 3 * 3 indices = torch.tensor(list(iter(sampler))) assert np.allclose(indices.numpy(), torch.tensor([0, 0, 1, 2, 4, 6, 7, 9, 11]).numpy()) def test_distributed_sampler_and_uniform_clip_sampler(self, tmpdir): video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[25, 25, 25]) video_clips = VideoClips(video_list, 5, 5) clip_sampler = UniformClipSampler(video_clips, 3) distributed_sampler_rank0 = DistributedSampler( clip_sampler, num_replicas=2, rank=0, group_size=3, ) indices = torch.tensor(list(iter(distributed_sampler_rank0))) assert len(distributed_sampler_rank0) == 6 assert np.allclose(indices.numpy(), torch.tensor([0, 2, 4, 10, 12, 14]).numpy()) distributed_sampler_rank1 = DistributedSampler( clip_sampler, num_replicas=2, rank=1, group_size=3, ) indices = torch.tensor(list(iter(distributed_sampler_rank1))) assert len(distributed_sampler_rank1) == 6 assert np.allclose(indices.numpy(), torch.tensor([5, 7, 9, 0, 2, 4]).numpy()) if __name__ == "__main__": pytest.main([__file__])msadapter/testing/ut/torchvision/test_datasets_utils.py000066400000000000000000000211521436713554100241650ustar00rootroot00000000000000import contextlib import gzip import os import pathlib import re import tarfile import zipfile import pytest import ms_adapter.torchvision.datasets.utils as utils from ms_adapter.torchvision.datasets.folder import make_dataset from ms_adapter.torchvision.datasets.utils import _COMPRESSED_FILE_OPENERS def get_file_path_2(*path_components: str) -> str: return os.path.join(*path_components) TEST_FILE = get_file_path_2( os.path.dirname(os.path.abspath(__file__)), "assets", "encode_jpeg", "grace_hopper_517x606.jpg" ) def patch_url_redirection(mocker, redirect_url): class Response: def __init__(self, url): self.url = url @contextlib.contextmanager def patched_opener(*args, **kwargs): yield Response(redirect_url) return mocker.patch("torchvision.datasets.utils.urllib.request.urlopen", side_effect=patched_opener) class TestDatasetsUtils: # def test_get_redirect_url(self, mocker): #TODO pytorch error, too # url = "https://url.org" # expected_redirect_url = "https://redirect.url.org" # # mock = patch_url_redirection(mocker, expected_redirect_url) # # actual = utils._get_redirect_url(url) # assert actual == expected_redirect_url # # assert mock.call_count == 2 # call_args_1, call_args_2 = mock.call_args_list # assert call_args_1[0][0].full_url == url # assert call_args_2[0][0].full_url == expected_redirect_url # # def test_get_redirect_url_max_hops_exceeded(self, mocker): # url = "https://url.org" # redirect_url = "https://redirect.url.org" # # mock = patch_url_redirection(mocker, redirect_url) # # with pytest.raises(RecursionError): # utils._get_redirect_url(url, max_hops=0) # # assert mock.call_count == 1 # assert mock.call_args[0][0].full_url == url def test_check_md5(self): fpath = TEST_FILE correct_md5 = "9c0bb82894bb3af7f7675ef2b3b6dcdc" false_md5 = "" assert utils.check_md5(fpath, correct_md5) assert not utils.check_md5(fpath, false_md5) def test_check_integrity(self): existing_fpath = TEST_FILE nonexisting_fpath = "" correct_md5 = "9c0bb82894bb3af7f7675ef2b3b6dcdc" false_md5 = "" assert utils.check_integrity(existing_fpath, correct_md5) assert not utils.check_integrity(existing_fpath, false_md5) assert utils.check_integrity(existing_fpath) assert not utils.check_integrity(nonexisting_fpath) def test_get_google_drive_file_id(self): url = "https://drive.google.com/file/d/1GO-BHUYRuvzr1Gtp2_fqXRsr9TIeYbhV/view" expected = "1GO-BHUYRuvzr1Gtp2_fqXRsr9TIeYbhV" actual = utils._get_google_drive_file_id(url) assert actual == expected def test_get_google_drive_file_id_invalid_url(self): url = "http://www.vision.caltech.edu/visipedia-data/CUB-200-2011/CUB_200_2011.tgz" assert utils._get_google_drive_file_id(url) is None @pytest.mark.parametrize( "file, expected", [ ("foo.tar.bz2", (".tar.bz2", ".tar", ".bz2")), ("foo.tar.xz", (".tar.xz", ".tar", ".xz")), ("foo.tar", (".tar", ".tar", None)), ("foo.tar.gz", (".tar.gz", ".tar", ".gz")), ("foo.tbz", (".tbz", ".tar", ".bz2")), ("foo.tbz2", (".tbz2", ".tar", ".bz2")), ("foo.tgz", (".tgz", ".tar", ".gz")), ("foo.bz2", (".bz2", None, ".bz2")), ("foo.gz", (".gz", None, ".gz")), ("foo.zip", (".zip", ".zip", None)), ("foo.xz", (".xz", None, ".xz")), ("foo.bar.tar.gz", (".tar.gz", ".tar", ".gz")), ("foo.bar.gz", (".gz", None, ".gz")), ("foo.bar.zip", (".zip", ".zip", None)), ], ) def test_detect_file_type(self, file, expected): assert utils._detect_file_type(file) == expected @pytest.mark.parametrize("file", ["foo", "foo.tar.baz", "foo.bar"]) def test_detect_file_type_incompatible(self, file): # tests detect file type for no extension, unknown compression and unknown partial extension with pytest.raises(RuntimeError): utils._detect_file_type(file) @pytest.mark.parametrize("extension", [".bz2", ".gz", ".xz"]) def test_decompress(self, extension, tmpdir): def create_compressed(root, content="this is the content"): file = os.path.join(root, "file") compressed = f"{file}{extension}" compressed_file_opener = _COMPRESSED_FILE_OPENERS[extension] with compressed_file_opener(compressed, "wb") as fh: fh.write(content.encode()) return compressed, file, content compressed, file, content = create_compressed(tmpdir) utils._decompress(compressed) assert os.path.exists(file) with open(file) as fh: assert fh.read() == content def test_decompress_no_compression(self): with pytest.raises(RuntimeError): utils._decompress("foo.tar") def test_decompress_remove_finished(self, tmpdir): def create_compressed(root, content="this is the content"): file = os.path.join(root, "file") compressed = f"{file}.gz" with gzip.open(compressed, "wb") as fh: fh.write(content.encode()) return compressed, file, content compressed, file, content = create_compressed(tmpdir) utils.extract_archive(compressed, tmpdir, remove_finished=True) assert not os.path.exists(compressed) # @pytest.mark.parametrize("extension", [".gz", ".xz"]) # @pytest.mark.parametrize("remove_finished", [True, False]) # def test_extract_archive_defer_to_decompress(self, extension, remove_finished, mocker): # filename = "foo" # file = f"{filename}{extension}" # # mocked = mocker.patch("torchvision.datasets.utils._decompress") # utils.extract_archive(file, remove_finished=remove_finished) # # mocked.assert_called_once_with(file, filename, remove_finished=remove_finished) def test_extract_zip(self, tmpdir): def create_archive(root, content="this is the content"): file = os.path.join(root, "dst.txt") archive = os.path.join(root, "archive.zip") with zipfile.ZipFile(archive, "w") as zf: zf.writestr(os.path.basename(file), content) return archive, file, content archive, file, content = create_archive(tmpdir) utils.extract_archive(archive, tmpdir) assert os.path.exists(file) with open(file) as fh: assert fh.read() == content @pytest.mark.parametrize( "extension, mode", [(".tar", "w"), (".tar.gz", "w:gz"), (".tgz", "w:gz"), (".tar.xz", "w:xz")] ) def test_extract_tar(self, extension, mode, tmpdir): def create_archive(root, extension, mode, content="this is the content"): src = os.path.join(root, "src.txt") dst = os.path.join(root, "dst.txt") archive = os.path.join(root, f"archive{extension}") with open(src, "w") as fh: fh.write(content) with tarfile.open(archive, mode=mode) as fh: fh.add(src, arcname=os.path.basename(dst)) return archive, dst, content archive, file, content = create_archive(tmpdir, extension, mode) utils.extract_archive(archive, tmpdir) assert os.path.exists(file) with open(file) as fh: assert fh.read() == content def test_verify_str_arg(self): assert "a" == utils.verify_str_arg("a", "arg", ("a",)) pytest.raises(ValueError, utils.verify_str_arg, 0, ("a",), "arg") pytest.raises(ValueError, utils.verify_str_arg, "b", ("a",), "arg") @pytest.mark.parametrize( ("kwargs", "expected_error_msg"), [ (dict(is_valid_file=lambda path: pathlib.Path(path).suffix in {".png", ".jpeg"}), "classes c"), (dict(extensions=".png"), re.escape("classes b, c. Supported extensions are: .png")), (dict(extensions=(".png", ".jpeg")), re.escape("classes c. Supported extensions are: .png, .jpeg")), ], ) def test_make_dataset_no_valid_files(tmpdir, kwargs, expected_error_msg): tmpdir = pathlib.Path(tmpdir) (tmpdir / "a").mkdir() (tmpdir / "a" / "a.png").touch() (tmpdir / "b").mkdir() (tmpdir / "b" / "b.jpeg").touch() (tmpdir / "c").mkdir() (tmpdir / "c" / "c.unknown").touch() with pytest.raises(FileNotFoundError, match=expected_error_msg): make_dataset(str(tmpdir), **kwargs) if __name__ == "__main__": pytest.main([__file__]) msadapter/testing/ut/torchvision/test_datasets_video_utils.py000066400000000000000000000103301436713554100253470ustar00rootroot00000000000000import pytest import numpy as np import ms_adapter.pytorch as torch from common_utils import get_list_of_videos from ms_adapter.torchvision import io from ms_adapter.torchvision.datasets.video_utils import unfold, VideoClips class TestVideo: def test_unfold(self): a = torch.tensor([0,1,2,3,4,5,6]) r = unfold(a, 3, 3, 1) expected = torch.tensor( [ [0, 1, 2], [3, 4, 5], ] ) assert np.allclose(r.numpy(), expected.numpy()) r = unfold(a, 3, 2, 1) expected = torch.tensor([[0, 1, 2], [2, 3, 4], [4, 5, 6]]) assert np.allclose(r.numpy(), expected.numpy()) r = unfold(a, 3, 2, 2) expected = torch.tensor( [ [0, 2, 4], [2, 4, 6], ] ) assert np.allclose(r.numpy(), expected.numpy()) @pytest.mark.skipif(not io.video._av_available(), reason="this test requires av") def test_video_clips(self, tmpdir): video_list = get_list_of_videos(tmpdir, num_videos=3) video_clips = VideoClips(video_list, 5, 5, num_workers=2) assert video_clips.num_clips() == 1 + 2 + 3 for i, (v_idx, c_idx) in enumerate([(0, 0), (1, 0), (1, 1), (2, 0), (2, 1), (2, 2)]): video_idx, clip_idx = video_clips.get_clip_location(i) assert video_idx == v_idx assert clip_idx == c_idx video_clips = VideoClips(video_list, 6, 6) assert video_clips.num_clips() == 0 + 1 + 2 for i, (v_idx, c_idx) in enumerate([(1, 0), (2, 0), (2, 1)]): video_idx, clip_idx = video_clips.get_clip_location(i) assert video_idx == v_idx assert clip_idx == c_idx video_clips = VideoClips(video_list, 6, 1) assert video_clips.num_clips() == 0 + (10 - 6 + 1) + (15 - 6 + 1) for i, v_idx, c_idx in [(0, 1, 0), (4, 1, 4), (5, 2, 0), (6, 2, 1)]: video_idx, clip_idx = video_clips.get_clip_location(i) assert video_idx == v_idx assert clip_idx == c_idx @pytest.mark.skipif(not io.video._av_available(), reason="this test requires av") def test_video_clips_custom_fps(self, tmpdir): video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[12, 12, 12], fps=[3, 4, 6]) num_frames = 4 for fps in [1, 3, 4, 10]: video_clips = VideoClips(video_list, num_frames, num_frames, fps, num_workers=2) for i in range(video_clips.num_clips()): video, audio, info, video_idx = video_clips.get_clip(i) assert video.shape[0] == num_frames assert info["video_fps"] == fps # TODO add tests checking that the content is right def test_compute_clips_for_video(self): video_pts = torch.tensor(np.arange(30)) # case 1: single clip num_frames = 13 orig_fps = 30 duration = float(len(video_pts)) / orig_fps new_fps = 13 clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, orig_fps, new_fps) resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps) assert len(clips) == 1 assert np.allclose(clips.numpy(), idxs.numpy()) assert np.allclose(idxs[0].numpy(), resampled_idxs.numpy()) # case 2: all frames appear only once num_frames = 4 orig_fps = 30 duration = float(len(video_pts)) / orig_fps new_fps = 12 clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, orig_fps, new_fps) resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps) assert len(clips) == 3 assert np.allclose(clips.numpy(), idxs.numpy()) assert np.allclose(idxs.flatten().numpy(), resampled_idxs.numpy()) # case 3: frames aren't enough for a clip num_frames = 32 orig_fps = 30 new_fps = 13 with pytest.warns(UserWarning): clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, orig_fps, new_fps) assert len(clips) == 0 assert len(idxs) == 0 if __name__ == "__main__": pytest.main([__file__]) msadapter/testing/ut/torchvision/test_functional_tensor.py000066400000000000000000001461361436713554100247030ustar00rootroot00000000000000import colorsys import itertools import math import os import re from functools import partial from typing import Sequence import numpy as np import pytest import torch as pytorch import torchvision.transforms.functional as tv_F import ms_adapter.pytorch as torch import ms_adapter.torchvision.transforms as T import ms_adapter.torchvision.transforms.functional as F import ms_adapter.torchvision.transforms.functional_pil as F_pil import ms_adapter.torchvision.transforms.functional_tensor as F_t from ms_adapter.pytorch.tensor import cast_to_adapter_tensor from common_utils import ( cpu_and_gpu, needs_cuda, _create_data, _create_data_batch, _assert_equal_tensor_to_pil, _assert_approx_equal_tensor_to_pil, _test_fn_on_batch, # assert_equal, ) from ms_adapter.torchvision.transforms import InterpolationMode NEAREST, BILINEAR, BICUBIC = InterpolationMode.NEAREST, InterpolationMode.BILINEAR, InterpolationMode.BICUBIC @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("fn", [F.get_image_size, F.get_image_num_channels, F.get_dimensions]) def test_image_sizes(device, fn): # script_F = torch.jit.script(fn) script_F = fn img_tensor, pil_img = _create_data(16, 18, 3, device=device) value_img = fn(img_tensor) value_pil_img = fn(pil_img) assert value_img == value_pil_img value_img_script = script_F(img_tensor) assert value_img == value_img_script batch_tensors = _create_data_batch(16, 18, 3, num_samples=4, device=device) value_img_batch = fn(batch_tensors) assert value_img == value_img_batch class TestRotate: ALL_DTYPES = [None, torch.float32, torch.float64, torch.float16] # scripted_rotate = torch.jit.script(F.rotate) IMG_W = 26 @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("height, width", [(7, 33), (26, IMG_W), (32, IMG_W)]) @pytest.mark.parametrize( "center", [ None, (int(IMG_W * 0.3), int(IMG_W * 0.4)), [int(IMG_W * 0.5), int(IMG_W * 0.6)], ], ) @pytest.mark.parametrize("dt", ALL_DTYPES) @pytest.mark.parametrize("angle", range(-180, 180, 34)) @pytest.mark.parametrize("expand", [True, False]) @pytest.mark.parametrize( "fill", [ None, [0, 0, 0], (1, 2, 3), [255, 255, 255], [ 1, ], (2.0,), ], ) @pytest.mark.parametrize("fn", [F.rotate]) def test_rotate(self, device, height, width, center, dt, angle, expand, fill, fn): tensor, pil_img = _create_data(height, width, device=device) if dt == torch.float16 and torch.device(device).type == "cpu": # skip float16 on CPU case return if dt is not None: tensor = tensor.to(dtype=dt) f_pil = int(fill[0]) if fill is not None and len(fill) == 1 else fill out_pil_img = F.rotate(pil_img, angle=angle, interpolation=NEAREST, expand=expand, center=center, fill=f_pil) out_pil_tensor = torch.tensor(np.array(out_pil_img).transpose((2, 0, 1))) out_tensor = fn(tensor, angle=angle, interpolation=NEAREST, expand=expand, center=center, fill=fill).cpu() if out_tensor.dtype != torch.uint8: out_tensor = out_tensor.to(torch.uint8) assert ( out_tensor.shape == out_pil_tensor.shape ), f"{(height, width, NEAREST, dt, angle, expand, center)}: {out_tensor.shape} vs {out_pil_tensor.shape}" num_diff_pixels = cast_to_adapter_tensor(out_tensor != out_pil_tensor).sum().item() / 3.0 ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2] # Tolerance : less than 3% of different pixels assert ratio_diff_pixels < 0.03, ( f"{(height, width, NEAREST, dt, angle, expand, center, fill)}: " f"{ratio_diff_pixels}\n{out_tensor[0, :7, :7]} vs \n" f"{out_pil_tensor[0, :7, :7]}" ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dt", ALL_DTYPES) def test_rotate_batch(self, device, dt): if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return batch_tensors = _create_data_batch(26, 36, num_samples=4, device=device) if dt is not None: batch_tensors = batch_tensors.to(dtype=dt) center = (20, 22) _test_fn_on_batch(batch_tensors, F.rotate, angle=32, interpolation=NEAREST, expand=True, center=center) def test_rotate_deprecation_resample(self): tensor, _ = _create_data(26, 26) # assert deprecation warning and non-BC with pytest.warns( UserWarning, match=re.escape( "The parameter 'resample' is deprecated since 0.12 and will be removed 0.14. " "Please use 'interpolation' instead." ), ): res1 = F.rotate(tensor, 45, resample=2) res2 = F.rotate(tensor, 45, interpolation=BILINEAR) # assert_equal(res1, res2) assert np.allclose(res1.numpy(), res2.numpy()) def test_rotate_interpolation_type(self): tensor, _ = _create_data(26, 26) # assert changed type warning with pytest.warns( UserWarning, match=re.escape( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ), ): res1 = F.rotate(tensor, 45, interpolation=2) res2 = F.rotate(tensor, 45, interpolation=BILINEAR) # assert_equal(res1, res2) assert np.allclose(res1.numpy(), res2.numpy()) class TestAffine: ALL_DTYPES = [None, torch.float32, torch.float64, torch.float16] # scripted_affine = torch.jit.script(F.affine) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("height, width", [(26, 26)]) @pytest.mark.parametrize("dt", ALL_DTYPES) @pytest.mark.parametrize( "angle, config", [ (90, {"k": 1, "dims": (-1, -2)}), (45, None), (30, None), (-30, None), (-45, None), (-90, {"k": -1, "dims": (-1, -2)}), (180, {"k": 2, "dims": (-1, -2)}), ], ) @pytest.mark.parametrize("fn", [F.affine]) def test_square_rotations(self, device, height, width, dt, angle, config, fn): # 2) Test rotation tensor, pil_img = _create_data(height, width, device=device) if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return if dt is not None: tensor = tensor.to(dtype=dt) out_pil_img = F.affine( pil_img, angle=angle, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], interpolation=NEAREST ) out_pil_tensor = torch.tensor(np.array(out_pil_img).transpose((2, 0, 1))) out_tensor = fn(tensor, angle=angle, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], interpolation=NEAREST) # TODO: # if config is not None: # # assert_equal(torch.rot90(tensor, **config), out_tensor) # assert np.allclose(torch.rot90(tensor, **config).numpy(), out_tensor.numpy()) if out_tensor.dtype != torch.uint8: out_tensor = out_tensor.to(torch.uint8) num_diff_pixels = cast_to_adapter_tensor(out_tensor != out_pil_tensor).sum().item() / 3.0 ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2] # Tolerance : less than 6% of different pixels assert ratio_diff_pixels < 0.06 @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("height, width", [(32, 26)]) @pytest.mark.parametrize("dt", ALL_DTYPES) @pytest.mark.parametrize("angle", [90, 45, 15, -30, -60, -120]) @pytest.mark.parametrize("fn", [F.affine]) @pytest.mark.parametrize("center", [None, [0, 0]]) def test_rect_rotations(self, device, height, width, dt, angle, fn, center): # Tests on rectangular images tensor, pil_img = _create_data(height, width, device=device) if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return if dt is not None: tensor = tensor.to(dtype=dt) out_pil_img = F.affine( pil_img, angle=angle, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], interpolation=NEAREST, center=center ) out_pil_tensor = torch.tensor(np.array(out_pil_img).transpose((2, 0, 1))) out_tensor = fn( tensor, angle=angle, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], interpolation=NEAREST, center=center ).cpu() if out_tensor.dtype != torch.uint8: out_tensor = out_tensor.to(torch.uint8) num_diff_pixels = cast_to_adapter_tensor(out_tensor != out_pil_tensor).sum().item() / 3.0 ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2] # Tolerance : less than 3% of different pixels assert ratio_diff_pixels < 0.03 @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("height, width", [(26, 26), (32, 26)]) @pytest.mark.parametrize("dt", ALL_DTYPES) @pytest.mark.parametrize("t", [[10, 12], (-12, -13)]) @pytest.mark.parametrize("fn", [F.affine]) def test_translations(self, device, height, width, dt, t, fn): # 3) Test translation tensor, pil_img = _create_data(height, width, device=device) if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return if dt is not None: tensor = tensor.to(dtype=dt) out_pil_img = F.affine(pil_img, angle=0, translate=t, scale=1.0, shear=[0.0, 0.0], interpolation=NEAREST) out_tensor = fn(tensor, angle=0, translate=t, scale=1.0, shear=[0.0, 0.0], interpolation=NEAREST) if out_tensor.dtype != torch.uint8: out_tensor = out_tensor.to(torch.uint8) _assert_equal_tensor_to_pil(out_tensor, out_pil_img) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("height, width", [(26, 26), (32, 26)]) @pytest.mark.parametrize("dt", ALL_DTYPES) @pytest.mark.parametrize( "a, t, s, sh, f", [ (45.5, [5, 6], 1.0, [0.0, 0.0], None), (33, (5, -4), 1.0, [0.0, 0.0], [0, 0, 0]), (45, [-5, 4], 1.2, [0.0, 0.0], (1, 2, 3)), (33, (-4, -8), 2.0, [0.0, 0.0], [255, 255, 255]), ( 85, (10, -10), 0.7, [0.0, 0.0], [ 1, ], ), ( 0, [0, 0], 1.0, [ 35.0, ], (2.0,), ), (-25, [0, 0], 1.2, [0.0, 15.0], None), (-45, [-10, 0], 0.7, [2.0, 5.0], None), (-45, [-10, -10], 1.2, [4.0, 5.0], None), (-90, [0, 0], 1.0, [0.0, 0.0], None), ], ) @pytest.mark.parametrize("fn", [F.affine]) def test_all_ops(self, device, height, width, dt, a, t, s, sh, f, fn): # 4) Test rotation + translation + scale + shear tensor, pil_img = _create_data(height, width, device=device) if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return if dt is not None: tensor = tensor.to(dtype=dt) f_pil = int(f[0]) if f is not None and len(f) == 1 else f out_pil_img = F.affine(pil_img, angle=a, translate=t, scale=s, shear=sh, interpolation=NEAREST, fill=f_pil) out_pil_tensor = torch.tensor(np.array(out_pil_img).transpose((2, 0, 1))) out_tensor = fn(tensor, angle=a, translate=t, scale=s, shear=sh, interpolation=NEAREST, fill=f).cpu() if out_tensor.dtype != torch.uint8: out_tensor = out_tensor.to(torch.uint8) num_diff_pixels = cast_to_adapter_tensor(out_tensor != out_pil_tensor).sum().item() / 3.0 ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2] # Tolerance : less than 5% (cpu), 6% (cuda) of different pixels tol = 0.06 if device == "cuda" else 0.05 assert ratio_diff_pixels < tol @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dt", ALL_DTYPES) def test_batches(self, device, dt): if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return batch_tensors = _create_data_batch(26, 36, num_samples=4, device=device) if dt is not None: batch_tensors = batch_tensors.to(dtype=dt) _test_fn_on_batch(batch_tensors, F.affine, angle=-43, translate=[-3, 4], scale=1.2, shear=[4.0, 5.0]) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_warnings(self, device): tensor, pil_img = _create_data(26, 26, device=device) # assert deprecation warning and non-BC with pytest.warns( UserWarning, match=re.escape( "The parameter 'resample' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'interpolation' instead." ), ): res1 = F.affine(tensor, 45, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=2) res2 = F.affine(tensor, 45, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], interpolation=BILINEAR) # assert_equal(res1, res2) assert np.allclose(res1.numpy(), res2.numpy()) # assert changed type warning with pytest.warns( UserWarning, match=re.escape( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ), ): res1 = F.affine(tensor, 45, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], interpolation=2) res2 = F.affine(tensor, 45, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], interpolation=BILINEAR) # assert_equal(res1, res2) assert np.allclose(res1.numpy(), res2.numpy()) with pytest.warns( UserWarning, match=re.escape( "The parameter 'fillcolor' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'fill' instead." ), ): res1 = F.affine(pil_img, 45, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], fillcolor=10) res2 = F.affine(pil_img, 45, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], fill=10) # we convert the PIL images to numpy as assert_equal doesn't work on PIL images. # assert_equal(np.asarray(res1), np.asarray(res2)) assert np.allclose(np.asarray(res1), np.asarray(res2)) def _get_data_dims_and_points_for_perspective(): # Ideally we would parametrize independently over data dims and points, but # we want to tests on some points that also depend on the data dims. # Pytest doesn't support covariant parametrization, so we do it somewhat manually here. data_dims = [(26, 34), (26, 26)] points = [ [[[0, 0], [33, 0], [33, 25], [0, 25]], [[3, 2], [32, 3], [30, 24], [2, 25]]], [[[3, 2], [32, 3], [30, 24], [2, 25]], [[0, 0], [33, 0], [33, 25], [0, 25]]], [[[3, 2], [32, 3], [30, 24], [2, 25]], [[5, 5], [30, 3], [33, 19], [4, 25]]], ] dims_and_points = list(itertools.product(data_dims, points)) # up to here, we could just have used 2 @parametrized. # Down below is the covarariant part as the points depend on the data dims. n = 10 for dim in data_dims: points += [(dim, T.RandomPerspective.get_params(dim[1], dim[0], i / n)) for i in range(n)] return dims_and_points @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dims_and_points", _get_data_dims_and_points_for_perspective()) @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64, torch.float16]) @pytest.mark.parametrize( "fill", ( None, [0, 0, 0], [1, 2, 3], [255, 255, 255], [ 1, ], (2.0,), ), ) @pytest.mark.parametrize("fn", [F.perspective]) def test_perspective_pil_vs_tensor(device, dims_and_points, dt, fill, fn): if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return data_dims, (spoints, epoints) = dims_and_points tensor, pil_img = _create_data(*data_dims, device=device) if dt is not None: tensor = tensor.to(dtype=dt) interpolation = NEAREST fill_pil = int(fill[0]) if fill is not None and len(fill) == 1 else fill out_pil_img = F.perspective( pil_img, startpoints=spoints, endpoints=epoints, interpolation=interpolation, fill=fill_pil ) out_pil_tensor = torch.tensor(np.array(out_pil_img).transpose((2, 0, 1))) out_tensor = fn(tensor, startpoints=spoints, endpoints=epoints, interpolation=interpolation, fill=fill).cpu() if out_tensor.dtype != torch.uint8: out_tensor = out_tensor.to(torch.uint8) num_diff_pixels = cast_to_adapter_tensor(out_tensor != out_pil_tensor).sum().item() / 3.0 ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2] # Tolerance : less than 5% of different pixels assert ratio_diff_pixels < 0.05 @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dims_and_points", _get_data_dims_and_points_for_perspective()) @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64, torch.float16]) def test_perspective_batch(device, dims_and_points, dt): if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return data_dims, (spoints, epoints) = dims_and_points batch_tensors = _create_data_batch(*data_dims, num_samples=4, device=device) if dt is not None: batch_tensors = batch_tensors.to(dtype=dt) # Ignore the equivalence between scripted and regular function on float16 cuda. The pixels at # the border may be entirely different due to small rounding errors. scripted_fn_atol = -1 if (dt == torch.float16 and device == "cuda") else 1e-8 _test_fn_on_batch( batch_tensors, F.perspective, scripted_fn_atol=scripted_fn_atol, startpoints=spoints, endpoints=epoints, interpolation=NEAREST, ) def test_perspective_interpolation_warning(): # assert changed type warning spoints = [[0, 0], [33, 0], [33, 25], [0, 25]] epoints = [[3, 2], [32, 3], [30, 24], [2, 25]] tensor = torch.randint(0, 256, (3, 26, 26)) with pytest.warns( UserWarning, match=re.escape( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ), ): res1 = F.perspective(tensor, startpoints=spoints, endpoints=epoints, interpolation=2) res2 = F.perspective(tensor, startpoints=spoints, endpoints=epoints, interpolation=BILINEAR) # assert_equal(res1, res2) assert np.allclose(res1.numpy(), res2.numpy()) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64, torch.float16]) @pytest.mark.parametrize( "size", [ 32, 26, [ 32, ], [32, 32], (32, 32), [26, 35], ], ) @pytest.mark.parametrize("max_size", [None, 34, 40, 1000]) # @pytest.mark.parametrize("interpolation", [BILINEAR, BICUBIC, NEAREST]) @pytest.mark.parametrize("interpolation", [BILINEAR, NEAREST]) def test_resize(device, dt, size, max_size, interpolation): if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return if max_size is not None and isinstance(size, Sequence) and len(size) != 1: return # unsupported torch.manual_seed(12) # script_fn = torch.jit.script(F.resize) tensor, pil_img = _create_data(26, 36, device=device) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) if dt is not None: # This is a trivial cast to float of uint8 data to test all cases tensor = tensor.to(dt) batch_tensors = batch_tensors.to(dt) resized_tensor = F.resize(tensor, size=size, interpolation=interpolation, max_size=max_size) resized_pil_img = F.resize(pil_img, size=size, interpolation=interpolation, max_size=max_size) assert resized_tensor.size()[1:] == resized_pil_img.size[::-1] if interpolation not in [ NEAREST, ]: # We can not check values if mode = NEAREST, as results are different # E.g. resized_tensor = [[a, a, b, c, d, d, e, ...]] # E.g. resized_pil_img = [[a, b, c, c, d, e, f, ...]] resized_tensor_f = resized_tensor # we need to cast to uint8 to compare with PIL image if resized_tensor_f.dtype == torch.uint8: resized_tensor_f = resized_tensor_f.to(torch.float) # Pay attention to high tolerance for MAE _assert_approx_equal_tensor_to_pil(resized_tensor_f, resized_pil_img, tol=8.0) if isinstance(size, int): script_size = [ size, ] else: script_size = size # resize_result = script_fn(tensor, size=script_size, interpolation=interpolation, max_size=max_size) # assert_equal(resized_tensor, resize_result) _test_fn_on_batch(batch_tensors, F.resize, size=script_size, interpolation=interpolation, max_size=max_size) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_resize_asserts(device): tensor, pil_img = _create_data(26, 36, device=device) # assert changed type warning with pytest.warns( UserWarning, match=re.escape( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ), ): res1 = F.resize(tensor, size=32, interpolation=2) res2 = F.resize(tensor, size=32, interpolation=BILINEAR) # assert_equal(res1, res2) assert np.allclose(res1.numpy(), res2.numpy()) for img in (tensor, pil_img): exp_msg = "max_size should only be passed if size specifies the length of the smaller edge" with pytest.raises(ValueError, match=exp_msg): F.resize(img, size=(32, 34), max_size=35) with pytest.raises(ValueError, match="max_size = 32 must be strictly greater"): F.resize(img, size=32, max_size=32) # TODO: # @pytest.mark.parametrize("device", cpu_and_gpu()) # @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64, torch.float16]) # @pytest.mark.parametrize("size", [[96, 72], [96, 420], [420, 72]]) # @pytest.mark.parametrize("interpolation", [BILINEAR, BICUBIC]) # def test_resize_antialias(device, dt, size, interpolation): # # if dt == torch.float16 and device == "cpu": # # skip float16 on CPU case # return # # torch.manual_seed(12) # # script_fn = torch.jit.script(F.resize) # tensor, pil_img = _create_data(320, 290, device=device) # # if dt is not None: # # This is a trivial cast to float of uint8 data to test all cases # tensor = tensor.to(dt) # # resized_tensor = F.resize(tensor, size=size, interpolation=interpolation, antialias=True) # resized_pil_img = F.resize(pil_img, size=size, interpolation=interpolation) # # assert resized_tensor.size()[1:] == resized_pil_img.size[::-1] # # resized_tensor_f = resized_tensor # # we need to cast to uint8 to compare with PIL image # if resized_tensor_f.dtype == torch.uint8: # resized_tensor_f = resized_tensor_f.to(torch.float) # # _assert_approx_equal_tensor_to_pil(resized_tensor_f, resized_pil_img, tol=0.5, msg=f"{size}, {interpolation}, {dt}") # # accepted_tol = 1.0 + 1e-5 # if interpolation == BICUBIC: # # this overall mean value to make the tests pass # # High value is mostly required for test cases with # # downsampling and upsampling where we can not exactly # # match PIL implementation. # accepted_tol = 15.0 # # _assert_approx_equal_tensor_to_pil( # resized_tensor_f, resized_pil_img, tol=accepted_tol, agg_method="max", msg=f"{size}, {interpolation}, {dt}" # ) # # # if isinstance(size, int): # # script_size = [ # # size, # # ] # # else: # # script_size = size # # # # resize_result = script_fn(tensor, size=script_size, interpolation=interpolation, antialias=True) # # assert_equal(resized_tensor, resize_result) # TODO: # @needs_cuda # @pytest.mark.parametrize("interpolation", [BILINEAR, BICUBIC]) # def test_assert_resize_antialias(interpolation): # # # Checks implementation on very large scales # # and catch TORCH_CHECK inside PyTorch implementation # torch.manual_seed(12) # tensor, _ = _create_data(1000, 1000, device="cuda") # # # Error message is not yet updated in pytorch nightly # # with pytest.raises(RuntimeError, match=r"Provided interpolation parameters can not be handled"): # with pytest.raises(RuntimeError, match=r"Too much shared memory required"): # F.resize(tensor, size=(5, 5), interpolation=interpolation, antialias=True) # TODO: # @pytest.mark.parametrize("device", cpu_and_gpu()) # @pytest.mark.parametrize("dt", [torch.float32, torch.float64, torch.float16]) # @pytest.mark.parametrize("size", [[10, 7], [10, 42], [42, 7]]) # @pytest.mark.parametrize("interpolation", [BILINEAR, BICUBIC]) # def test_interpolate_antialias_backward(device, dt, size, interpolation): # # if dt == torch.float16 and device == "cpu": # # skip float16 on CPU case # return # # torch.manual_seed(12) # x = (torch.rand(1, 32, 29, 3, dtype=torch.double, device=device).permute(0, 3, 1, 2).requires_grad_(True),) # resize = partial(F.resize, size=size, interpolation=interpolation, antialias=True) # assert torch.autograd.gradcheck(resize, x, eps=1e-8, atol=1e-6, rtol=1e-6, fast_mode=False) # # x = (torch.rand(1, 3, 32, 29, dtype=torch.double, device=device, requires_grad=True),) # assert torch.autograd.gradcheck(resize, x, eps=1e-8, atol=1e-6, rtol=1e-6, fast_mode=False) def check_functional_vs_PIL_vs_scripted( fn, fn_pil, fn_t, config, device, dtype, channels=3, tol=2.0 + 1e-10, agg_method="max" ): # script_fn = torch.jit.script(fn) torch.manual_seed(15) tensor, pil_img = _create_data(26, 34, channels=channels, device=device) batch_tensors = _create_data_batch(16, 18, num_samples=4, channels=channels, device=device) if dtype is not None: tensor = F.convert_image_dtype(tensor, dtype) batch_tensors = F.convert_image_dtype(batch_tensors, dtype) out_fn_t = fn_t(tensor, **config) out_pil = fn_pil(pil_img, **config) # out_scripted = script_fn(tensor, **config) # assert out_fn_t.dtype == out_scripted.dtype assert out_fn_t.size()[1:] == out_pil.size[::-1] rbg_tensor = out_fn_t if out_fn_t.dtype != torch.uint8: rbg_tensor = F.convert_image_dtype(out_fn_t, torch.uint8) # Check that max difference does not exceed 2 in [0, 255] range # Exact matching is not possible due to incompatibility convert_image_dtype and PIL results _assert_approx_equal_tensor_to_pil(rbg_tensor.float(), out_pil, tol=tol, agg_method=agg_method) atol = 1e-6 if out_fn_t.dtype == torch.uint8 and "cuda" in torch.device(device).type: atol = 1.0 # assert out_fn_t.allclose(out_scripted, atol=atol) # FIXME: fn will be scripted again in _test_fn_on_batch. We could avoid that. _test_fn_on_batch(batch_tensors, fn, scripted_fn_atol=atol, **config) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("config", [{"brightness_factor": f} for f in (0.1, 0.5, 1.0, 1.34, 2.5)]) @pytest.mark.parametrize("channels", [1, 3]) def test_adjust_brightness(device, dtype, config, channels): check_functional_vs_PIL_vs_scripted( F.adjust_brightness, F_pil.adjust_brightness, F_t.adjust_brightness, config, device, dtype, channels, ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("channels", [1, 3]) def test_invert(device, dtype, channels): check_functional_vs_PIL_vs_scripted( F.invert, F_pil.invert, F_t.invert, {}, device, dtype, channels, tol=1.0, agg_method="max" ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("config", [{"bits": bits} for bits in range(0, 8)]) @pytest.mark.parametrize("channels", [1, 3]) def test_posterize(device, config, channels): check_functional_vs_PIL_vs_scripted( F.posterize, F_pil.posterize, F_t.posterize, config, device, dtype=None, channels=channels, tol=1.0, agg_method="max", ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("config", [{"threshold": threshold} for threshold in [0, 64, 128, 192, 255]]) @pytest.mark.parametrize("channels", [1, 3]) def test_solarize1(device, config, channels): check_functional_vs_PIL_vs_scripted( F.solarize, F_pil.solarize, F_t.solarize, config, device, dtype=None, channels=channels, tol=1.0, agg_method="max", ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (torch.float32, torch.float64)) @pytest.mark.parametrize("config", [{"threshold": threshold} for threshold in [0.0, 0.25, 0.5, 0.75, 1.0]]) @pytest.mark.parametrize("channels", [1, 3]) def test_solarize2(device, dtype, config, channels): check_functional_vs_PIL_vs_scripted( F.solarize, lambda img, threshold: F_pil.solarize(img, 255 * threshold), F_t.solarize, config, device, dtype, channels, tol=1.0, agg_method="max", ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("threshold", [0.0, 0.25, 0.5, 0.75, 1.0]) def test_solarize_threshold1_bound(threshold, device): img = torch.rand((3, 12, 23)).to(device) F_t.solarize(img, threshold) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("threshold", [1.5]) def test_solarize_threshold1_upper_bound(threshold, device): img = torch.rand((3, 12, 23)).to(device) with pytest.raises(TypeError, match="Threshold should be less than bound of img."): F_t.solarize(img, threshold) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("threshold", [0, 64, 128, 192, 255]) def test_solarize_threshold2_bound(threshold, device): img = torch.randint(0, 256, (3, 12, 23)).to(device) F_t.solarize(img, threshold) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("threshold", [260]) def test_solarize_threshold2_upper_bound(threshold, device): img = torch.randint(0, 256, (3, 12, 23)).to(device) with pytest.raises(TypeError, match="Threshold should be less than bound of img."): F_t.solarize(img, threshold) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("config", [{"sharpness_factor": f} for f in [0.2, 0.5, 1.0, 1.5, 2.0]]) @pytest.mark.parametrize("channels", [1, 3]) def test_adjust_sharpness(device, dtype, config, channels): check_functional_vs_PIL_vs_scripted( F.adjust_sharpness, F_pil.adjust_sharpness, F_t.adjust_sharpness, config, device, dtype, channels, ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("channels", [1, 3]) def test_autocontrast(device, dtype, channels): check_functional_vs_PIL_vs_scripted( F.autocontrast, F_pil.autocontrast, F_t.autocontrast, {}, device, dtype, channels, tol=1.0, agg_method="max" ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("channels", [1, 3]) def test_autocontrast_equal_minmax(device, dtype, channels): a = _create_data_batch(32, 32, num_samples=1, channels=channels, device=device) a = a / 2.0 + 0.3 assert (F.autocontrast(a)[0] == F.autocontrast(a[0])).all() a[0, 0] = 0.7 assert (F.autocontrast(a)[0] == F.autocontrast(a[0])).all() @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("channels", [1, 3]) def test_equalize(device, channels): # torch.use_deterministic_algorithms(False) check_functional_vs_PIL_vs_scripted( F.equalize, F_pil.equalize, F_t.equalize, {}, device, dtype=None, channels=channels, tol=1.0, agg_method="max", ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("config", [{"contrast_factor": f} for f in [0.2, 0.5, 1.0, 1.5, 2.0]]) @pytest.mark.parametrize("channels", [1, 3]) def test_adjust_contrast(device, dtype, config, channels): check_functional_vs_PIL_vs_scripted( F.adjust_contrast, F_pil.adjust_contrast, F_t.adjust_contrast, config, device, dtype, channels ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("config", [{"saturation_factor": f} for f in [0.5, 0.75, 1.0, 1.5, 2.0]]) @pytest.mark.parametrize("channels", [1, 3]) def test_adjust_saturation(device, dtype, config, channels): check_functional_vs_PIL_vs_scripted( F.adjust_saturation, F_pil.adjust_saturation, F_t.adjust_saturation, config, device, dtype, channels ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("config", [{"hue_factor": f} for f in [-0.45, -0.25, 0.0, 0.25, 0.45]]) @pytest.mark.parametrize("channels", [1, 3]) def test_adjust_hue(device, dtype, config, channels): check_functional_vs_PIL_vs_scripted( F.adjust_hue, F_pil.adjust_hue, F_t.adjust_hue, config, device, dtype, channels, tol=16.1, agg_method="max" ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64)) @pytest.mark.parametrize("config", [{"gamma": g1, "gain": g2} for g1, g2 in zip([0.8, 1.0, 1.2], [0.7, 1.0, 1.3])]) @pytest.mark.parametrize("channels", [1, 3]) def test_adjust_gamma(device, dtype, config, channels): check_functional_vs_PIL_vs_scripted( F.adjust_gamma, F_pil.adjust_gamma, F_t.adjust_gamma, config, device, dtype, channels, ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64, torch.float16]) @pytest.mark.parametrize( "pad", [ 2, [ 3, ], [0, 3], (3, 3), [4, 2, 4, 3], ], ) @pytest.mark.parametrize( "config", [ {"padding_mode": "constant", "fill": 0}, {"padding_mode": "constant", "fill": 10}, {"padding_mode": "constant", "fill": 20}, {"padding_mode": "edge"}, {"padding_mode": "reflect"}, {"padding_mode": "symmetric"}, ], ) def test_pad(device, dt, pad, config): # script_fn = torch.jit.script(F.pad) tensor, pil_img = _create_data(7, 8, device=device) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) if dt == torch.float16 and device == "cpu": # skip float16 on CPU case return if dt is not None: # This is a trivial cast to float of uint8 data to test all cases tensor = tensor.to(dt) batch_tensors = batch_tensors.to(dt) pad_tensor = F_t.pad(tensor, pad, **config) pad_pil_img = F_pil.pad(pil_img, pad, **config) pad_tensor_8b = pad_tensor # we need to cast to uint8 to compare with PIL image if pad_tensor_8b.dtype != torch.uint8: pad_tensor_8b = pad_tensor_8b.to(torch.uint8) _assert_equal_tensor_to_pil(pad_tensor_8b, pad_pil_img, msg=f"{pad}, {config}") if isinstance(pad, int): script_pad = [ pad, ] else: script_pad = pad # pad_tensor_script = script_fn(tensor, script_pad, **config) # assert_equal(pad_tensor, pad_tensor_script, msg=f"{pad}, {config}") _test_fn_on_batch(batch_tensors, F.pad, padding=script_pad, **config) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("mode", [NEAREST, BILINEAR, BICUBIC]) def test_resized_crop(device, mode): # test values of F.resized_crop in several cases: # 1) resize to the same size, crop to the same size => should be identity tensor, _ = _create_data(26, 36, device=device) out_tensor = F.resized_crop(tensor, top=0, left=0, height=26, width=36, size=[26, 36], interpolation=mode) # assert_equal(tensor, out_tensor, msg=f"{out_tensor[0, :5, :5]} vs {tensor[0, :5, :5]}") assert np.allclose(tensor.numpy(), out_tensor.numpy()) # 2) resize by half and crop a TL corner tensor, _ = _create_data(26, 36, device=device) out_tensor = F.resized_crop(tensor, top=0, left=0, height=20, width=30, size=[10, 15], interpolation=NEAREST) expected_out_tensor = tensor[:, :20:2, :30:2] # assert_equal( # expected_out_tensor, # out_tensor, # msg=f"{expected_out_tensor[0, :10, :10]} vs {out_tensor[0, :10, :10]}", # ) assert np.allclose(expected_out_tensor.numpy(), out_tensor.numpy()) batch_tensors = _create_data_batch(26, 36, num_samples=4, device=device) _test_fn_on_batch( batch_tensors, F.resized_crop, top=1, left=2, height=20, width=30, size=[10, 15], interpolation=NEAREST ) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize( "func, args", [ (F_t.get_dimensions, ()), (F_t.get_image_size, ()), (F_t.get_image_num_channels, ()), (F_t.vflip, ()), (F_t.hflip, ()), (F_t.crop, (1, 2, 4, 5)), (F_t.adjust_brightness, (0.0,)), (F_t.adjust_contrast, (1.0,)), (F_t.adjust_hue, (-0.5,)), (F_t.adjust_saturation, (2.0,)), (F_t.pad, ([2], 2, "constant")), (F_t.resize, ([10, 11],)), (F_t.perspective, ([0.2])), (F_t.gaussian_blur, ((2, 2), (0.7, 0.5))), (F_t.invert, ()), (F_t.posterize, (0,)), (F_t.solarize, (0.3,)), (F_t.adjust_sharpness, (0.3,)), (F_t.autocontrast, ()), (F_t.equalize, ()), ], ) def test_assert_image_tensor(device, func, args): shape = (100,) tensor = torch.rand(*shape, dtype=torch.float, device=device) with pytest.raises(Exception, match=r"Tensor is not a torch image."): func(tensor, *args) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_vflip(device): # script_vflip = torch.jit.script(F.vflip) img_tensor, pil_img = _create_data(16, 18, device=device) vflipped_img = F.vflip(img_tensor) vflipped_pil_img = F.vflip(pil_img) _assert_equal_tensor_to_pil(vflipped_img, vflipped_pil_img) # # scriptable function test # vflipped_img_script = script_vflip(img_tensor) # assert_equal(vflipped_img, vflipped_img_script) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) _test_fn_on_batch(batch_tensors, F.vflip) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_hflip(device): # script_hflip = torch.jit.script(F.hflip) img_tensor, pil_img = _create_data(16, 18, device=device) hflipped_img = F.hflip(img_tensor) hflipped_pil_img = F.hflip(pil_img) _assert_equal_tensor_to_pil(hflipped_img, hflipped_pil_img) # # scriptable function test # hflipped_img_script = script_hflip(img_tensor) # assert_equal(hflipped_img, hflipped_img_script) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) _test_fn_on_batch(batch_tensors, F.hflip) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize( "top, left, height, width", [ (1, 2, 4, 5), # crop inside top-left corner (2, 12, 3, 4), # crop inside top-right corner (8, 3, 5, 6), # crop inside bottom-left corner (8, 11, 4, 3), # crop inside bottom-right corner ], ) def test_crop(device, top, left, height, width): # script_crop = torch.jit.script(F.crop) img_tensor, pil_img = _create_data(16, 18, device=device) pil_img_cropped = F.crop(pil_img, top, left, height, width) img_tensor_cropped = F.crop(img_tensor, top, left, height, width) _assert_equal_tensor_to_pil(img_tensor_cropped, pil_img_cropped) # img_tensor_cropped = script_crop(img_tensor, top, left, height, width) # _assert_equal_tensor_to_pil(img_tensor_cropped, pil_img_cropped) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) _test_fn_on_batch(batch_tensors, F.crop, top=top, left=left, height=height, width=width) # TODO: # # @pytest.mark.parametrize("device", cpu_and_gpu()) # @pytest.mark.parametrize("image_size", ("small", "large")) # @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64, torch.float16]) # @pytest.mark.parametrize("ksize", [(3, 3), [3, 5], (23, 23)]) # @pytest.mark.parametrize("sigma", [[0.5, 0.5], (0.5, 0.5), (0.8, 0.8), (1.7, 1.7)]) # @pytest.mark.parametrize("fn", [F.gaussian_blur]) # def test_gaussian_blur(device, image_size, dt, ksize, sigma, fn): # # # true_cv2_results = { # # # np_img = np.arange(3 * 10 * 12, dtype="uint8").reshape((10, 12, 3)) # # # cv2.GaussianBlur(np_img, ksize=(3, 3), sigmaX=0.8) # # "3_3_0.8": ... # # # cv2.GaussianBlur(np_img, ksize=(3, 3), sigmaX=0.5) # # "3_3_0.5": ... # # # cv2.GaussianBlur(np_img, ksize=(3, 5), sigmaX=0.8) # # "3_5_0.8": ... # # # cv2.GaussianBlur(np_img, ksize=(3, 5), sigmaX=0.5) # # "3_5_0.5": ... # # # np_img2 = np.arange(26 * 28, dtype="uint8").reshape((26, 28)) # # # cv2.GaussianBlur(np_img2, ksize=(23, 23), sigmaX=1.7) # # "23_23_1.7": ... # # } # # p = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "gaussian_blur_opencv_results.pt") # # true_cv2_results = torch.load(p) # # # # if image_size == "small": # tensor = ( # torch.tensor(np.arange(3 * 10 * 12, dtype="uint8").reshape((10, 12, 3))).permute(2, 0, 1) # ) # pytorch_tensor = \ # pytorch.tensor(np.arange(3 * 10 * 12, dtype="uint8").reshape((10, 12, 3))).permute(2, 0, 1) # # else: # tensor = torch.tensor(np.arange(26 * 28, dtype="uint8").reshape((1, 26, 28))) # pytorch_tensor = pytorch.tensor(np.arange(26 * 28, dtype="uint8").reshape((1, 26, 28))) # # if dt == torch.float16 and device == "cpu": # # skip float16 on CPU case # return # # if dt is not None: # tensor = tensor.to(dtype=dt) # pytorch_tensor = pytorch_tensor.to(dtype=dt) # # # _ksize = (ksize, ksize) if isinstance(ksize, int) else ksize # # _sigma = sigma[0] if sigma is not None else None # # shape = tensor.shape # # gt_key = f"{shape[-2]}_{shape[-1]}_{shape[-3]}__{_ksize[0]}_{_ksize[1]}_{_sigma}" # # if gt_key not in true_cv2_results: # # return # # # # true_out = ( # # torch.tensor(true_cv2_results[gt_key]).reshape(shape[-2], shape[-1], shape[-3]).permute(2, 0, 1).to(tensor) # # ) # # torch_out = tv_F.gaussian_blur(pytorch_tensor, kernel_size=ksize, sigma=sigma).clone() # print("------------torch_out:", torch_out.shape) # out = fn(tensor, kernel_size=ksize, sigma=sigma).clone() # print("------------torch_out:", torch_out.shape) # print("------------out:", out.numpy().shape) # # # torch.testing.assert_close(out, true_out, rtol=0.0, atol=1.0, msg=f"{ksize}, {sigma}") # assert np.allclose(out.numpy(), torch_out.numpy()) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_hsv2rgb(device): # scripted_fn = torch.jit.script(F_t._hsv2rgb) shape = (3, 100, 150) for _ in range(10): hsv_img = torch.rand(*shape, dtype=torch.float, device=device) rgb_img = F_t._hsv2rgb(hsv_img) ft_img = rgb_img.permute(1, 2, 0).flatten(0, 1) ( h, s, v, ) = hsv_img.unbind(0) h = h.flatten().cpu().numpy() s = s.flatten().cpu().numpy() v = v.flatten().cpu().numpy() rgb = [] for h1, s1, v1 in zip(h, s, v): rgb.append(colorsys.hsv_to_rgb(h1, s1, v1)) colorsys_img = torch.tensor(rgb, dtype=torch.float32, device=device) # torch.testing.assert_close(ft_img, colorsys_img, rtol=0.0, atol=1e-5) assert np.allclose(ft_img.numpy(), colorsys_img.numpy(), rtol=0.0, atol=1e-5) # s_rgb_img = scripted_fn(hsv_img) # torch.testing.assert_close(rgb_img, s_rgb_img) batch_tensors = _create_data_batch(120, 100, num_samples=4, device=device).float() _test_fn_on_batch(batch_tensors, F_t._hsv2rgb) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_rgb2hsv(device): # scripted_fn = torch.jit.script(F_t._rgb2hsv) shape = (3, 150, 100) for _ in range(10): rgb_img = torch.rand(*shape, dtype=torch.float, device=device) hsv_img = F_t._rgb2hsv(rgb_img) ft_hsv_img = hsv_img.permute(1, 2, 0).flatten(0, 1) ( r, g, b, ) = rgb_img.unbind(dim=-3) r = r.flatten().cpu().numpy() g = g.flatten().cpu().numpy() b = b.flatten().cpu().numpy() hsv = [] for r1, g1, b1 in zip(r, g, b): hsv.append(colorsys.rgb_to_hsv(r1, g1, b1)) colorsys_img = torch.tensor(hsv, dtype=torch.float32, device=device) ft_hsv_img_h, ft_hsv_img_sv = torch.split(ft_hsv_img, [1, 2], dim=1) colorsys_img_h, colorsys_img_sv = torch.split(colorsys_img, [1, 2], dim=1) max_diff_h = ((colorsys_img_h * 2 * math.pi).sin() - (ft_hsv_img_h * 2 * math.pi).sin()).abs().max() max_diff_sv = (colorsys_img_sv - ft_hsv_img_sv).abs().max() max_diff = max(max_diff_h, max_diff_sv) assert max_diff < 1e-5 # s_hsv_img = scripted_fn(rgb_img) # torch.testing.assert_close(hsv_img, s_hsv_img, rtol=1e-5, atol=1e-7) batch_tensors = _create_data_batch(120, 100, num_samples=4, device=device).float() _test_fn_on_batch(batch_tensors, F_t._rgb2hsv) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("num_output_channels", (3, 1)) def test_rgb_to_grayscale(device, num_output_channels): # script_rgb_to_grayscale = torch.jit.script(F.rgb_to_grayscale) img_tensor, pil_img = _create_data(32, 34, device=device) gray_pil_image = F.to_grayscale(pil_img, num_output_channels=num_output_channels) gray_tensor = F.rgb_to_grayscale(img_tensor, num_output_channels=num_output_channels) _assert_approx_equal_tensor_to_pil(gray_tensor.float(), gray_pil_image, tol=1.0 + 1e-10, agg_method="max") # s_gray_tensor = script_rgb_to_grayscale(img_tensor, num_output_channels=num_output_channels) # assert_equal(s_gray_tensor, gray_tensor) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) _test_fn_on_batch(batch_tensors, F.rgb_to_grayscale, num_output_channels=num_output_channels) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_center_crop(device): # script_center_crop = torch.jit.script(F.center_crop) img_tensor, pil_img = _create_data(32, 34, device=device) cropped_pil_image = F.center_crop(pil_img, [10, 11]) cropped_tensor = F.center_crop(img_tensor, [10, 11]) _assert_equal_tensor_to_pil(cropped_tensor, cropped_pil_image) # cropped_tensor = script_center_crop(img_tensor, [10, 11]) # _assert_equal_tensor_to_pil(cropped_tensor, cropped_pil_image) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) _test_fn_on_batch(batch_tensors, F.center_crop, output_size=[10, 11]) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_five_crop(device): # script_five_crop = torch.jit.script(F.five_crop) img_tensor, pil_img = _create_data(32, 34, device=device) cropped_pil_images = F.five_crop(pil_img, [10, 11]) cropped_tensors = F.five_crop(img_tensor, [10, 11]) for i in range(5): _assert_equal_tensor_to_pil(cropped_tensors[i], cropped_pil_images[i]) # cropped_tensors = script_five_crop(img_tensor, [10, 11]) # for i in range(5): # _assert_equal_tensor_to_pil(cropped_tensors[i], cropped_pil_images[i]) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) tuple_transformed_batches = F.five_crop(batch_tensors, [10, 11]) for i in range(len(batch_tensors)): img_tensor = batch_tensors[i, ...] tuple_transformed_imgs = F.five_crop(img_tensor, [10, 11]) assert len(tuple_transformed_imgs) == len(tuple_transformed_batches) for j in range(len(tuple_transformed_imgs)): true_transformed_img = tuple_transformed_imgs[j] transformed_img = tuple_transformed_batches[j][i, ...] # assert_equal(true_transformed_img, transformed_img) assert np.allclose(true_transformed_img.numpy(), transformed_img.numpy()) # scriptable function test # s_tuple_transformed_batches = script_five_crop(batch_tensors, [10, 11]) # for transformed_batch, s_transformed_batch in zip(tuple_transformed_batches, s_tuple_transformed_batches): # assert_equal(transformed_batch, s_transformed_batch) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_ten_crop(device): # script_ten_crop = torch.jit.script(F.ten_crop) img_tensor, pil_img = _create_data(32, 34, device=device) cropped_pil_images = F.ten_crop(pil_img, [10, 11]) cropped_tensors = F.ten_crop(img_tensor, [10, 11]) for i in range(10): _assert_equal_tensor_to_pil(cropped_tensors[i], cropped_pil_images[i]) # cropped_tensors = script_ten_crop(img_tensor, [10, 11]) # for i in range(10): # _assert_equal_tensor_to_pil(cropped_tensors[i], cropped_pil_images[i]) batch_tensors = _create_data_batch(16, 18, num_samples=4, device=device) tuple_transformed_batches = F.ten_crop(batch_tensors, [10, 11]) for i in range(len(batch_tensors)): img_tensor = batch_tensors[i, ...] tuple_transformed_imgs = F.ten_crop(img_tensor, [10, 11]) assert len(tuple_transformed_imgs) == len(tuple_transformed_batches) for j in range(len(tuple_transformed_imgs)): true_transformed_img = tuple_transformed_imgs[j] transformed_img = tuple_transformed_batches[j][i, ...] # assert_equal(true_transformed_img, transformed_img) assert np.allclose(true_transformed_img.numpy(), transformed_img.numpy()) # # scriptable function test # s_tuple_transformed_batches = script_ten_crop(batch_tensors, [10, 11]) # for transformed_batch, s_transformed_batch in zip(tuple_transformed_batches, s_tuple_transformed_batches): # assert_equal(transformed_batch, s_transformed_batch) if __name__ == "__main__": pytest.main([__file__]) msadapter/testing/ut/torchvision/test_ops.py000066400000000000000000001623561436713554100217520ustar00rootroot00000000000000import math import os from abc import ABC, abstractmethod from functools import lru_cache from itertools import product from typing import Callable, List, Tuple import mindspore as ms import numpy as np import pytest import ms_adapter.pytorch as torch import ms_adapter.pytorch.nn.functional as F from common_utils import cpu_and_gpu, needs_cuda from PIL import Image from ms_adapter.pytorch import nn, Tensor from ms_adapter.pytorch.nn.modules.utils import _pair from ms_adapter.torchvision import ops class RoIOpTesterModuleWrapper(nn.Module): def __init__(self, obj): super().__init__() self.layer = obj self.n_inputs = 2 def forward(self, a, b): self.layer(a, b) class MultiScaleRoIAlignModuleWrapper(nn.Module): def __init__(self, obj): super().__init__() self.layer = obj self.n_inputs = 3 def forward(self, a, b, c): self.layer(a, b, c) class DeformConvModuleWrapper(nn.Module): def __init__(self, obj): super().__init__() self.layer = obj self.n_inputs = 3 def forward(self, a, b, c): self.layer(a, b, c) class StochasticDepthWrapper(nn.Module): def __init__(self, obj): super().__init__() self.layer = obj self.n_inputs = 1 def forward(self, a): self.layer(a) class DropBlockWrapper(nn.Module): def __init__(self, obj): super().__init__() self.layer = obj self.n_inputs = 1 def forward(self, a): self.layer(a) class RoIOpTester(ABC): dtype = torch.float32 @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("contiguous", (True, False)) def test_forward(self, device, contiguous, x_dtype=None, rois_dtype=None, **kwargs): x_dtype = self.dtype if x_dtype is None else x_dtype rois_dtype = self.dtype if rois_dtype is None else rois_dtype pool_size = 5 # n_channels % (pool_size ** 2) == 0 required for PS opeartions. n_channels = 2 * (pool_size ** 2) x = torch.rand(2, n_channels, 10, 10, dtype=x_dtype, device=device) if not contiguous: x = x.permute(0, 1, 3, 2) rois = torch.tensor( [[0, 0, 0, 9, 9], [0, 0, 5, 4, 9], [0, 5, 5, 9, 9], [1, 0, 0, 9, 9]], # format is (xyxy) dtype=rois_dtype, device=device, ) pool_h, pool_w = pool_size, pool_size y = self.fn(x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=2, **kwargs) # the following should be true whether we're running an autocast test or not. assert y.dtype == x.dtype gt_y = self.expected_fn( x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=2, device=device, dtype=self.dtype, **kwargs ) tol = 1e-3 if (x_dtype is torch.half or rois_dtype is torch.half) else 1e-5 # torch.testing.assert_close(gt_y.to(y), y, rtol=tol, atol=tol) assert np.allclose(gt_y.to(y).numpy(), y.numpy(), rtol=tol, atol=tol) def _helper_boxes_shape(self, func): # test boxes as Tensor[N, 5] with pytest.raises(AssertionError): a = torch.linspace(1, 8 * 8, 8 * 8).reshape(1, 1, 8, 8) boxes = torch.tensor([[0, 0, 3, 3]], dtype=a.dtype) func(a, boxes, output_size=(2, 2)) @abstractmethod def fn(*args, **kwargs): pass @abstractmethod def make_obj(*args, **kwargs): pass # @abstractmethod # def get_script_fn(*args, **kwargs): # pass @abstractmethod def expected_fn(*args, **kwargs): pass # class TestRoiPool(RoIOpTester): # def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs): # return ops.RoIPool((pool_h, pool_w), spatial_scale)(x, rois) # # def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, wrap=False): # obj = ops.RoIPool((pool_h, pool_w), spatial_scale) # return RoIOpTesterModuleWrapper(obj) if wrap else obj # # # def get_script_fn(self, rois, pool_size): # # scriped = torch.jit.script(ops.roi_pool) # # return lambda x: scriped(x, rois, pool_size) # # def expected_fn( # self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, device=None, dtype=torch.float64 # ): # if device is None: # device = torch.device("cpu") # # n_channels = x.size(1) # y = torch.zeros(rois.size(0), n_channels, pool_h, pool_w, dtype=dtype, device=device) # # def get_slice(k, block): # return slice(int(np.floor(k * block)), int(np.ceil((k + 1) * block))) # # for roi_idx, roi in enumerate(rois): # batch_idx = int(roi[0]) # j_begin, i_begin, j_end, i_end = (int(round(x.item() * spatial_scale)) for x in roi[1:]) # roi_x = x[batch_idx, :, i_begin : i_end + 1, j_begin : j_end + 1] # # roi_h, roi_w = roi_x.shape[-2:] # bin_h = roi_h / pool_h # bin_w = roi_w / pool_w # # for i in range(0, pool_h): # for j in range(0, pool_w): # bin_x = roi_x[:, get_slice(i, bin_h), get_slice(j, bin_w)] # if bin_x.numel() > 0: # y[roi_idx, :, i, j] = bin_x.reshape(n_channels, -1).max(dim=1)[0] # return y # # def test_boxes_shape(self): # self._helper_boxes_shape(ops.roi_pool) # class TestPSRoIPool(RoIOpTester): # def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs): # return ops.PSRoIPool((pool_h, pool_w), 1)(x, rois) # # def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, wrap=False): # obj = ops.PSRoIPool((pool_h, pool_w), spatial_scale) # return RoIOpTesterModuleWrapper(obj) if wrap else obj # # # def get_script_fn(self, rois, pool_size): # # scriped = torch.jit.script(ops.ps_roi_pool) # # return lambda x: scriped(x, rois, pool_size) # # def expected_fn( # self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, device=None, dtype=torch.float64 # ): # if device is None: # device = torch.device("cpu") # n_input_channels = x.size(1) # assert n_input_channels % (pool_h * pool_w) == 0, "input channels must be divisible by ph * pw" # n_output_channels = int(n_input_channels / (pool_h * pool_w)) # y = torch.zeros(rois.size(0), n_output_channels, pool_h, pool_w, dtype=dtype, device=device) # # def get_slice(k, block): # return slice(int(np.floor(k * block)), int(np.ceil((k + 1) * block))) # # for roi_idx, roi in enumerate(rois): # batch_idx = int(roi[0]) # j_begin, i_begin, j_end, i_end = (int(round(x.item() * spatial_scale)) for x in roi[1:]) # roi_x = x[batch_idx, :, i_begin : i_end + 1, j_begin : j_end + 1] # # roi_height = max(i_end - i_begin, 1) # roi_width = max(j_end - j_begin, 1) # bin_h, bin_w = roi_height / float(pool_h), roi_width / float(pool_w) # # for i in range(0, pool_h): # for j in range(0, pool_w): # bin_x = roi_x[:, get_slice(i, bin_h), get_slice(j, bin_w)] # if bin_x.numel() > 0: # area = bin_x.size(-2) * bin_x.size(-1) # for c_out in range(0, n_output_channels): # c_in = c_out * (pool_h * pool_w) + pool_w * i + j # t = torch.sum(bin_x[c_in, :, :]) # y[roi_idx, c_out, i, j] = t / area # return y # # def test_boxes_shape(self): # self._helper_boxes_shape(ops.ps_roi_pool) def bilinear_interpolate(data, y, x, snap_border=False): height, width = data.shape if snap_border: if -1 < y <= 0: y = 0 elif height - 1 <= y < height: y = height - 1 if -1 < x <= 0: x = 0 elif width - 1 <= x < width: x = width - 1 y_low = int(math.floor(y)) x_low = int(math.floor(x)) y_high = y_low + 1 x_high = x_low + 1 wy_h = y - y_low wx_h = x - x_low wy_l = 1 - wy_h wx_l = 1 - wx_h val = 0 for wx, xp in zip((wx_l, wx_h), (x_low, x_high)): for wy, yp in zip((wy_l, wy_h), (y_low, y_high)): if 0 <= yp < height and 0 <= xp < width: val += wx * wy * data[yp, xp] return val class TestRoIAlign(RoIOpTester): def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, aligned=False, **kwargs): return ops.RoIAlign( (pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio, aligned=aligned )(x, rois) def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, sampling_ratio=-1, aligned=False, wrap=False): obj = ops.RoIAlign( (pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio, aligned=aligned ) return RoIOpTesterModuleWrapper(obj) if wrap else obj def expected_fn( self, in_data, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, aligned=False, device=None, dtype=torch.float64, ): if device is None: device = torch.device("cpu") n_channels = in_data.size(1) out_data = torch.zeros(rois.size(0), n_channels, pool_h, pool_w, dtype=dtype, device=device) offset = 0.5 if aligned else 0.0 for r, roi in enumerate(rois): batch_idx = int(roi[0]) j_begin, i_begin, j_end, i_end = (x.item() * spatial_scale - offset for x in roi[1:]) roi_h = i_end - i_begin roi_w = j_end - j_begin bin_h = roi_h / pool_h bin_w = roi_w / pool_w for i in range(0, pool_h): start_h = i_begin + i * bin_h grid_h = sampling_ratio if sampling_ratio > 0 else int(np.ceil(bin_h)) for j in range(0, pool_w): start_w = j_begin + j * bin_w grid_w = sampling_ratio if sampling_ratio > 0 else int(np.ceil(bin_w)) for channel in range(0, n_channels): val = 0 for iy in range(0, grid_h): y = start_h + (iy + 0.5) * bin_h / grid_h for ix in range(0, grid_w): x = start_w + (ix + 0.5) * bin_w / grid_w val += bilinear_interpolate(in_data[batch_idx, channel, :, :], y, x, snap_border=True) val /= grid_h * grid_w out_data[r, channel, i, j] = val return out_data def test_boxes_shape(self): self._helper_boxes_shape(ops.roi_align) @pytest.mark.parametrize("aligned", (True, False)) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("contiguous", (True, False)) def test_forward(self, device, contiguous, aligned, x_dtype=None, rois_dtype=None): torch.random.manual_seed(3) super().test_forward( device=device, contiguous=contiguous, x_dtype=x_dtype, rois_dtype=rois_dtype, aligned=aligned ) def _make_rois(self, img_size, num_imgs, dtype, num_rois=1000): rois = torch.randint(0, img_size // 2, size=(num_rois, 5)).to(dtype) rois[:, 0] = torch.randint(0, num_imgs, size=(num_rois,)) # set batch index rois[:, 3:] += rois[:, 1:3] # make sure boxes aren't degenerate return rois # class TestPSRoIAlign(RoIOpTester): # def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs): # return ops.PSRoIAlign((pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio)(x, rois) # # def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, sampling_ratio=-1, wrap=False): # obj = ops.PSRoIAlign((pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio) # return RoIOpTesterModuleWrapper(obj) if wrap else obj # # # def get_script_fn(self, rois, pool_size): # # scriped = torch.jit.script(ops.ps_roi_align) # # return lambda x: scriped(x, rois, pool_size) # # def expected_fn( # self, in_data, rois, pool_h, pool_w, device, spatial_scale=1, sampling_ratio=-1, dtype=torch.float64 # ): # if device is None: # device = torch.device("cpu") # n_input_channels = in_data.size(1) # assert n_input_channels % (pool_h * pool_w) == 0, "input channels must be divisible by ph * pw" # n_output_channels = int(n_input_channels / (pool_h * pool_w)) # out_data = torch.zeros(rois.size(0), n_output_channels, pool_h, pool_w, dtype=dtype, device=device) # # for r, roi in enumerate(rois): # batch_idx = int(roi[0]) # j_begin, i_begin, j_end, i_end = (x.item() * spatial_scale - 0.5 for x in roi[1:]) # # roi_h = i_end - i_begin # roi_w = j_end - j_begin # bin_h = roi_h / pool_h # bin_w = roi_w / pool_w # # for i in range(0, pool_h): # start_h = i_begin + i * bin_h # grid_h = sampling_ratio if sampling_ratio > 0 else int(np.ceil(bin_h)) # for j in range(0, pool_w): # start_w = j_begin + j * bin_w # grid_w = sampling_ratio if sampling_ratio > 0 else int(np.ceil(bin_w)) # for c_out in range(0, n_output_channels): # c_in = c_out * (pool_h * pool_w) + pool_w * i + j # # val = 0 # for iy in range(0, grid_h): # y = start_h + (iy + 0.5) * bin_h / grid_h # for ix in range(0, grid_w): # x = start_w + (ix + 0.5) * bin_w / grid_w # val += bilinear_interpolate(in_data[batch_idx, c_in, :, :], y, x, snap_border=True) # val /= grid_h * grid_w # # out_data[r, c_out, i, j] = val # return out_data # # def test_boxes_shape(self): # self._helper_boxes_shape(ops.ps_roi_align) class TestMultiScaleRoIAlign: def make_obj(self, fmap_names=None, output_size=(7, 7), sampling_ratio=2, wrap=False): if fmap_names is None: fmap_names = ["0"] obj = ops.poolers.MultiScaleRoIAlign(fmap_names, output_size, sampling_ratio) return MultiScaleRoIAlignModuleWrapper(obj) if wrap else obj def test_msroialign_repr(self): fmap_names = ["0"] output_size = (7, 7) sampling_ratio = 2 # Pass mock feature map names t = self.make_obj(fmap_names, output_size, sampling_ratio, wrap=False) # Check integrity of object __repr__ attribute expected_string = ( f"MultiScaleRoIAlign(featmap_names={fmap_names}, output_size={output_size}, " f"sampling_ratio={sampling_ratio})" ) assert repr(t) == expected_string class TestNMS: def _reference_nms(self, boxes, scores, iou_threshold): """ Args: box_scores (N, 5): boxes in corner-form and probabilities. iou_threshold: intersection over union threshold. Returns: picked: a list of indexes of the kept boxes """ picked = [] _, indexes = scores.sort(descending=True) while len(indexes) > 0: current = indexes[0] picked.append(current.item()) if len(indexes) == 1: break current_box = boxes[current, :] indexes = indexes[1:] rest_boxes = boxes[indexes, :] iou = ops.box_iou(rest_boxes, current_box.unsqueeze(0)).squeeze(1) indexes = indexes[iou <= iou_threshold] return torch.as_tensor(picked) def _create_tensors_with_iou(self, N, iou_thresh): # force last box to have a pre-defined iou with the first box # let b0 be [x0, y0, x1, y1], and b1 be [x0, y0, x1 + d, y1], # then, in order to satisfy ops.iou(b0, b1) == iou_thresh, # we need to have d = (x1 - x0) * (1 - iou_thresh) / iou_thresh # Adjust the threshold upward a bit with the intent of creating # at least one box that exceeds (barely) the threshold and so # should be suppressed. boxes = torch.rand(N, 4) * 100 boxes[:, 2:] += boxes[:, :2] boxes[-1, :] = boxes[0, :] x0, y0, x1, y1 = boxes[-1].tolist() iou_thresh += 1e-5 boxes[-1, 2] += (x1 - x0) * (1 - iou_thresh) / iou_thresh scores = torch.rand(N) return boxes, scores @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8)) @pytest.mark.parametrize("seed", range(10)) def test_nms_ref(self, iou, seed): torch.random.manual_seed(seed) err_msg = "NMS incompatible between CPU and reference implementation for IoU={}" boxes, scores = self._create_tensors_with_iou(1000, iou) keep_ref = self._reference_nms(boxes, scores, iou) keep = ops.nms(boxes, scores, iou) assert np.allclose(keep.numpy(), keep_ref.numpy()) def test_nms_input_errors(self): with pytest.raises(RuntimeError): ops.nms(torch.rand(4), torch.rand(3), 0.5) with pytest.raises(RuntimeError): ops.nms(torch.rand(3, 5), torch.rand(3), 0.5) with pytest.raises(RuntimeError): ops.nms(torch.rand(3, 4), torch.rand(3, 2), 0.5) with pytest.raises(RuntimeError): ops.nms(torch.rand(3, 4), torch.rand(4), 0.5) @needs_cuda @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8)) def test_nms_cuda(self, iou, dtype=torch.float64): tol = 1e-3 if dtype is torch.half else 1e-5 err_msg = "NMS incompatible between CPU and CUDA for IoU={}" boxes, scores = self._create_tensors_with_iou(1000, iou) r_cpu = ops.nms(boxes, scores, iou) r_cuda = ops.nms(boxes.cuda(), scores.cuda(), iou) is_eq = np.allclose(r_cpu.numpy(), r_cuda.cpu().numpy()) if not is_eq: # if the indices are not the same, ensure that it's because the scores # are duplicate is_eq = np.allclose(scores[r_cpu].numpy(), scores[r_cuda.cpu()].numpy(), rtol=tol, atol=tol) assert is_eq, err_msg.format(iou) @needs_cuda def test_nms_cuda_float16(self): boxes = torch.tensor( [ [285.3538, 185.5758, 1193.5110, 851.4551], [285.1472, 188.7374, 1192.4984, 851.0669], [279.2440, 197.9812, 1189.4746, 849.2019], ] ).cuda() scores = torch.tensor([0.6370, 0.7569, 0.3966]).cuda() iou_thres = 0.2 keep32 = ops.nms(boxes, scores, iou_thres) keep16 = ops.nms(boxes.to(torch.float16), scores.to(torch.float16), iou_thres) assert np.allclose(keep32.numpy(), keep16.numpy()) @pytest.mark.parametrize("seed", range(10)) def test_batched_nms_implementations(self, seed): """Make sure that both implementations of batched_nms yield identical results""" torch.random.manual_seed(seed) num_boxes = 1000 iou_threshold = 0.9 boxes = torch.cat((torch.rand(num_boxes, 2), torch.rand(num_boxes, 2) + 10), dim=1) assert max(boxes[:, 0]) < min(boxes[:, 2]) # x1 < x2 assert max(boxes[:, 1]) < min(boxes[:, 3]) # y1 < y2 scores = torch.rand(num_boxes) idxs = torch.randint(0, 4, size=(num_boxes,)) keep_vanilla = ops.boxes._batched_nms_vanilla(boxes, scores, idxs, iou_threshold) keep_trick = ops.boxes._batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold) assert np.allclose(keep_vanilla.numpy(), keep_trick.numpy()) # Also make sure an empty tensor is returned if boxes is empty empty = torch.empty((0,), dtype=torch.int64) assert np.allclose(empty.numpy(), ops.batched_nms(empty, None, None, None).numpy()) class TestDeformConv: dtype = torch.float32 def expected_fn(self, x, weight, offset, mask, bias, stride=1, padding=0, dilation=1): stride_h, stride_w = _pair(stride) pad_h, pad_w = _pair(padding) dil_h, dil_w = _pair(dilation) weight_h, weight_w = weight.shape[-2:] n_batches, n_in_channels, in_h, in_w = x.shape n_out_channels = weight.shape[0] out_h = (in_h + 2 * pad_h - (dil_h * (weight_h - 1) + 1)) // stride_h + 1 out_w = (in_w + 2 * pad_w - (dil_w * (weight_w - 1) + 1)) // stride_w + 1 n_offset_grps = offset.shape[1] // (2 * weight_h * weight_w) in_c_per_offset_grp = n_in_channels // n_offset_grps n_weight_grps = n_in_channels // weight.shape[1] in_c_per_weight_grp = weight.shape[1] out_c_per_weight_grp = n_out_channels // n_weight_grps out = torch.zeros(n_batches, n_out_channels, out_h, out_w, device=x.device, dtype=x.dtype) for b in range(n_batches): for c_out in range(n_out_channels): for i in range(out_h): for j in range(out_w): for di in range(weight_h): for dj in range(weight_w): for c in range(in_c_per_weight_grp): weight_grp = c_out // out_c_per_weight_grp c_in = weight_grp * in_c_per_weight_grp + c offset_grp = c_in // in_c_per_offset_grp mask_idx = offset_grp * (weight_h * weight_w) + di * weight_w + dj offset_idx = 2 * mask_idx pi = stride_h * i - pad_h + dil_h * di + offset[b, offset_idx, i, j] pj = stride_w * j - pad_w + dil_w * dj + offset[b, offset_idx + 1, i, j] mask_value = 1.0 if mask is not None: mask_value = mask[b, mask_idx, i, j] out[b, c_out, i, j] += ( mask_value * weight[c_out, c, di, dj] * bilinear_interpolate(x[b, c_in, :, :], pi, pj) ) out += bias.view(1, n_out_channels, 1, 1) return out @lru_cache(maxsize=None) def get_fn_args(self, device, contiguous, batch_sz, dtype): n_in_channels = 6 n_out_channels = 2 n_weight_grps = 2 n_offset_grps = 3 stride = (2, 1) pad = (1, 0) dilation = (2, 1) stride_h, stride_w = stride pad_h, pad_w = pad dil_h, dil_w = dilation weight_h, weight_w = (3, 2) in_h, in_w = (5, 4) out_h = (in_h + 2 * pad_h - (dil_h * (weight_h - 1) + 1)) // stride_h + 1 out_w = (in_w + 2 * pad_w - (dil_w * (weight_w - 1) + 1)) // stride_w + 1 x = torch.rand(batch_sz, n_in_channels, in_h, in_w, device=device, dtype=dtype, requires_grad=True) offset = torch.randn( batch_sz, n_offset_grps * 2 * weight_h * weight_w, out_h, out_w, device=device, dtype=dtype, requires_grad=True, ) mask = torch.randn( batch_sz, n_offset_grps * weight_h * weight_w, out_h, out_w, device=device, dtype=dtype, requires_grad=True ) weight = torch.randn( n_out_channels, n_in_channels // n_weight_grps, weight_h, weight_w, device=device, dtype=dtype, requires_grad=True, ) bias = torch.randn(n_out_channels, device=device, dtype=dtype, requires_grad=True) if not contiguous: x = x.permute(0, 1, 3, 2).contiguous().permute(0, 1, 3, 2) offset = offset.permute(1, 3, 0, 2).contiguous().permute(2, 0, 3, 1) mask = mask.permute(1, 3, 0, 2).contiguous().permute(2, 0, 3, 1) weight = weight.permute(3, 2, 0, 1).contiguous().permute(2, 3, 1, 0) return x, weight, offset, mask, bias, stride, pad, dilation def make_obj(self, in_channels=6, out_channels=2, kernel_size=(3, 2), groups=2, wrap=False): obj = ops.DeformConv2d( in_channels, out_channels, kernel_size, stride=(2, 1), padding=(1, 0), dilation=(2, 1), groups=groups ) return DeformConvModuleWrapper(obj) if wrap else obj @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("contiguous", (True, False)) @pytest.mark.parametrize("batch_sz", (33,)) def test_forward(self, device, contiguous, batch_sz, dtype=None): dtype = dtype or self.dtype x, _, offset, mask, _, stride, padding, dilation = self.get_fn_args(device, contiguous, batch_sz, dtype) in_channels = 6 out_channels = 2 kernel_size = (3, 2) groups = 2 tol = 2e-3 if dtype is torch.half else 1e-5 layer = self.make_obj(in_channels, out_channels, kernel_size, groups, wrap=False) res = layer(x, offset, mask) weight = layer.weight.data bias = layer.bias.data expected = self.expected_fn(x, weight, offset, mask, bias, stride=stride, padding=padding, dilation=dilation) assert np.allclose(res.to(expected).numpy(), expected.numpy(), rtol=tol, atol=tol) class TestFrozenBNT: def test_frozenbatchnorm2d_repr(self): num_features = 32 eps = 1e-5 t = ops.misc.FrozenBatchNorm2d(num_features, eps=eps) # Check integrity of object __repr__ attribute expected_string = f"FrozenBatchNorm2d({num_features}, eps={eps})" assert repr(t) == expected_string @pytest.mark.parametrize("seed", range(10)) def test_frozenbatchnorm2d_eps(self, seed): torch.random.manual_seed(seed) sample_size = (4, 32, 28, 28) x = torch.rand(sample_size) weight=torch.rand(sample_size[1]) bias=torch.rand(sample_size[1]) running_mean=torch.rand(sample_size[1]) running_var=torch.rand(sample_size[1]) num_batches_tracked=torch.tensor(100) # Check that default eps is equal to the one of BN fbn = ops.misc.FrozenBatchNorm2d(sample_size[1]) fbn.weight = torch.nn.Parameter(weight, requires_grad=False) fbn.bias = torch.nn.Parameter(bias, requires_grad=False) fbn.running_mean = torch.nn.Parameter(running_mean, requires_grad=False) fbn.running_var = torch.nn.Parameter(running_var, requires_grad=False) fbn.num_batches_tracked = torch.nn.Parameter(num_batches_tracked, requires_grad=False) # fbn.load_state_dict(state_dict, strict=False) bn = torch.nn.BatchNorm2d(sample_size[1]).eval() bn.weight = torch.nn.Parameter(weight, requires_grad=False) bn.bias = torch.nn.Parameter(bias, requires_grad=False) bn.running_mean = torch.nn.Parameter(running_mean, requires_grad=False) bn.running_var = torch.nn.Parameter(running_var, requires_grad=False) bn.num_batches_tracked = torch.nn.Parameter(num_batches_tracked, requires_grad=False) # bn.load_state_dict(state_dict) # Difference is expected to fall in an acceptable range assert np.allclose(fbn(x).numpy(), bn(x).numpy(), rtol=1e-5, atol=1e-6) # Check computation for eps > 0 fbn = ops.misc.FrozenBatchNorm2d(sample_size[1], eps=1e-5) fbn.weight = torch.nn.Parameter(weight, requires_grad=False) fbn.bias = torch.nn.Parameter(bias, requires_grad=False) fbn.running_mean = torch.nn.Parameter(running_mean, requires_grad=False) fbn.running_var = torch.nn.Parameter(running_var, requires_grad=False) fbn.num_batches_tracked = torch.nn.Parameter(num_batches_tracked, requires_grad=False) # fbn.load_state_dict(state_dict, strict=False) bn = torch.nn.BatchNorm2d(sample_size[1], eps=1e-5).eval() bn.weight = torch.nn.Parameter(weight, requires_grad=False) bn.bias = torch.nn.Parameter(bias, requires_grad=False) bn.running_mean = torch.nn.Parameter(running_mean, requires_grad=False) bn.running_var = torch.nn.Parameter(running_var, requires_grad=False) bn.num_batches_tracked = torch.nn.Parameter(num_batches_tracked, requires_grad=False) # bn.load_state_dict(state_dict) assert np.allclose(fbn(x).numpy(), bn(x).numpy(), rtol=1e-5, atol=1e-6) class TestBoxConversion: def _get_box_sequences(): # Define here the argument type of `boxes` supported by region pooling operations box_tensor = torch.tensor([[0, 0, 0, 100, 100], [1, 0, 0, 100, 100]], dtype=torch.float) box_list = [ torch.tensor([[0, 0, 100, 100]], dtype=torch.float), torch.tensor([[0, 0, 100, 100]], dtype=torch.float), ] box_tuple = tuple(box_list) return box_tensor, box_list, box_tuple @pytest.mark.parametrize("box_sequence", _get_box_sequences()) def test_check_roi_boxes_shape(self, box_sequence): # Ensure common sequences of tensors are supported ops._utils.check_roi_boxes_shape(box_sequence) @pytest.mark.parametrize("box_sequence", _get_box_sequences()) def test_convert_boxes_to_roi_format(self, box_sequence): # Ensure common sequences of tensors yield the same result ref_tensor = None if ref_tensor is None: ref_tensor = box_sequence else: assert np.allclose(ref_tensor.numpy(), ops._utils.convert_boxes_to_roi_format(box_sequence).numpy()) class TestBox: def test_bbox_same(self): box_tensor = torch.tensor( [[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float ) exp_xyxy = torch.tensor([[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float) assert exp_xyxy.size() == torch.Size([4, 4]) assert np.allclose(ops.box_convert(box_tensor, in_fmt="xyxy", out_fmt="xyxy").numpy(), exp_xyxy.numpy()) assert np.allclose(ops.box_convert(box_tensor, in_fmt="xywh", out_fmt="xywh").numpy(), exp_xyxy.numpy()) assert np.allclose(ops.box_convert(box_tensor, in_fmt="cxcywh", out_fmt="cxcywh").numpy(), exp_xyxy.numpy()) def test_bbox_xyxy_xywh(self): # Simple test convert boxes to xywh and back. Make sure they are same. # box_tensor is in x1 y1 x2 y2 format. box_tensor = torch.tensor( [[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float ) exp_xywh = torch.tensor([[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 20, 20], [23, 35, 70, 60]], dtype=torch.float) assert exp_xywh.size() == torch.Size([4, 4]) box_xywh = ops.box_convert(box_tensor, in_fmt="xyxy", out_fmt="xywh") assert np.allclose(box_xywh.numpy(), exp_xywh.numpy()) # Reverse conversion box_xyxy = ops.box_convert(box_xywh, in_fmt="xywh", out_fmt="xyxy") assert np.allclose(box_xyxy.numpy(), box_tensor.numpy()) def test_bbox_xyxy_cxcywh(self): # Simple test convert boxes to xywh and back. Make sure they are same. # box_tensor is in x1 y1 x2 y2 format. box_tensor = torch.tensor( [[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float ) exp_cxcywh = torch.tensor( [[50, 50, 100, 100], [0, 0, 0, 0], [20, 25, 20, 20], [58, 65, 70, 60]], dtype=torch.float ) assert exp_cxcywh.size() == torch.Size([4, 4]) box_cxcywh = ops.box_convert(box_tensor, in_fmt="xyxy", out_fmt="cxcywh") assert np.allclose(box_cxcywh.numpy(), exp_cxcywh.numpy()) # Reverse conversion box_xyxy = ops.box_convert(box_cxcywh, in_fmt="cxcywh", out_fmt="xyxy") assert np.allclose(box_xyxy.numpy(), box_tensor.numpy()) def test_bbox_xywh_cxcywh(self): box_tensor = torch.tensor( [[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 20, 20], [23, 35, 70, 60]], dtype=torch.float ) # This is wrong exp_cxcywh = torch.tensor( [[50, 50, 100, 100], [0, 0, 0, 0], [20, 25, 20, 20], [58, 65, 70, 60]], dtype=torch.float ) assert exp_cxcywh.size() == torch.Size([4, 4]) box_cxcywh = ops.box_convert(box_tensor, in_fmt="xywh", out_fmt="cxcywh") assert np.allclose(box_cxcywh.numpy(), exp_cxcywh.numpy()) # Reverse conversion box_xywh = ops.box_convert(box_cxcywh, in_fmt="cxcywh", out_fmt="xywh") assert np.allclose(box_xywh.numpy(), box_tensor.numpy()) @pytest.mark.parametrize("inv_infmt", ["xwyh", "cxwyh"]) @pytest.mark.parametrize("inv_outfmt", ["xwcx", "xhwcy"]) def test_bbox_invalid(self, inv_infmt, inv_outfmt): box_tensor = torch.tensor( [[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 20, 20], [23, 35, 70, 60]], dtype=torch.float ) with pytest.raises(ValueError): ops.box_convert(box_tensor, inv_infmt, inv_outfmt) class BoxTestBase(ABC): @abstractmethod def _target_fn(self) -> Tuple[bool, Callable]: pass def _perform_box_operation(self, box: Tensor, run_as_script: bool = False) -> Tensor: is_binary_fn = self._target_fn()[0] target_fn = self._target_fn()[1] box_operation = torch.jit.script(target_fn) if run_as_script else target_fn return box_operation(box, box) if is_binary_fn else box_operation(box) def _run_test(self, test_input: List, dtypes, tolerance: float, expected: List) -> None: def assert_close(box: Tensor, expected: Tensor, tolerance): out = self._perform_box_operation(box) # torch.testing.assert_close(out, expected, rtol=0.0, check_dtype=False, atol=tolerance) assert np.allclose(out.numpy(), expected.numpy(), rtol=0.0, atol=tolerance) for dtype in dtypes: actual_box = torch.tensor(test_input, dtype=dtype) expected_box = torch.tensor(expected) assert_close(actual_box, expected_box, tolerance) class TestBoxArea(BoxTestBase): def _target_fn(self) -> Tuple[bool, Callable]: return (False, ops.box_area) def _generate_int_input() -> List[List[int]]: return [[0, 0, 100, 100], [0, 0, 0, 0]] def _generate_int_expected() -> List[int]: return [10000, 0] def _generate_float_input(index: int) -> List[List[float]]: return [ [ [285.3538, 185.5758, 1193.5110, 851.4551], [285.1472, 188.7374, 1192.4984, 851.0669], [279.2440, 197.9812, 1189.4746, 849.2019], ], [[285.25, 185.625, 1194.0, 851.5], [285.25, 188.75, 1192.0, 851.0], [279.25, 198.0, 1189.0, 849.0]], ][index] def _generate_float_expected(index: int) -> List[float]: return [[604723.0806, 600965.4666, 592761.0085], [605113.875, 600495.1875, 592247.25]][index] @pytest.mark.parametrize( "test_input, dtypes, tolerance, expected", [ pytest.param( _generate_int_input(), [torch.int8, torch.int16, torch.int32, torch.int64], 1e-4, _generate_int_expected(), ), pytest.param(_generate_float_input(0), [torch.float32, torch.float64], 0.05, _generate_float_expected(0)), pytest.param(_generate_float_input(1), [torch.float16], 1e-4, _generate_float_expected(1)), ], ) def test_box_area(self, test_input: List, dtypes, tolerance: float, expected: List) -> None: self._run_test(test_input, dtypes, tolerance, expected) class TestBoxIou(BoxTestBase): def _target_fn(self) -> Tuple[bool, Callable]: return (True, ops.box_iou) def _generate_int_input() -> List[List[int]]: return [[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]] def _generate_int_expected() -> List[List[float]]: return [[1.0, 0.25, 0.0], [0.25, 1.0, 0.0], [0.0, 0.0, 1.0]] def _generate_float_input() -> List[List[float]]: return [ [285.3538, 185.5758, 1193.5110, 851.4551], [285.1472, 188.7374, 1192.4984, 851.0669], [279.2440, 197.9812, 1189.4746, 849.2019], ] def _generate_float_expected() -> List[List[float]]: return [[1.0, 0.9933, 0.9673], [0.9933, 1.0, 0.9737], [0.9673, 0.9737, 1.0]] @pytest.mark.parametrize( "test_input, dtypes, tolerance, expected", [ pytest.param( _generate_int_input(), [torch.int16, torch.int32, torch.int64], 1e-4, _generate_int_expected() ), pytest.param(_generate_float_input(), [torch.float16], 0.002, _generate_float_expected()), pytest.param(_generate_float_input(), [torch.float32, torch.float64], 1e-4, _generate_float_expected()), ], ) def test_iou(self, test_input: List, dtypes, tolerance: float, expected: List) -> None: self._run_test(test_input, dtypes, tolerance, expected) class TestGenBoxIou(BoxTestBase): def _target_fn(self) -> Tuple[bool, Callable]: return (True, ops.generalized_box_iou) def _generate_int_input() -> List[List[int]]: return [[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]] def _generate_int_expected() -> List[List[float]]: return [[1.0, 0.25, -0.7778], [0.25, 1.0, -0.8611], [-0.7778, -0.8611, 1.0]] def _generate_float_input() -> List[List[float]]: return [ [285.3538, 185.5758, 1193.5110, 851.4551], [285.1472, 188.7374, 1192.4984, 851.0669], [279.2440, 197.9812, 1189.4746, 849.2019], ] def _generate_float_expected() -> List[List[float]]: return [[1.0, 0.9933, 0.9673], [0.9933, 1.0, 0.9737], [0.9673, 0.9737, 1.0]] @pytest.mark.parametrize( "test_input, dtypes, tolerance, expected", [ pytest.param( _generate_int_input(), [torch.int16, torch.int32, torch.int64], 1e-4, _generate_int_expected() ), pytest.param(_generate_float_input(), [torch.float16], 0.002, _generate_float_expected()), pytest.param(_generate_float_input(), [torch.float32, torch.float64], 0.001, _generate_float_expected()), ], ) def test_gen_iou(self, test_input: List, dtypes, tolerance: float, expected: List) -> None: self._run_test(test_input, dtypes, tolerance, expected) class TestDistanceBoxIoU(BoxTestBase): def _target_fn(self): return (True, ops.distance_box_iou) def _generate_int_input(): return [[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]] def _generate_int_expected(): return [ [1.0000, 0.1875, -0.4444], [0.1875, 1.0000, -0.5625], [-0.4444, -0.5625, 1.0000], ] def _generate_float_input(): return [ [285.3538, 185.5758, 1193.5110, 851.4551], [285.1472, 188.7374, 1192.4984, 851.0669], [279.2440, 197.9812, 1189.4746, 849.2019], ] def _generate_float_expected(): return [[1.0, 0.9933, 0.9673], [0.9933, 1.0, 0.9737], [0.9673, 0.9737, 1.0]] @pytest.mark.parametrize( "test_input, dtypes, tolerance, expected", [ pytest.param( _generate_int_input(), [torch.int16, torch.int32, torch.int64], 1e-4, _generate_int_expected() ), pytest.param(_generate_float_input(), [torch.float16], 0.002, _generate_float_expected()), pytest.param(_generate_float_input(), [torch.float32, torch.float64], 0.001, _generate_float_expected()), ], ) def test_distance_iou(self, test_input, dtypes, tolerance, expected): self._run_test(test_input, dtypes, tolerance, expected) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", [torch.float32, torch.half]) def test_distance_iou_loss(dtype, device): box1 = torch.tensor([-1, -1, 1, 1], dtype=dtype, device=device) box2 = torch.tensor([0, 0, 1, 1], dtype=dtype, device=device) box3 = torch.tensor([0, 1, 1, 2], dtype=dtype, device=device) box4 = torch.tensor([1, 1, 2, 2], dtype=dtype, device=device) box1s = torch.stack( [box2, box2], dim=0, ) box2s = torch.stack( [box3, box4], dim=0, ) def assert_distance_iou_loss(box1, box2, expected_output, reduction="none"): output = ops.distance_box_iou_loss(box1, box2, reduction=reduction) # TODO: When passing the dtype, the torch.half fails as usual. expected_output = torch.tensor(expected_output, device=device) tol = 1e-5 if dtype != torch.half else 1e-3 # torch.testing.assert_close(output, expected_output, rtol=tol, atol=tol) assert np.allclose(output.numpy(), expected_output.numpy(), rtol=tol, atol=tol) assert_distance_iou_loss(box1, box1, 0.0) assert_distance_iou_loss(box1, box2, 0.8125) assert_distance_iou_loss(box1, box3, 1.1923) assert_distance_iou_loss(box1, box4, 1.2500) assert_distance_iou_loss(box1s, box2s, 1.2250, reduction="mean") assert_distance_iou_loss(box1s, box2s, 2.4500, reduction="sum") class TestCompleteBoxIou(BoxTestBase): def _target_fn(self) -> Tuple[bool, Callable]: return (True, ops.complete_box_iou) def _generate_int_input() -> List[List[int]]: return [[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]] def _generate_int_expected() -> List[List[float]]: return [ [1.0000, 0.1875, -0.4444], [0.1875, 1.0000, -0.5625], [-0.4444, -0.5625, 1.0000], ] def _generate_float_input() -> List[List[float]]: return [ [285.3538, 185.5758, 1193.5110, 851.4551], [285.1472, 188.7374, 1192.4984, 851.0669], [279.2440, 197.9812, 1189.4746, 849.2019], ] def _generate_float_expected() -> List[List[float]]: return [[1.0, 0.9933, 0.9673], [0.9933, 1.0, 0.9737], [0.9673, 0.9737, 1.0]] @pytest.mark.parametrize( "test_input, dtypes, tolerance, expected", [ pytest.param( _generate_int_input(), [torch.int16, torch.int32, torch.int64], 1e-4, _generate_int_expected() ), pytest.param(_generate_float_input(), [torch.float32, torch.float64], 0.002, _generate_float_expected()), pytest.param(_generate_float_input(), [torch.float32, torch.float64], 0.001, _generate_float_expected()), ], ) def test_complete_iou(self, test_input: List, dtypes, tolerance: float, expected: List) -> None: self._run_test(test_input, dtypes, tolerance, expected) class TestMasksToBoxes: def test_masks_box(self): def masks_box_check(masks, expected, tolerance=1e-4): out = ops.masks_to_boxes(masks) assert out.dtype == torch.float assert np.allclose(out.numpy(), expected.numpy(), rtol=0.0, atol=tolerance) # Check for int type boxes. def _get_image(): assets_directory = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") mask_path = os.path.join(assets_directory, "masks.tiff") image = Image.open(mask_path) return image def _create_masks(image, masks): for index in range(image.n_frames): image.seek(index) frame = np.array(image) masks[index] = torch.tensor(frame) return masks expected = torch.tensor( [ [127, 2, 165, 40], [2, 50, 44, 92], [56, 63, 98, 100], [139, 68, 175, 104], [160, 112, 198, 145], [49, 138, 99, 182], [108, 148, 152, 213], ], dtype=torch.float, ) image = _get_image() for dtype in [torch.float16, torch.float32, torch.float64]: masks = torch.zeros((image.n_frames, image.height, image.width), dtype=dtype) masks = _create_masks(image, masks) masks_box_check(masks, expected) class TestStochasticDepth: @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("p", [0.2, 0.5, 0.8]) @pytest.mark.parametrize("mode", ["batch", "row"]) def test_stochastic_depth_random(self, seed, mode, p): torch.manual_seed(seed) stats = pytest.importorskip("scipy.stats") batch_size = 5 x = torch.ones((batch_size, 3, 4, 4)) layer = ops.StochasticDepth(p=p, mode=mode) layer.__repr__() trials = 250 num_samples = 0 counts = 0 for _ in range(trials): out = layer(x) non_zero_count = out.sum(dim=(1, 2, 3)).nonzero().size(0) if mode == "batch": if non_zero_count == 0: counts += 1 num_samples += 1 elif mode == "row": counts += batch_size - non_zero_count num_samples += batch_size p_value = stats.binom_test(counts, num_samples, p=p) assert p_value > 0.01 @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("p", (0, 1)) @pytest.mark.parametrize("mode", ["batch", "row"]) def test_stochastic_depth(self, seed, mode, p): torch.manual_seed(seed) batch_size = 5 x = torch.ones((batch_size, 3, 4, 4)) layer = ops.StochasticDepth(p=p, mode=mode) out = layer(x) if p == 0: assert out.equal(x) elif p == 1: assert out.equal(torch.zeros_like(x)) def make_obj(self, p, mode, wrap=False): obj = ops.StochasticDepth(p, mode) return StochasticDepthWrapper(obj) if wrap else obj class TestDropBlock: @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("p", [0, 0.5]) @pytest.mark.parametrize("block_size", [5, 11]) @pytest.mark.parametrize("inplace", [True, False]) def test_drop_block(self, seed, dim, p, block_size, inplace): torch.manual_seed(seed) batch_size = 5 channels = 3 height = 11 width = height depth = height if dim == 2: x = torch.ones((batch_size, channels, height, width)) layer = ops.DropBlock2d(p=p, block_size=block_size, inplace=inplace) feature_size = height * width elif dim == 3: x = torch.ones((batch_size, channels, depth, height, width)) layer = ops.DropBlock3d(p=p, block_size=block_size, inplace=inplace) feature_size = depth * height * width layer.__repr__() out = layer(x) if p == 0: assert out.equal(x) if block_size == height: for b, c in product(range(batch_size), range(channels)): assert out[b, c].count_nonzero() in (0, feature_size) @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("p", [0.1, 0.2]) @pytest.mark.parametrize("block_size", [3]) @pytest.mark.parametrize("inplace", [False]) def test_drop_block_random(self, seed, dim, p, block_size, inplace): torch.manual_seed(seed) batch_size = 5 channels = 3 height = 11 width = height depth = height if dim == 2: x = torch.ones((batch_size, channels, height, width)) layer = ops.DropBlock2d(p=p, block_size=block_size, inplace=inplace) elif dim == 3: x = torch.ones((batch_size, channels, depth, height, width)) layer = ops.DropBlock3d(p=p, block_size=block_size, inplace=inplace) trials = 250 num_samples = 0 counts = 0 cell_numel = torch.tensor(x.shape).prod() for _ in range(trials): # with torch.no_grad(): # out = layer(x) out = layer(x) out = ms.ops.stop_gradient(out) non_zero_count = out.nonzero().size(0) counts += cell_numel - non_zero_count num_samples += cell_numel assert abs(p - counts / num_samples) / p < 0.15 def make_obj(self, dim, p, block_size, inplace, wrap=False): if dim == 2: obj = ops.DropBlock2d(p, block_size, inplace) elif dim == 3: obj = ops.DropBlock3d(p, block_size, inplace) return DropBlockWrapper(obj) if wrap else obj class TestFocalLoss: def _generate_diverse_input_target_pair(self, shape=(5, 2), **kwargs): def logit(p: Tensor) -> Tensor: return torch.log(p / (1 - p)) def generate_tensor_with_range_type(shape, range_type, **kwargs): if range_type != "random_binary": low, high = { "small": (0.0, 0.2), "big": (0.8, 1.0), "zeros": (0.0, 0.0), "ones": (1.0, 1.0), "random": (0.0, 1.0), }[range_type] # return torch.testing.make_tensor(shape, low=low, high=high, **kwargs) if low == high and low == 0: return torch.zeros(shape) elif low == high and low == 1: return torch.ones(shape) else: return torch.ones(shape).uniform_(low, high) else: return torch.randint(0, 2, shape, **kwargs) # This function will return inputs and targets with shape: (shape[0]*9, shape[1]) inputs = [] targets = [] for input_range_type, target_range_type in [ ("small", "zeros"), ("small", "ones"), ("small", "random_binary"), ("big", "zeros"), ("big", "ones"), ("big", "random_binary"), ("random", "zeros"), ("random", "ones"), ("random", "random_binary"), ]: inputs.append(logit(generate_tensor_with_range_type(shape, input_range_type, **kwargs))) targets.append(generate_tensor_with_range_type(shape, target_range_type, **kwargs)) return torch.cat(inputs), torch.cat(targets) @pytest.mark.parametrize("alpha", [-1.0, 0.0, 0.58, 1.0]) @pytest.mark.parametrize("gamma", [0, 2]) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", [torch.float32, torch.half]) @pytest.mark.parametrize("seed", [0, 1]) def test_correct_ratio(self, alpha, gamma, device, dtype, seed) -> None: if device == "cpu" and dtype is torch.half: pytest.skip("Currently torch.half is not fully supported on cpu") # For testing the ratio with manual calculation, we require the reduction to be "none" reduction = "none" torch.random.manual_seed(seed) inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device) focal_loss = ops.sigmoid_focal_loss(inputs, targets, gamma=gamma, alpha=alpha, reduction=reduction) ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction=reduction) assert torch.all( focal_loss <= ce_loss ), "focal loss must be less or equal to cross entropy loss with same input" loss_ratio = (focal_loss / ce_loss).squeeze() prob = torch.sigmoid(inputs) p_t = prob * targets + (1 - prob) * (1 - targets) correct_ratio = (1.0 - p_t) ** gamma if alpha >= 0: alpha_t = alpha * targets + (1 - alpha) * (1 - targets) correct_ratio = correct_ratio * alpha_t tol = 1e-3 if dtype is torch.half else 1e-5 # torch.testing.assert_close(correct_ratio, loss_ratio, rtol=tol, atol=tol) assert np.allclose(correct_ratio.numpy(), loss_ratio.numpy(), rtol=tol, atol=tol) @pytest.mark.parametrize("reduction", ["mean", "sum"]) @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", [torch.float32, torch.half]) @pytest.mark.parametrize("seed", [2, 3]) def test_equal_ce_loss(self, reduction, device, dtype, seed) -> None: if device == "cpu" and dtype is torch.half: pytest.skip("Currently torch.half is not fully supported on cpu") # focal loss should be equal ce_loss if alpha=-1 and gamma=0 alpha = -1 gamma = 0 torch.random.manual_seed(seed) inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device) inputs_fl = inputs.clone().requires_grad_() targets_fl = targets.clone() inputs_ce = inputs.clone().requires_grad_() targets_ce = targets.clone() focal_loss = ops.sigmoid_focal_loss(inputs_fl, targets_fl, gamma=gamma, alpha=alpha, reduction=reduction) ce_loss = F.binary_cross_entropy_with_logits(inputs_ce, targets_ce, reduction=reduction) tol = 1e-3 if dtype is torch.half else 1e-5 assert np.allclose(focal_loss.numpy(), ce_loss.numpy(), rtol=tol, atol=tol) class TestGeneralizedBoxIouLoss: # We refer to original test: https://github.com/facebookresearch/fvcore/blob/main/tests/test_giou_loss.py @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", [torch.float32, torch.half]) def test_giou_loss(self, dtype, device) -> None: box1 = torch.tensor([-1, -1, 1, 1], dtype=dtype, device=device) box2 = torch.tensor([0, 0, 1, 1], dtype=dtype, device=device) box3 = torch.tensor([0, 1, 1, 2], dtype=dtype, device=device) box4 = torch.tensor([1, 1, 2, 2], dtype=dtype, device=device) box1s = torch.stack([box2, box2], dim=0) box2s = torch.stack([box3, box4], dim=0) def assert_giou_loss(box1, box2, expected_loss, reduction="none"): tol = 1e-3 if dtype is torch.half else 1e-5 computed_loss = ops.generalized_box_iou_loss(box1, box2, reduction=reduction) expected_loss = torch.tensor(expected_loss, device=device) assert np.allclose(computed_loss.numpy(), expected_loss.numpy(), rtol=tol, atol=tol) # Identical boxes should have loss of 0 assert_giou_loss(box1, box1, 0.0) # quarter size box inside other box = IoU of 0.25 assert_giou_loss(box1, box2, 0.75) # Two side by side boxes, area=union # IoU=0 and GIoU=0 (loss 1.0) assert_giou_loss(box2, box3, 1.0) # Two diagonally adjacent boxes, area=2*union # IoU=0 and GIoU=-0.5 (loss 1.5) assert_giou_loss(box2, box4, 1.5) # Test batched loss and reductions assert_giou_loss(box1s, box2s, 2.5, reduction="sum") assert_giou_loss(box1s, box2s, 1.25, reduction="mean") @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", [torch.float32, torch.half]) def test_empty_inputs(self, dtype, device) -> None: box1 = torch.randn([0, 4], dtype=dtype).requires_grad_() box2 = torch.randn([0, 4], dtype=dtype).requires_grad_() loss = ops.generalized_box_iou_loss(box1, box2, reduction="mean") # loss.backward() tol = 1e-3 if dtype is torch.half else 1e-5 assert np.allclose(loss.numpy(), torch.tensor(0.0).numpy(), rtol=tol, atol=tol) class TestCIOULoss: @pytest.mark.parametrize("dtype", [torch.float32, torch.half]) @pytest.mark.parametrize("device", cpu_and_gpu()) def test_ciou_loss(self, dtype, device): box1 = torch.tensor([-1, -1, 1, 1], dtype=dtype, device=device) box2 = torch.tensor([0, 0, 1, 1], dtype=dtype, device=device) box3 = torch.tensor([0, 1, 1, 2], dtype=dtype, device=device) box4 = torch.tensor([1, 1, 2, 2], dtype=dtype, device=device) box1s = torch.stack([box2, box2], dim=0) box2s = torch.stack([box3, box4], dim=0) def assert_ciou_loss(box1, box2, expected_output, reduction="none"): output = ops.complete_box_iou_loss(box1, box2, reduction=reduction) # TODO: When passing the dtype, the torch.half test doesn't pass... expected_output = torch.tensor(expected_output, device=device) tol = 1e-5 if dtype != torch.half else 1e-3 # torch.testing.assert_close(output, expected_output, rtol=tol, atol=tol) assert np.allclose(output.numpy(), expected_output.numpy(), rtol=tol, atol=tol) assert_ciou_loss(box1, box1, 0.0) assert_ciou_loss(box1, box2, 0.8125) assert_ciou_loss(box1, box3, 1.1923) assert_ciou_loss(box1, box4, 1.2500) assert_ciou_loss(box1s, box2s, 1.2250, reduction="mean") assert_ciou_loss(box1s, box2s, 2.4500, reduction="sum") @pytest.mark.parametrize("device", cpu_and_gpu()) @pytest.mark.parametrize("dtype", [torch.float32, torch.half]) def test_empty_inputs(self, dtype, device) -> None: box1 = torch.randn([0, 4], dtype=dtype).requires_grad_() box2 = torch.randn([0, 4], dtype=dtype).requires_grad_() loss = ops.complete_box_iou_loss(box1, box2, reduction="mean") # loss.backward() tol = 1e-3 if dtype is torch.half else 1e-5 # torch.testing.assert_close(loss, torch.tensor(0.0), rtol=tol, atol=tol) assert np.allclose(loss.numpy(), torch.tensor(0.0).numpy(), rtol=tol, atol=tol) # assert box1.grad is not None, "box1.grad should not be None after backward is called" # assert box2.grad is not None, "box2.grad should not be None after backward is called" loss = ops.complete_box_iou_loss(box1, box2, reduction="none") assert loss.numel() == 0, "ciou_loss for two empty box should be empty" if __name__ == "__main__": pytest.main([__file__]) msadapter/testing/ut/torchvision/test_transforms.py000066400000000000000000002476221436713554100233470ustar00rootroot00000000000000import math import os import random import re from functools import partial import numpy as np import pytest # import torch # import torchvision.transforms as transforms # import torchvision.transforms._pil_constants as _pil_constants # import torchvision.transforms.functional as F # import torchvision.transforms.functional_tensor as F_t import ms_adapter.pytorch as torch import ms_adapter.torchvision.transforms as transforms import ms_adapter.torchvision.transforms._pil_constants as _pil_constants import ms_adapter.torchvision.transforms.functional as F import ms_adapter.torchvision.transforms.functional_tensor as F_t from ms_adapter.pytorch.tensor import cast_to_ms_tensor, cast_to_adapter_tensor from PIL import Image try: import accimage except ImportError: accimage = None try: from scipy import stats except ImportError: stats = None from common_utils import cycle_over, int_dtypes, float_dtypes GRACE_HOPPER = os.path.join( os.path.dirname(os.path.abspath(__file__)), "assets", "encode_jpeg", "grace_hopper_517x606.jpg" ) def _get_grayscale_test_image(img, fill=None): img = img.convert("L") fill = (fill[0],) if isinstance(fill, tuple) else fill return img, fill class TestConvertImageDtype: @pytest.mark.parametrize("input_dtype, output_dtype", cycle_over(float_dtypes())) def test_float_to_float(self, input_dtype, output_dtype): input_image = torch.tensor((0.0, 1.0), dtype=input_dtype) transform = transforms.ConvertImageDtype(output_dtype) # transform_script = torch.jit.script(F.convert_image_dtype) output_image = transform(input_image) # output_image_script = transform_script(input_image, output_dtype) # torch.testing.assert_close(output_image_script, output_image, rtol=0.0, atol=1e-6) actual_min, actual_max = output_image.tolist() desired_min, desired_max = 0.0, 1.0 assert abs(actual_min - desired_min) < 1e-7 assert abs(actual_max - desired_max) < 1e-7 @pytest.mark.parametrize("input_dtype", float_dtypes()) @pytest.mark.parametrize("output_dtype", int_dtypes()) def test_float_to_int(self, input_dtype, output_dtype): input_image = torch.tensor((0.0, 1.0), dtype=input_dtype) transform = transforms.ConvertImageDtype(output_dtype) # transform_script = torch.jit.script(F.convert_image_dtype) if (input_dtype == torch.float32 and output_dtype in (torch.int32, torch.int64)) or ( input_dtype == torch.float64 and output_dtype == torch.int64 ): with pytest.raises(RuntimeError): transform(input_image) else: output_image = transform(input_image) # output_image_script = transform_script(input_image, output_dtype) # torch.testing.assert_close(output_image_script, output_image, rtol=0.0, atol=1e-6) actual_min, actual_max = output_image.tolist() desired_min, desired_max = 0, torch.iinfo(output_dtype).max assert actual_min == desired_min assert actual_max == desired_max @pytest.mark.parametrize("input_dtype", int_dtypes()) @pytest.mark.parametrize("output_dtype", float_dtypes()) def test_int_to_float(self, input_dtype, output_dtype): input_image = torch.tensor((0, torch.iinfo(input_dtype).max), dtype=input_dtype) transform = transforms.ConvertImageDtype(output_dtype) # transform_script = torch.jit.script(F.convert_image_dtype) output_image = transform(input_image) # output_image_script = transform_script(input_image, output_dtype) # torch.testing.assert_close(output_image_script, output_image, rtol=0.0, atol=1e-6) actual_min, actual_max = output_image.tolist() desired_min, desired_max = 0.0, 1.0 assert abs(actual_min - desired_min) < 1e-7 assert actual_min >= desired_min assert abs(actual_max - desired_max) < 1e-7 assert actual_max <= desired_max @pytest.mark.parametrize("input_dtype, output_dtype", cycle_over(int_dtypes())) def test_dtype_int_to_int(self, input_dtype, output_dtype): input_max = torch.iinfo(input_dtype).max input_image = torch.tensor((0, input_max), dtype=input_dtype) output_max = torch.iinfo(output_dtype).max transform = transforms.ConvertImageDtype(output_dtype) # transform_script = torch.jit.script(F.convert_image_dtype) output_image = transform(input_image) # output_image_script = transform_script(input_image, output_dtype) # torch.testing.assert_close( # output_image_script, # output_image, # rtol=0.0, # atol=1e-6, # msg=f"{output_image_script} vs {output_image}", # ) actual_min, actual_max = output_image.tolist() desired_min, desired_max = 0, output_max # see https://github.com/pytorch/vision/pull/2078#issuecomment-641036236 for details if input_max >= output_max: error_term = 0 else: error_term = 1 - (torch.iinfo(output_dtype).max + 1) // (torch.iinfo(input_dtype).max + 1) assert actual_min == desired_min assert actual_max == (desired_max + error_term) @pytest.mark.parametrize("input_dtype, output_dtype", cycle_over(int_dtypes())) def test_int_to_int_consistency(self, input_dtype, output_dtype): input_max = torch.iinfo(input_dtype).max input_image = torch.tensor((0, input_max), dtype=input_dtype) output_max = torch.iinfo(output_dtype).max if output_max <= input_max: return transform = transforms.ConvertImageDtype(output_dtype) inverse_transfrom = transforms.ConvertImageDtype(input_dtype) output_image = inverse_transfrom(transform(input_image)) actual_min, actual_max = output_image.tolist() desired_min, desired_max = 0, input_max assert actual_min == desired_min assert actual_max == desired_max class TestToTensor: @pytest.mark.parametrize("channels", [1, 3, 4]) def test_to_tensor(self, channels): height, width = 4, 4 trans = transforms.ToTensor() np_rng = np.random.RandomState(0) input_data = torch.ByteTensor(channels, height, width).random_adapter(0, 255).float().div(255) img = transforms.ToPILImage()(input_data) output = trans(img) # torch.testing.assert_close(output, input_data) assert np.allclose(output.numpy(), input_data.numpy()) ndarray = np_rng.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8) output = trans(ndarray) expected_output = ndarray.transpose((2, 0, 1)) / 255.0 # torch.testing.assert_close(output.numpy(), expected_output, check_dtype=False) assert np.allclose(output.numpy(), expected_output) ndarray = np_rng.rand(height, width, channels).astype(np.float32) output = trans(ndarray) expected_output = ndarray.transpose((2, 0, 1)) # torch.testing.assert_close(output.numpy(), expected_output, check_dtype=False) assert np.allclose(output.numpy(), expected_output) # separate test for mode '1' PIL images # TODO: input_data = torch.ByteTensor(1, height, width).bernoulli_adapter() input_data = torch.ByteTensor(1, height, width) img = transforms.ToPILImage()(input_data.mul(255)).convert("1") output = trans(img) # torch.testing.assert_close(input_data, output, check_dtype=False) assert np.allclose(input_data.numpy(), output.numpy()) def test_to_tensor_errors(self): height, width = 4, 4 trans = transforms.ToTensor() np_rng = np.random.RandomState(0) with pytest.raises(TypeError): trans(np_rng.rand(1, height, width).tolist()) with pytest.raises(ValueError): trans(np_rng.rand(height)) with pytest.raises(ValueError): trans(np_rng.rand(1, 1, height, width)) # TODO: # @pytest.mark.parametrize("dtype", [torch.float16, torch.float, torch.double]) # def test_to_tensor_with_other_default_dtypes(self, dtype): # np_rng = np.random.RandomState(0) # # # TODO: current_def_dtype = torch.get_default_dtype() # current_def_dtype = torch.float # # t = transforms.ToTensor() # np_arr = np_rng.randint(0, 255, (32, 32, 3), dtype=np.uint8) # img = Image.fromarray(np_arr) # # # torch.set_default_dtype(dtype) # res = t(img) # assert res.dtype == dtype, f"{res.dtype} vs {dtype}" # # torch.set_default_dtype(current_def_dtype) @pytest.mark.parametrize("channels", [1, 3, 4]) def test_pil_to_tensor(self, channels): height, width = 4, 4 trans = transforms.PILToTensor() np_rng = np.random.RandomState(0) input_data = torch.ByteTensor(channels, height, width).random_adapter(0, 255) img = transforms.ToPILImage()(input_data) output = trans(img) # torch.testing.assert_close(input_data, output) assert np.allclose(input_data.numpy(), output.numpy()) input_data = np_rng.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8) img = transforms.ToPILImage()(input_data) output = trans(img) expected_output = input_data.transpose((2, 0, 1)) # torch.testing.assert_close(output.numpy(), expected_output) assert np.allclose(output.numpy(), expected_output) input_data = torch.as_tensor(np_rng.rand(channels, height, width).astype(np.float32)) img = transforms.ToPILImage()(input_data) # CHW -> HWC and (* 255).byte() output = trans(img) # HWC -> CHW expected_output = (input_data * 255).byte() # torch.testing.assert_close(output, expected_output) assert np.allclose(output.numpy(), expected_output.numpy()) # separate test for mode '1' PIL images # TODO: input_data = torch.ByteTensor(1, height, width).bernoulli_adapter() input_data = torch.ByteTensor(1, height, width) img = transforms.ToPILImage()(input_data.mul(255)).convert("1") output = trans(img).to(torch.uint8).bool().to(torch.uint8) # torch.testing.assert_close(input_data, output) assert np.allclose(input_data.numpy(), output.numpy()) def test_pil_to_tensor_errors(self): height, width = 4, 4 trans = transforms.PILToTensor() np_rng = np.random.RandomState(0) with pytest.raises(TypeError): trans(np_rng.rand(1, height, width).tolist()) with pytest.raises(TypeError): trans(np_rng.rand(1, height, width)) def test_randomresized_params(): height = random.randint(24, 32) * 2 width = random.randint(24, 32) * 2 img = torch.ones(3, height, width) to_pil_image = transforms.ToPILImage() img = to_pil_image(img) size = 100 epsilon = 0.05 min_scale = 0.25 for _ in range(10): scale_min = max(round(random.random(), 2), min_scale) scale_range = (scale_min, scale_min + round(random.random(), 2)) aspect_min = max(round(random.random(), 2), epsilon) aspect_ratio_range = (aspect_min, aspect_min + round(random.random(), 2)) randresizecrop = transforms.RandomResizedCrop(size, scale_range, aspect_ratio_range) i, j, h, w = randresizecrop.get_params(img, scale_range, aspect_ratio_range) aspect_ratio_obtained = w / h assert ( min(aspect_ratio_range) - epsilon <= aspect_ratio_obtained and aspect_ratio_obtained <= max(aspect_ratio_range) + epsilon ) or aspect_ratio_obtained == 1.0 assert isinstance(i, int) assert isinstance(j, int) assert isinstance(h, int) assert isinstance(w, int) @pytest.mark.parametrize( "height, width", [ # height, width # square image (28, 28), (27, 27), # rectangular image: h < w (28, 34), (29, 35), # rectangular image: h > w (34, 28), (35, 29), ], ) @pytest.mark.parametrize( "osize", [ # single integer 22, 27, 28, 36, # single integer in tuple/list [ 22, ], (27,), ], ) @pytest.mark.parametrize("max_size", (None, 37, 1000)) def test_resize(height, width, osize, max_size): img = Image.new("RGB", size=(width, height), color=127) t = transforms.Resize(osize, max_size=max_size) result = t(img) msg = f"{height}, {width} - {osize} - {max_size}" osize = osize[0] if isinstance(osize, (list, tuple)) else osize # If size is an int, smaller edge of the image will be matched to this number. # i.e, if height > width, then image will be rescaled to (size * height / width, size). if height < width: exp_w, exp_h = (int(osize * width / height), osize) # (w, h) if max_size is not None and max_size < exp_w: exp_w, exp_h = max_size, int(max_size * exp_h / exp_w) assert result.size == (exp_w, exp_h), msg elif width < height: exp_w, exp_h = (osize, int(osize * height / width)) # (w, h) if max_size is not None and max_size < exp_h: exp_w, exp_h = int(max_size * exp_w / exp_h), max_size assert result.size == (exp_w, exp_h), msg else: exp_w, exp_h = (osize, osize) # (w, h) if max_size is not None and max_size < osize: exp_w, exp_h = max_size, max_size assert result.size == (exp_w, exp_h), msg @pytest.mark.parametrize( "height, width", [ # height, width # square image (28, 28), (27, 27), # rectangular image: h < w (28, 34), (29, 35), # rectangular image: h > w (34, 28), (35, 29), ], ) @pytest.mark.parametrize( "osize", [ # two integers sequence output [22, 22], [22, 28], [22, 36], [27, 22], [36, 22], [28, 28], [28, 37], [37, 27], [37, 37], ], ) def test_resize_sequence_output(height, width, osize): img = Image.new("RGB", size=(width, height), color=127) oheight, owidth = osize t = transforms.Resize(osize) result = t(img) assert (owidth, oheight) == result.size def test_resize_antialias_error(): osize = [37, 37] img = Image.new("RGB", size=(35, 29), color=127) with pytest.warns(UserWarning, match=r"Anti-alias option is always applied for PIL Image input"): t = transforms.Resize(osize, antialias=False) t(img) @pytest.mark.parametrize("height, width", ((32, 64), (64, 32))) def test_resize_size_equals_small_edge_size(height, width): # Non-regression test for https://github.com/pytorch/vision/issues/5405 # max_size used to be ignored if size == small_edge_size max_size = 40 img = Image.new("RGB", size=(width, height), color=127) small_edge = min(height, width) t = transforms.Resize(small_edge, max_size=max_size) result = t(img) assert max(result.size) == max_size class TestPad: @pytest.mark.parametrize("fill", [85, 85.0]) def test_pad(self, fill): height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 img = torch.ones(3, height, width, dtype=torch.uint8) padding = random.randint(1, 20) result = transforms.Compose( [ transforms.ToPILImage(), transforms.Pad(padding, fill=fill), transforms.PILToTensor(), ] )(img) assert result.size(1) == height + 2 * padding assert result.size(2) == width + 2 * padding # check that all elements in the padded region correspond # to the pad value h_padded = result[:, :padding, :] w_padded = result[:, :, :padding] # torch.testing.assert_close(h_padded, torch.full_like(h_padded, fill_value=fill), rtol=0.0, atol=0.0) # torch.testing.assert_close(w_padded, torch.full_like(w_padded, fill_value=fill), rtol=0.0, atol=0.0) assert np.allclose(h_padded.numpy(), torch.full_like(h_padded, fill_value=fill).numpy(), rtol=0.0, atol=0.0) assert np.allclose(w_padded.numpy(), torch.full_like(w_padded, fill_value=fill).numpy(), rtol=0.0, atol=0.0) pytest.raises(ValueError, transforms.Pad(padding, fill=(1, 2)), transforms.ToPILImage()(img)) def test_pad_with_tuple_of_pad_values(self): height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 img = transforms.ToPILImage()(torch.ones(3, height, width)) padding = tuple(random.randint(1, 20) for _ in range(2)) output = transforms.Pad(padding)(img) assert output.size == (width + padding[0] * 2, height + padding[1] * 2) padding = [random.randint(1, 20) for _ in range(4)] output = transforms.Pad(padding)(img) assert output.size[0] == width + padding[0] + padding[2] assert output.size[1] == height + padding[1] + padding[3] # Checking if Padding can be printed as string transforms.Pad(padding).__repr__() def test_pad_with_non_constant_padding_modes(self): """Unit tests for edge, reflect, symmetric padding""" img = torch.zeros(3, 27, 27).byte() # TODO: # img[:, :, 0] = 1 # Constant value added to leftmost edge img_dtype = img.dtype img = img.to(torch.float32) img[:, :, 0] = 1.0 img = img.to(img_dtype) img = transforms.ToPILImage()(img) img = F.pad(img, 1, (200, 200, 200)) # pad 3 to all sidess edge_padded_img = F.pad(img, 3, padding_mode="edge") # First 6 elements of leftmost edge in the middle of the image, values are in order: # edge_pad, edge_pad, edge_pad, constant_pad, constant value added to leftmost edge, 0 edge_middle_slice = np.asarray(edge_padded_img).transpose(2, 0, 1)[0][17][:6] # assert_equal(edge_middle_slice, np.asarray([200, 200, 200, 200, 1, 0], dtype=np.uint8)) assert np.allclose(edge_middle_slice, np.asarray([200, 200, 200, 200, 1, 0], dtype=np.uint8)) assert transforms.PILToTensor()(edge_padded_img).size() == (3, 35, 35) # Pad 3 to left/right, 2 to top/bottom reflect_padded_img = F.pad(img, (3, 2), padding_mode="reflect") # First 6 elements of leftmost edge in the middle of the image, values are in order: # reflect_pad, reflect_pad, reflect_pad, constant_pad, constant value added to leftmost edge, 0 reflect_middle_slice = np.asarray(reflect_padded_img).transpose(2, 0, 1)[0][17][:6] # assert_equal(reflect_middle_slice, np.asarray([0, 0, 1, 200, 1, 0], dtype=np.uint8)) assert np.allclose(reflect_middle_slice, np.asarray([0, 0, 1, 200, 1, 0], dtype=np.uint8)) assert transforms.PILToTensor()(reflect_padded_img).size() == (3, 33, 35) # Pad 3 to left, 2 to top, 2 to right, 1 to bottom symmetric_padded_img = F.pad(img, (3, 2, 2, 1), padding_mode="symmetric") # First 6 elements of leftmost edge in the middle of the image, values are in order: # sym_pad, sym_pad, sym_pad, constant_pad, constant value added to leftmost edge, 0 symmetric_middle_slice = np.asarray(symmetric_padded_img).transpose(2, 0, 1)[0][17][:6] # assert_equal(symmetric_middle_slice, np.asarray([0, 1, 200, 200, 1, 0], dtype=np.uint8)) assert np.allclose(symmetric_middle_slice, np.asarray([0, 1, 200, 200, 1, 0], dtype=np.uint8)) assert transforms.PILToTensor()(symmetric_padded_img).size() == (3, 32, 34) # Check negative padding explicitly for symmetric case, since it is not # implemented for tensor case to compare to # Crop 1 to left, pad 2 to top, pad 3 to right, crop 3 to bottom symmetric_padded_img_neg = F.pad(img, (-1, 2, 3, -3), padding_mode="symmetric") symmetric_neg_middle_left = np.asarray(symmetric_padded_img_neg).transpose(2, 0, 1)[0][17][:3] symmetric_neg_middle_right = np.asarray(symmetric_padded_img_neg).transpose(2, 0, 1)[0][17][-4:] # assert_equal(symmetric_neg_middle_left, np.asarray([1, 0, 0], dtype=np.uint8)) # assert_equal(symmetric_neg_middle_right, np.asarray([200, 200, 0, 0], dtype=np.uint8)) assert np.allclose(symmetric_neg_middle_left, np.asarray([1, 0, 0], dtype=np.uint8)) assert np.allclose(symmetric_neg_middle_right, np.asarray([200, 200, 0, 0], dtype=np.uint8)) assert transforms.PILToTensor()(symmetric_padded_img_neg).size() == (3, 28, 31) def test_pad_raises_with_invalid_pad_sequence_len(self): with pytest.raises(ValueError): transforms.Pad(()) with pytest.raises(ValueError): transforms.Pad((1, 2, 3)) with pytest.raises(ValueError): transforms.Pad((1, 2, 3, 4, 5)) def test_pad_with_mode_F_images(self): pad = 2 transform = transforms.Pad(pad) img = Image.new("F", (10, 10)) padded_img = transform(img) # assert_equal(padded_img.size, [edge_size + 2 * pad for edge_size in img.size]) assert np.allclose(padded_img.size, [edge_size + 2 * pad for edge_size in img.size]) #??? @pytest.mark.parametrize( "fn, trans, kwargs", [ (F.invert, transforms.RandomInvert, {}), (F.posterize, transforms.RandomPosterize, {"bits": 4}), (F.solarize, transforms.RandomSolarize, {"threshold": 192}), (F.adjust_sharpness, transforms.RandomAdjustSharpness, {"sharpness_factor": 2.0}), (F.autocontrast, transforms.RandomAutocontrast, {}), (F.equalize, transforms.RandomEqualize, {}), (F.vflip, transforms.RandomVerticalFlip, {}), (F.hflip, transforms.RandomHorizontalFlip, {}), (partial(F.to_grayscale, num_output_channels=3), transforms.RandomGrayscale, {}), ], ) @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("p", (0, 1)) def test_randomness(fn, trans, kwargs, seed, p): torch.manual_seed(seed) img = transforms.ToPILImage()(torch.rand(3, 16, 18)) expected_transformed_img = fn(img, **kwargs) randomly_transformed_img = trans(p=p, **kwargs)(img) if p == 0: assert randomly_transformed_img == img elif p == 1: assert randomly_transformed_img == expected_transformed_img trans(**kwargs).__repr__() def test_autocontrast_equal_minmax(): img_tensor = torch.tensor([[[10]], [[128]], [[245]]], dtype=torch.uint8).expand(3, 32, 32) img_pil = F.to_pil_image(img_tensor) img_tensor = F.autocontrast(img_tensor) img_pil = F.autocontrast(img_pil) # torch.testing.assert_close(img_tensor, F.pil_to_tensor(img_pil)) assert np.allclose(img_tensor.numpy(), F.pil_to_tensor(img_pil).numpy()) class TestToPil: def _get_1_channel_tensor_various_types(): img_data_float = torch.Tensor(1, 4, 4).uniform_adpater() expected_output = img_data_float.mul(255).int().float().div(255).numpy() yield img_data_float, expected_output, "L" img_data_byte = torch.ByteTensor(1, 4, 4).random_adapter(0, 255) expected_output = img_data_byte.float().div(255.0).numpy() yield img_data_byte, expected_output, "L" img_data_short = torch.ShortTensor(1, 4, 4).random_adapter() expected_output = img_data_short.numpy() yield img_data_short, expected_output, "I;16" img_data_int = torch.IntTensor(1, 4, 4).random_adapter() expected_output = img_data_int.numpy() yield img_data_int, expected_output, "I" def _get_2d_tensor_various_types(): img_data_float = torch.Tensor(4, 4).uniform_adpater() expected_output = img_data_float.mul(255).int().float().div(255).numpy() yield img_data_float, expected_output, "L" img_data_byte = torch.ByteTensor(4, 4).random_adapter(0, 255) expected_output = img_data_byte.float().div(255.0).numpy() yield img_data_byte, expected_output, "L" img_data_short = torch.ShortTensor(4, 4).random_adapter() expected_output = img_data_short.numpy() yield img_data_short, expected_output, "I;16" img_data_int = torch.IntTensor(4, 4).random_adapter() expected_output = img_data_int.numpy() yield img_data_int, expected_output, "I" @pytest.mark.parametrize("with_mode", [False, True]) @pytest.mark.parametrize("img_data, expected_output, expected_mode", _get_1_channel_tensor_various_types()) def test_1_channel_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode): transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() to_tensor = transforms.ToTensor() img = transform(img_data) assert img.mode == expected_mode # torch.testing.assert_close(expected_output, to_tensor(img).numpy()) assert np.allclose(expected_output, to_tensor(img).numpy()) def test_1_channel_float_tensor_to_pil_image(self): img_data = torch.Tensor(1, 4, 4).uniform_adpater() # 'F' mode for torch.FloatTensor img_F_mode = transforms.ToPILImage(mode="F")(img_data) assert img_F_mode.mode == "F" # torch.testing.assert_close( # np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode="F")), np.array(img_F_mode) # ) assert np.allclose(np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode="F")), np.array(img_F_mode)) @pytest.mark.parametrize("with_mode", [False, True]) @pytest.mark.parametrize( "img_data, expected_mode", [ (torch.Tensor(4, 4, 1).uniform_adpater().numpy(), "F"), (torch.ByteTensor(4, 4, 1).random_adapter(0, 255).numpy(), "L"), (torch.ShortTensor(4, 4, 1).random_adapter().numpy(), "I;16"), (torch.IntTensor(4, 4, 1).random_adapter().numpy(), "I"), ], ) def test_1_channel_ndarray_to_pil_image(self, with_mode, img_data, expected_mode): transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() img = transform(img_data) assert img.mode == expected_mode # note: we explicitly convert img's dtype because pytorch doesn't support uint16 # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array # torch.testing.assert_close(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype)) assert np.allclose(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype)) @pytest.mark.parametrize("expected_mode", [None, "LA"]) def test_2_channel_ndarray_to_pil_image(self, expected_mode): img_data = torch.ByteTensor(4, 4, 2).random_adapter(0, 255).numpy() if expected_mode is None: img = transforms.ToPILImage()(img_data) assert img.mode == "LA" # default should assume LA else: img = transforms.ToPILImage(mode=expected_mode)(img_data) assert img.mode == expected_mode split = img.split() for i in range(2): # torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i])) assert np.allclose(img_data[:, :, i], np.asarray(split[i])) def test_2_channel_ndarray_to_pil_image_error(self): img_data = torch.ByteTensor(4, 4, 2).random_adapter(0, 255).numpy() transforms.ToPILImage().__repr__() # should raise if we try a mode for 4 or 1 or 3 channel images with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): transforms.ToPILImage(mode="RGBA")(img_data) with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): transforms.ToPILImage(mode="P")(img_data) with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): transforms.ToPILImage(mode="RGB")(img_data) @pytest.mark.parametrize("expected_mode", [None, "LA"]) def test_2_channel_tensor_to_pil_image(self, expected_mode): img_data = torch.Tensor(2, 4, 4).uniform_adpater() expected_output = img_data.mul(255).int().float().div(255) if expected_mode is None: img = transforms.ToPILImage()(img_data) assert img.mode == "LA" # default should assume LA else: img = transforms.ToPILImage(mode=expected_mode)(img_data) assert img.mode == expected_mode split = img.split() for i in range(2): # torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) assert np.allclose(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) def test_2_channel_tensor_to_pil_image_error(self): img_data = torch.Tensor(2, 4, 4).uniform_adpater() # should raise if we try a mode for 4 or 1 or 3 channel images with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): transforms.ToPILImage(mode="RGBA")(img_data) with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): transforms.ToPILImage(mode="P")(img_data) with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): transforms.ToPILImage(mode="RGB")(img_data) @pytest.mark.parametrize("with_mode", [False, True]) @pytest.mark.parametrize("img_data, expected_output, expected_mode", _get_2d_tensor_various_types()) def test_2d_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode): transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() to_tensor = transforms.ToTensor() img = transform(img_data) assert img.mode == expected_mode # torch.testing.assert_close(expected_output, to_tensor(img).numpy()[0]) assert np.allclose(expected_output, to_tensor(img).numpy()[0]) @pytest.mark.parametrize("with_mode", [False, True]) @pytest.mark.parametrize( "img_data, expected_mode", [ (torch.Tensor(4, 4).uniform_adpater().numpy(), "F"), (torch.ByteTensor(4, 4).random_adapter(0, 255).numpy(), "L"), (torch.ShortTensor(4, 4).random_adapter().numpy(), "I;16"), (torch.IntTensor(4, 4).random_adapter().numpy(), "I"), ], ) def test_2d_ndarray_to_pil_image(self, with_mode, img_data, expected_mode): transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() img = transform(img_data) assert img.mode == expected_mode np.testing.assert_allclose(img_data, img) @pytest.mark.parametrize("expected_mode", [None, "RGB", "HSV", "YCbCr"]) def test_3_channel_tensor_to_pil_image(self, expected_mode): img_data = torch.Tensor(3, 4, 4).uniform_adpater() expected_output = img_data.mul(255).int().float().div(255) if expected_mode is None: img = transforms.ToPILImage()(img_data) assert img.mode == "RGB" # default should assume RGB else: img = transforms.ToPILImage(mode=expected_mode)(img_data) assert img.mode == expected_mode split = img.split() for i in range(3): # torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) assert np.allclose(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) def test_3_channel_tensor_to_pil_image_error(self): img_data = torch.Tensor(3, 4, 4).uniform_adpater() error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs" # should raise if we try a mode for 4 or 1 or 2 channel images with pytest.raises(ValueError, match=error_message_3d): transforms.ToPILImage(mode="RGBA")(img_data) with pytest.raises(ValueError, match=error_message_3d): transforms.ToPILImage(mode="P")(img_data) with pytest.raises(ValueError, match=error_message_3d): transforms.ToPILImage(mode="LA")(img_data) with pytest.raises(ValueError, match=r"pic should be 2/3 dimensional. Got \d+ dimensions."): transforms.ToPILImage()(torch.Tensor(1, 3, 4, 4).uniform_adpater()) @pytest.mark.parametrize("expected_mode", [None, "RGB", "HSV", "YCbCr"]) def test_3_channel_ndarray_to_pil_image(self, expected_mode): img_data = torch.ByteTensor(4, 4, 3).random_adapter(0, 255).numpy() if expected_mode is None: img = transforms.ToPILImage()(img_data) assert img.mode == "RGB" # default should assume RGB else: img = transforms.ToPILImage(mode=expected_mode)(img_data) assert img.mode == expected_mode split = img.split() for i in range(3): # torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i])) assert np.allclose(img_data[:, :, i], np.asarray(split[i])) def test_3_channel_ndarray_to_pil_image_error(self): img_data = torch.ByteTensor(4, 4, 3).random_adapter(0, 255).numpy() # Checking if ToPILImage can be printed as string transforms.ToPILImage().__repr__() error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs" # should raise if we try a mode for 4 or 1 or 2 channel images with pytest.raises(ValueError, match=error_message_3d): transforms.ToPILImage(mode="RGBA")(img_data) with pytest.raises(ValueError, match=error_message_3d): transforms.ToPILImage(mode="P")(img_data) with pytest.raises(ValueError, match=error_message_3d): transforms.ToPILImage(mode="LA")(img_data) @pytest.mark.parametrize("expected_mode", [None, "RGBA", "CMYK", "RGBX"]) def test_4_channel_tensor_to_pil_image(self, expected_mode): img_data = torch.Tensor(4, 4, 4).uniform_adpater() expected_output = img_data.mul(255).int().float().div(255) if expected_mode is None: img = transforms.ToPILImage()(img_data) assert img.mode == "RGBA" # default should assume RGBA else: img = transforms.ToPILImage(mode=expected_mode)(img_data) assert img.mode == expected_mode split = img.split() for i in range(4): # torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) assert np.allclose(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) def test_4_channel_tensor_to_pil_image_error(self): img_data = torch.Tensor(4, 4, 4).uniform_adpater() error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs" # should raise if we try a mode for 3 or 1 or 2 channel images with pytest.raises(ValueError, match=error_message_4d): transforms.ToPILImage(mode="RGB")(img_data) with pytest.raises(ValueError, match=error_message_4d): transforms.ToPILImage(mode="P")(img_data) with pytest.raises(ValueError, match=error_message_4d): transforms.ToPILImage(mode="LA")(img_data) @pytest.mark.parametrize("expected_mode", [None, "RGBA", "CMYK", "RGBX"]) def test_4_channel_ndarray_to_pil_image(self, expected_mode): img_data = torch.ByteTensor(4, 4, 4).random_adapter(0, 255).numpy() if expected_mode is None: img = transforms.ToPILImage()(img_data) assert img.mode == "RGBA" # default should assume RGBA else: img = transforms.ToPILImage(mode=expected_mode)(img_data) assert img.mode == expected_mode split = img.split() for i in range(4): # torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i])) assert np.allclose(img_data[:, :, i], np.asarray(split[i])) def test_4_channel_ndarray_to_pil_image_error(self): img_data = torch.ByteTensor(4, 4, 4).random_adapter(0, 255).numpy() error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs" # should raise if we try a mode for 3 or 1 or 2 channel images with pytest.raises(ValueError, match=error_message_4d): transforms.ToPILImage(mode="RGB")(img_data) with pytest.raises(ValueError, match=error_message_4d): transforms.ToPILImage(mode="P")(img_data) with pytest.raises(ValueError, match=error_message_4d): transforms.ToPILImage(mode="LA")(img_data) def test_ndarray_bad_types_to_pil_image(self): trans = transforms.ToPILImage() reg_msg = r"Input type \w+ is not supported" with pytest.raises(TypeError, match=reg_msg): trans(np.ones([4, 4, 1], np.int64)) with pytest.raises(TypeError, match=reg_msg): trans(np.ones([4, 4, 1], np.uint16)) with pytest.raises(TypeError, match=reg_msg): trans(np.ones([4, 4, 1], np.uint32)) with pytest.raises(TypeError, match=reg_msg): trans(np.ones([4, 4, 1], np.float64)) with pytest.raises(ValueError, match=r"pic should be 2/3 dimensional. Got \d+ dimensions."): transforms.ToPILImage()(np.ones([1, 4, 4, 3])) with pytest.raises(ValueError, match=r"pic should not have > 4 channels. Got \d+ channels."): transforms.ToPILImage()(np.ones([4, 4, 6])) def test_tensor_bad_types_to_pil_image(self): with pytest.raises(ValueError, match=r"pic should be 2/3 dimensional. Got \d+ dimensions."): transforms.ToPILImage()(torch.ones(1, 3, 4, 4)) with pytest.raises(ValueError, match=r"pic should not have > 4 channels. Got \d+ channels."): transforms.ToPILImage()(torch.ones(6, 4, 4)) def test_adjust_brightness(): x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") # test 0 y_pil = F.adjust_brightness(x_pil, 1) y_np = np.array(y_pil) # torch.testing.assert_close(y_np, x_np) assert np.allclose(y_np, x_np) # test 1 y_pil = F.adjust_brightness(x_pil, 0.5) y_np = np.array(y_pil) y_ans = [0, 2, 6, 27, 67, 113, 18, 4, 117, 45, 127, 0] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 2 y_pil = F.adjust_brightness(x_pil, 2) y_np = np.array(y_pil) y_ans = [0, 10, 26, 108, 255, 255, 74, 16, 255, 180, 255, 2] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) def test_adjust_contrast(): x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") # test 0 y_pil = F.adjust_contrast(x_pil, 1) y_np = np.array(y_pil) # torch.testing.assert_close(y_np, x_np) assert np.allclose(y_np, x_np) # test 1 y_pil = F.adjust_contrast(x_pil, 0.5) y_np = np.array(y_pil) y_ans = [43, 45, 49, 70, 110, 156, 61, 47, 160, 88, 170, 43] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 2 y_pil = F.adjust_contrast(x_pil, 2) y_np = np.array(y_pil) y_ans = [0, 0, 0, 22, 184, 255, 0, 0, 255, 94, 255, 0] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) @pytest.mark.skipif(Image.__version__ >= "7", reason="Temporarily disabled") def test_adjust_saturation(): x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") # test 0 y_pil = F.adjust_saturation(x_pil, 1) y_np = np.array(y_pil) # torch.testing.assert_close(y_np, x_np) assert np.allclose(y_np, x_np) # test 1 y_pil = F.adjust_saturation(x_pil, 0.5) y_np = np.array(y_pil) y_ans = [2, 4, 8, 87, 128, 173, 39, 25, 138, 133, 215, 88] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 2 y_pil = F.adjust_saturation(x_pil, 2) y_np = np.array(y_pil) y_ans = [0, 6, 22, 0, 149, 255, 32, 0, 255, 4, 255, 0] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) def test_adjust_hue(): x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") with pytest.raises(ValueError): F.adjust_hue(x_pil, -0.7) F.adjust_hue(x_pil, 1) # test 0: almost same as x_data but not exact. # probably because hsv <-> rgb floating point ops y_pil = F.adjust_hue(x_pil, 0) y_np = np.array(y_pil) y_ans = [0, 5, 13, 54, 139, 226, 35, 8, 234, 91, 255, 1] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 1 y_pil = F.adjust_hue(x_pil, 0.25) y_np = np.array(y_pil) y_ans = [13, 0, 12, 224, 54, 226, 234, 8, 99, 1, 222, 255] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 2 y_pil = F.adjust_hue(x_pil, -0.25) y_np = np.array(y_pil) y_ans = [0, 13, 2, 54, 226, 58, 8, 234, 152, 255, 43, 1] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) def test_adjust_sharpness(): x_shape = [4, 4, 3] x_data = [ 75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 0, 0, 65, 108, 101, 120, 97, 110, 100, 101, 114, 32, 86, 114, 121, 110, 105, 111, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117, ] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") # test 0 y_pil = F.adjust_sharpness(x_pil, 1) y_np = np.array(y_pil) # torch.testing.assert_close(y_np, x_np) assert np.allclose(y_np, x_np) # test 1 y_pil = F.adjust_sharpness(x_pil, 0.5) y_np = np.array(y_pil) y_ans = [ 75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 30, 30, 74, 103, 96, 114, 97, 110, 100, 101, 114, 32, 81, 103, 108, 102, 101, 107, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117, ] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 2 y_pil = F.adjust_sharpness(x_pil, 2) y_np = np.array(y_pil) y_ans = [ 75, 121, 114, 105, 97, 107, 105, 32, 66, 111, 117, 114, 99, 104, 97, 0, 0, 46, 118, 111, 132, 97, 110, 100, 101, 114, 32, 95, 135, 146, 126, 112, 119, 116, 105, 115, 0, 0, 73, 32, 108, 111, 118, 101, 32, 121, 111, 117, ] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 3 x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") x_th = torch.tensor(x_np.transpose(2, 0, 1)) y_pil = F.adjust_sharpness(x_pil, 2) y_np = np.array(y_pil).transpose(2, 0, 1) y_th = F.adjust_sharpness(x_th, 2) # torch.testing.assert_close(y_np, y_th.numpy()) assert np.allclose(y_np, y_th.numpy()) def test_adjust_gamma(): x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") # test 0 y_pil = F.adjust_gamma(x_pil, 1) y_np = np.array(y_pil) # torch.testing.assert_close(y_np, x_np) assert np.allclose(y_np, x_np) # test 1 y_pil = F.adjust_gamma(x_pil, 0.5) y_np = np.array(y_pil) y_ans = [0, 35, 57, 117, 186, 241, 97, 45, 245, 152, 255, 16] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) # test 2 y_pil = F.adjust_gamma(x_pil, 2) y_np = np.array(y_pil) y_ans = [0, 0, 0, 11, 71, 201, 5, 0, 215, 31, 255, 0] y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape) # torch.testing.assert_close(y_np, y_ans) assert np.allclose(y_np, y_ans) def test_adjusts_L_mode(): x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_rgb = Image.fromarray(x_np, mode="RGB") x_l = x_rgb.convert("L") assert F.adjust_brightness(x_l, 2).mode == "L" assert F.adjust_saturation(x_l, 2).mode == "L" assert F.adjust_contrast(x_l, 2).mode == "L" assert F.adjust_hue(x_l, 0.4).mode == "L" assert F.adjust_sharpness(x_l, 2).mode == "L" assert F.adjust_gamma(x_l, 0.5).mode == "L" def test_rotate(): x = np.zeros((100, 100, 3), dtype=np.uint8) x[40, 40] = [255, 255, 255] with pytest.raises(TypeError, match=r"img should be PIL Image"): F.rotate(x, 10) img = F.to_pil_image(x) result = F.rotate(img, 45) assert result.size == (100, 100) r, c, ch = np.where(result) assert all(x in r for x in [49, 50]) assert all(x in c for x in [36]) assert all(x in ch for x in [0, 1, 2]) result = F.rotate(img, 45, expand=True) assert result.size == (142, 142) r, c, ch = np.where(result) assert all(x in r for x in [70, 71]) assert all(x in c for x in [57]) assert all(x in ch for x in [0, 1, 2]) result = F.rotate(img, 45, center=(40, 40)) assert result.size == (100, 100) r, c, ch = np.where(result) assert all(x in r for x in [40]) assert all(x in c for x in [40]) assert all(x in ch for x in [0, 1, 2]) result_a = F.rotate(img, 90) result_b = F.rotate(img, -270) # assert_equal(np.array(result_a), np.array(result_b)) assert np.allclose(np.array(result_a), np.array(result_b)) @pytest.mark.parametrize("mode", ["L", "RGB", "F"]) def test_rotate_fill(mode): img = F.to_pil_image(np.ones((100, 100, 3), dtype=np.uint8) * 255, "RGB") num_bands = len(mode) wrong_num_bands = num_bands + 1 fill = 127 img_conv = img.convert(mode) img_rot = F.rotate(img_conv, 45.0, fill=fill) pixel = img_rot.getpixel((0, 0)) if not isinstance(pixel, tuple): pixel = (pixel,) assert pixel == tuple([fill] * num_bands) with pytest.raises(ValueError): F.rotate(img_conv, 45.0, fill=tuple([fill] * wrong_num_bands)) def test_gaussian_blur_asserts(): np_img = np.ones((100, 100, 3), dtype=np.uint8) * 255 img = F.to_pil_image(np_img, "RGB") with pytest.raises(ValueError, match=r"If kernel_size is a sequence its length should be 2"): F.gaussian_blur(img, [3]) with pytest.raises(ValueError, match=r"If kernel_size is a sequence its length should be 2"): F.gaussian_blur(img, [3, 3, 3]) with pytest.raises(ValueError, match=r"Kernel size should be a tuple/list of two integers"): transforms.GaussianBlur([3, 3, 3]) with pytest.raises(ValueError, match=r"kernel_size should have odd and positive integers"): F.gaussian_blur(img, [4, 4]) with pytest.raises(ValueError, match=r"Kernel size value should be an odd and positive number"): transforms.GaussianBlur([4, 4]) with pytest.raises(ValueError, match=r"kernel_size should have odd and positive integers"): F.gaussian_blur(img, [-3, -3]) with pytest.raises(ValueError, match=r"Kernel size value should be an odd and positive number"): transforms.GaussianBlur([-3, -3]) with pytest.raises(ValueError, match=r"If sigma is a sequence, its length should be 2"): F.gaussian_blur(img, 3, [1, 1, 1]) with pytest.raises(ValueError, match=r"sigma should be a single number or a list/tuple with length 2"): transforms.GaussianBlur(3, [1, 1, 1]) with pytest.raises(ValueError, match=r"sigma should have positive values"): F.gaussian_blur(img, 3, -1.0) with pytest.raises(ValueError, match=r"If sigma is a single number, it must be positive"): transforms.GaussianBlur(3, -1.0) with pytest.raises(TypeError, match=r"kernel_size should be int or a sequence of integers"): F.gaussian_blur(img, "kernel_size_string") with pytest.raises(ValueError, match=r"Kernel size should be a tuple/list of two integers"): transforms.GaussianBlur("kernel_size_string") with pytest.raises(TypeError, match=r"sigma should be either float or sequence of floats"): F.gaussian_blur(img, 3, "sigma_string") with pytest.raises(ValueError, match=r"sigma should be a single number or a list/tuple with length 2"): transforms.GaussianBlur(3, "sigma_string") def test_lambda(): trans = transforms.Lambda(lambda x: x.add(10)) x = torch.randn(10) y = trans(x) # assert_equal(y, torch.add(x, 10)) assert np.allclose(y.numpy(), torch.add(x, 10).numpy()) trans = transforms.Lambda(lambda x: x.add_(10)) x = torch.randn(10) y = trans(x) # assert_equal(y, x) assert np.allclose(y.numpy(), x.numpy()) # Checking if Lambda can be printed as string trans.__repr__() def test_to_grayscale(): """Unit tests for grayscale transform""" x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") x_pil_2 = x_pil.convert("L") gray_np = np.array(x_pil_2) # Test Set: Grayscale an image with desired number of output channels # Case 1: RGB -> 1 channel grayscale trans1 = transforms.Grayscale(num_output_channels=1) gray_pil_1 = trans1(x_pil) gray_np_1 = np.array(gray_pil_1) assert gray_pil_1.mode == "L", "mode should be L" assert gray_np_1.shape == tuple(x_shape[0:2]), "should be 1 channel" # assert_equal(gray_np, gray_np_1) assert np.allclose(gray_np, gray_np_1) # Case 2: RGB -> 3 channel grayscale trans2 = transforms.Grayscale(num_output_channels=3) gray_pil_2 = trans2(x_pil) gray_np_2 = np.array(gray_pil_2) assert gray_pil_2.mode == "RGB", "mode should be RGB" assert gray_np_2.shape == tuple(x_shape), "should be 3 channel" # assert_equal(gray_np_2[:, :, 0], gray_np_2[:, :, 1]) # assert_equal(gray_np_2[:, :, 1], gray_np_2[:, :, 2]) # assert_equal(gray_np, gray_np_2[:, :, 0]) assert np.allclose(gray_np_2[:, :, 0], gray_np_2[:, :, 1]) assert np.allclose(gray_np_2[:, :, 1], gray_np_2[:, :, 2]) assert np.allclose(gray_np, gray_np_2[:, :, 0]) # Case 3: 1 channel grayscale -> 1 channel grayscale trans3 = transforms.Grayscale(num_output_channels=1) gray_pil_3 = trans3(x_pil_2) gray_np_3 = np.array(gray_pil_3) assert gray_pil_3.mode == "L", "mode should be L" assert gray_np_3.shape == tuple(x_shape[0:2]), "should be 1 channel" # assert_equal(gray_np, gray_np_3) assert np.allclose(gray_np, gray_np_3) # Case 4: 1 channel grayscale -> 3 channel grayscale trans4 = transforms.Grayscale(num_output_channels=3) gray_pil_4 = trans4(x_pil_2) gray_np_4 = np.array(gray_pil_4) assert gray_pil_4.mode == "RGB", "mode should be RGB" assert gray_np_4.shape == tuple(x_shape), "should be 3 channel" # assert_equal(gray_np_4[:, :, 0], gray_np_4[:, :, 1]) # assert_equal(gray_np_4[:, :, 1], gray_np_4[:, :, 2]) # assert_equal(gray_np, gray_np_4[:, :, 0]) assert np.allclose(gray_np_4[:, :, 0], gray_np_4[:, :, 1]) assert np.allclose(gray_np_4[:, :, 1], gray_np_4[:, :, 2]) assert np.allclose(gray_np, gray_np_4[:, :, 0]) # Checking if Grayscale can be printed as string trans4.__repr__() @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("p", (0, 1)) def test_random_apply(p, seed): torch.manual_seed(seed) random_apply_transform = transforms.RandomApply([transforms.RandomRotation((45, 50))], p=p) img = transforms.ToPILImage()(torch.rand(3, 30, 40)) out = random_apply_transform(img) if p == 0: assert out == img elif p == 1: assert out != img # Checking if RandomApply can be printed as string random_apply_transform.__repr__() @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("proba_passthrough", (0, 1)) def test_random_choice(proba_passthrough, seed): random.seed(seed) # RandomChoice relies on python builtin random.choice, not pytorch random_choice_transform = transforms.RandomChoice( [ lambda x: x, # passthrough transforms.RandomRotation((45, 50)), ], p=[proba_passthrough, 1 - proba_passthrough], ) img = transforms.ToPILImage()(torch.rand(3, 30, 40)) out = random_choice_transform(img) if proba_passthrough == 1: assert out == img elif proba_passthrough == 0: assert out != img # Checking if RandomChoice can be printed as string random_choice_transform.__repr__() @pytest.mark.skipif(stats is None, reason="scipy.stats not available") def test_random_order(): random_state = random.getstate() random.seed(42) random_order_transform = transforms.RandomOrder([transforms.Resize(20), transforms.CenterCrop(10)]) img = transforms.ToPILImage()(torch.rand(3, 25, 25)) num_samples = 250 num_normal_order = 0 resize_crop_out = transforms.CenterCrop(10)(transforms.Resize(20)(img)) for _ in range(num_samples): out = random_order_transform(img) if out == resize_crop_out: num_normal_order += 1 p_value = stats.binom_test(num_normal_order, num_samples, p=0.5) random.setstate(random_state) assert p_value > 0.0001 # Checking if RandomOrder can be printed as string random_order_transform.__repr__() def test_linear_transformation(): num_samples = 1000 x = torch.randn(num_samples, 3, 10, 10) flat_x = x.view(x.size(0), x.size(1) * x.size(2) * x.size(3)) # compute principal components sigma = torch.mm(flat_x.t(), flat_x) / flat_x.size(0) u, s, _ = np.linalg.svd(sigma.numpy()) zca_epsilon = 1e-10 # avoid division by 0 d = torch.Tensor(np.diag(1.0 / np.sqrt(s + zca_epsilon))) u = torch.Tensor(u) principal_components = torch.mm(torch.mm(u, d), u.t()) mean_vector = torch.sum(flat_x, dim=0) / flat_x.size(0) # initialize whitening matrix whitening = transforms.LinearTransformation(principal_components, mean_vector) # estimate covariance and mean using weak law of large number num_features = flat_x.size(1) cov = 0.0 mean = 0.0 for i in x: xwhite = whitening(i) xwhite = xwhite.view(1, -1).numpy() cov += np.dot(xwhite, xwhite.T) / num_features mean += np.sum(xwhite) / num_features # if rtol for std = 1e-3 then rtol for cov = 2e-3 as std**2 = cov # torch.testing.assert_close( # cov / num_samples, np.identity(1), rtol=2e-3, atol=1e-8, check_dtype=False, msg="cov not close to 1" # ) # torch.testing.assert_close( # mean / num_samples, 0, rtol=1e-3, atol=1e-8, check_dtype=False, msg="mean not close to 0" # ) assert np.allclose((cov / num_samples), np.identity(1), rtol=2e-3, atol=1e-8) assert np.allclose((mean / num_samples), 0, rtol=1e-3, atol=1e-8) # Checking if LinearTransformation can be printed as string whitening.__repr__() @pytest.mark.parametrize("dtype", int_dtypes()) def test_max_value(dtype): assert F_t._max_value(dtype) == torch.iinfo(dtype).max # remove float testing as it can lead to errors such as # runtime error: 5.7896e+76 is outside the range of representable values of type 'float' # for dtype in float_dtypes(): # self.assertGreater(F_t._max_value(dtype), torch.finfo(dtype).max) @pytest.mark.xfail( reason="torch.iinfo() is not supported by torchscript. See https://github.com/pytorch/pytorch/issues/41492." ) def test_max_value_iinfo(): # @torch.jit.script def max_value(image: torch.Tensor) -> int: return 1 if image.is_floating_point() else torch.iinfo(image.dtype).max @pytest.mark.parametrize("should_vflip", [True, False]) @pytest.mark.parametrize("single_dim", [True, False]) def test_ten_crop(should_vflip, single_dim): to_pil_image = transforms.ToPILImage() h = random.randint(5, 25) w = random.randint(5, 25) crop_h = random.randint(1, h) crop_w = random.randint(1, w) if single_dim: crop_h = min(crop_h, crop_w) crop_w = crop_h transform = transforms.TenCrop(crop_h, vertical_flip=should_vflip) five_crop = transforms.FiveCrop(crop_h) else: transform = transforms.TenCrop((crop_h, crop_w), vertical_flip=should_vflip) five_crop = transforms.FiveCrop((crop_h, crop_w)) img = to_pil_image(torch.FloatTensor(3, h, w).uniform_adpater()) results = transform(img) expected_output = five_crop(img) # Checking if FiveCrop and TenCrop can be printed as string transform.__repr__() five_crop.__repr__() if should_vflip: vflipped_img = img.transpose(_pil_constants.FLIP_TOP_BOTTOM) expected_output += five_crop(vflipped_img) else: hflipped_img = img.transpose(_pil_constants.FLIP_LEFT_RIGHT) expected_output += five_crop(hflipped_img) assert len(results) == 10 assert results == expected_output @pytest.mark.parametrize("single_dim", [True, False]) def test_five_crop(single_dim): to_pil_image = transforms.ToPILImage() h = random.randint(5, 25) w = random.randint(5, 25) crop_h = random.randint(1, h) crop_w = random.randint(1, w) if single_dim: crop_h = min(crop_h, crop_w) crop_w = crop_h transform = transforms.FiveCrop(crop_h) else: transform = transforms.FiveCrop((crop_h, crop_w)) img = torch.FloatTensor(3, h, w).uniform_adpater() results = transform(to_pil_image(img)) assert len(results) == 5 for crop in results: assert crop.size == (crop_w, crop_h) to_pil_image = transforms.ToPILImage() tl = to_pil_image(img[:, 0:crop_h, 0:crop_w]) tr = to_pil_image(img[:, 0:crop_h, w - crop_w :]) bl = to_pil_image(img[:, h - crop_h :, 0:crop_w]) br = to_pil_image(img[:, h - crop_h :, w - crop_w :]) center = transforms.CenterCrop((crop_h, crop_w))(to_pil_image(img)) expected_output = (tl, tr, bl, br, center) assert results == expected_output @pytest.mark.parametrize("policy", transforms.AutoAugmentPolicy) @pytest.mark.parametrize("fill", [None, 85, (128, 128, 128)]) @pytest.mark.parametrize("grayscale", [True, False]) def test_autoaugment(policy, fill, grayscale): random.seed(42) img = Image.open(GRACE_HOPPER) if grayscale: img, fill = _get_grayscale_test_image(img, fill) transform = transforms.AutoAugment(policy=policy, fill=fill) for _ in range(100): img = transform(img) transform.__repr__() @pytest.mark.parametrize("num_ops", [1, 2, 3]) @pytest.mark.parametrize("magnitude", [7, 9, 11]) @pytest.mark.parametrize("fill", [None, 85, (128, 128, 128)]) @pytest.mark.parametrize("grayscale", [True, False]) def test_randaugment(num_ops, magnitude, fill, grayscale): random.seed(42) img = Image.open(GRACE_HOPPER) if grayscale: img, fill = _get_grayscale_test_image(img, fill) transform = transforms.RandAugment(num_ops=num_ops, magnitude=magnitude, fill=fill) for _ in range(100): img = transform(img) transform.__repr__() @pytest.mark.parametrize("fill", [None, 85, (128, 128, 128)]) @pytest.mark.parametrize("num_magnitude_bins", [10, 13, 30]) @pytest.mark.parametrize("grayscale", [True, False]) def test_trivialaugmentwide(fill, num_magnitude_bins, grayscale): random.seed(42) img = Image.open(GRACE_HOPPER) if grayscale: img, fill = _get_grayscale_test_image(img, fill) transform = transforms.TrivialAugmentWide(fill=fill, num_magnitude_bins=num_magnitude_bins) for _ in range(100): img = transform(img) transform.__repr__() # TODO: # @pytest.mark.parametrize("fill", [None, 85, (128, 128, 128)]) # @pytest.mark.parametrize("severity", [1, 10]) # @pytest.mark.parametrize("mixture_width", [1, 2]) # @pytest.mark.parametrize("chain_depth", [-1, 2]) # @pytest.mark.parametrize("all_ops", [True, False]) # @pytest.mark.parametrize("grayscale", [True, False]) # def test_augmix(fill, severity, mixture_width, chain_depth, all_ops, grayscale): # random.seed(42) # img = Image.open(GRACE_HOPPER) # if grayscale: # img, fill = _get_grayscale_test_image(img, fill) # transform = transforms.AugMix( # fill=fill, severity=severity, mixture_width=mixture_width, chain_depth=chain_depth, all_ops=all_ops # ) # for _ in range(100): # img = transform(img) # transform.__repr__() def test_random_crop(): height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 oheight = random.randint(5, (height - 2) / 2) * 2 owidth = random.randint(5, (width - 2) / 2) * 2 img = torch.ones(3, height, width, dtype=torch.uint8) result = transforms.Compose( [ transforms.ToPILImage(), transforms.RandomCrop((oheight, owidth)), transforms.PILToTensor(), ] )(img) assert result.size(1) == oheight assert result.size(2) == owidth padding = random.randint(1, 20) result = transforms.Compose( [ transforms.ToPILImage(), transforms.RandomCrop((oheight, owidth), padding=padding), transforms.PILToTensor(), ] )(img) assert result.size(1) == oheight assert result.size(2) == owidth result = transforms.Compose( [transforms.ToPILImage(), transforms.RandomCrop((height, width)), transforms.PILToTensor()] )(img) assert result.size(1) == height assert result.size(2) == width # torch.testing.assert_close(result, img) assert np.allclose(result.numpy(), img.numpy()) result = transforms.Compose( [ transforms.ToPILImage(), transforms.RandomCrop((height + 1, width + 1), pad_if_needed=True), transforms.PILToTensor(), ] )(img) assert result.size(1) == height + 1 assert result.size(2) == width + 1 t = transforms.RandomCrop(48) img = torch.ones(3, 32, 32) with pytest.raises(ValueError, match=r"Required crop size .+ is larger then input image size .+"): t(img) def test_center_crop(): height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 oheight = random.randint(5, (height - 2) / 2) * 2 owidth = random.randint(5, (width - 2) / 2) * 2 img = torch.ones(3, height, width, dtype=torch.uint8) oh1 = (height - oheight) // 2 ow1 = (width - owidth) // 2 # TODO: # imgnarrow = img[:, oh1 : oh1 + oheight, ow1 : ow1 + owidth] # imgnarrow.fill_(0) img = img.to(torch.float32) img[:, oh1: oh1 + oheight, ow1: ow1 + owidth] = 0. img = img.to(torch.uint8) result = transforms.Compose( [ transforms.ToPILImage(), transforms.CenterCrop((oheight, owidth)), transforms.PILToTensor(), ] )(img) assert result.sum() == 0 oheight += 1 owidth += 1 result = transforms.Compose( [ transforms.ToPILImage(), transforms.CenterCrop((oheight, owidth)), transforms.PILToTensor(), ] )(img) sum1 = result.sum() assert sum1 > 1 oheight += 1 owidth += 1 result = transforms.Compose( [ transforms.ToPILImage(), transforms.CenterCrop((oheight, owidth)), transforms.PILToTensor(), ] )(img) sum2 = result.sum() assert sum2 > 0 assert sum2 > sum1 @pytest.mark.parametrize("odd_image_size", (True, False)) @pytest.mark.parametrize("delta", (1, 3, 5)) @pytest.mark.parametrize("delta_width", (-2, -1, 0, 1, 2)) @pytest.mark.parametrize("delta_height", (-2, -1, 0, 1, 2)) def test_center_crop_2(odd_image_size, delta, delta_width, delta_height): """Tests when center crop size is larger than image size, along any dimension""" # Since height is independent of width, we can ignore images with odd height and even width and vice-versa. input_image_size = (random.randint(10, 32) * 2, random.randint(10, 32) * 2) if odd_image_size: input_image_size = (input_image_size[0] + 1, input_image_size[1] + 1) delta_height *= delta delta_width *= delta img = torch.ones(3, *input_image_size, dtype=torch.uint8) crop_size = (input_image_size[0] + delta_height, input_image_size[1] + delta_width) # Test both transforms, one with PIL input and one with tensor output_pil = transforms.Compose( [transforms.ToPILImage(), transforms.CenterCrop(crop_size), transforms.PILToTensor()], )(img) assert output_pil.size()[1:3] == crop_size output_tensor = transforms.CenterCrop(crop_size)(img) assert output_tensor.size()[1:3] == crop_size # Ensure output for PIL and Tensor are equal # assert_equal( # output_tensor, # output_pil, # msg=f"image_size: {input_image_size} crop_size: {crop_size}", # ) assert np.allclose(output_tensor.numpy(), output_pil.numpy()) # Check if content in center of both image and cropped output is same. center_size = (min(crop_size[0], input_image_size[0]), min(crop_size[1], input_image_size[1])) crop_center_tl, input_center_tl = [0, 0], [0, 0] for index in range(2): if crop_size[index] > input_image_size[index]: crop_center_tl[index] = (crop_size[index] - input_image_size[index]) // 2 else: input_center_tl[index] = (input_image_size[index] - crop_size[index]) // 2 output_center = output_pil[ :, crop_center_tl[0] : crop_center_tl[0] + center_size[0], crop_center_tl[1] : crop_center_tl[1] + center_size[1], ] img_center = img[ :, input_center_tl[0] : input_center_tl[0] + center_size[0], input_center_tl[1] : input_center_tl[1] + center_size[1], ] # assert_equal(output_center, img_center) assert np.allclose(output_center.numpy(), img_center.numpy()) def test_color_jitter(): color_jitter = transforms.ColorJitter(2, 2, 2, 0.1) x_shape = [2, 2, 3] x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape) x_pil = Image.fromarray(x_np, mode="RGB") x_pil_2 = x_pil.convert("L") for _ in range(10): y_pil = color_jitter(x_pil) assert y_pil.mode == x_pil.mode y_pil_2 = color_jitter(x_pil_2) assert y_pil_2.mode == x_pil_2.mode # Checking if ColorJitter can be printed as string color_jitter.__repr__() @pytest.mark.parametrize("seed", range(10)) @pytest.mark.skipif(stats is None, reason="scipy.stats not available") def test_random_erasing(seed): torch.manual_seed(seed) img = torch.ones(3, 128, 128) t = transforms.RandomErasing(scale=(0.1, 0.1), ratio=(1 / 3, 3.0)) y, x, h, w, v = t.get_params( img, t.scale, t.ratio, [ t.value, ], ) aspect_ratio = h / w # Add some tolerance due to the rounding and int conversion used in the transform tol = 0.05 assert 1 / 3 - tol <= aspect_ratio <= 3 + tol # Make sure that h > w and h < w are equaly likely (log-scale sampling) aspect_ratios = [] random.seed(42) trial = 1000 for _ in range(trial): y, x, h, w, v = t.get_params( img, t.scale, t.ratio, [ t.value, ], ) aspect_ratios.append(h / w) count_bigger_then_ones = len([1 for aspect_ratio in aspect_ratios if aspect_ratio > 1]) p_value = stats.binom_test(count_bigger_then_ones, trial, p=0.5) assert p_value > 0.0001 # Checking if RandomErasing can be printed as string t.__repr__() def test_random_rotation(): with pytest.raises(ValueError): transforms.RandomRotation(-0.7) with pytest.raises(ValueError): transforms.RandomRotation([-0.7]) with pytest.raises(ValueError): transforms.RandomRotation([-0.7, 0, 0.7]) t = transforms.RandomRotation(0, fill=None) assert t.fill == 0 t = transforms.RandomRotation(10) angle = t.get_params(t.degrees) assert angle > -10 and angle < 10 t = transforms.RandomRotation((-10, 10)) angle = t.get_params(t.degrees) assert -10 < angle < 10 # Checking if RandomRotation can be printed as string t.__repr__() # assert deprecation warning and non-BC with pytest.warns( UserWarning, match=re.escape( "The parameter 'resample' is deprecated since 0.12 and will be removed 0.14. " "Please use 'interpolation' instead." ), ): t = transforms.RandomRotation((-10, 10), resample=2) assert t.interpolation == transforms.InterpolationMode.BILINEAR # assert changed type warning with pytest.warns( UserWarning, match=re.escape( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ), ): t = transforms.RandomRotation((-10, 10), interpolation=2) assert t.interpolation == transforms.InterpolationMode.BILINEAR def test_random_rotation_error(): # assert fill being either a Sequence or a Number with pytest.raises(TypeError): transforms.RandomRotation(0, fill={}) def test_randomperspective(): for _ in range(10): height = random.randint(24, 32) * 2 width = random.randint(24, 32) * 2 img = torch.ones(3, height, width) to_pil_image = transforms.ToPILImage() img = to_pil_image(img) perp = transforms.RandomPerspective() startpoints, endpoints = perp.get_params(width, height, 0.5) tr_img = F.perspective(img, startpoints, endpoints) tr_img2 = F.convert_image_dtype(F.pil_to_tensor(F.perspective(tr_img, endpoints, startpoints))) tr_img = F.convert_image_dtype(F.pil_to_tensor(tr_img)) assert img.size[0] == width assert img.size[1] == height assert torch.nn.functional.mse_loss( tr_img, F.convert_image_dtype(F.pil_to_tensor(img)) ) + 0.3 > torch.nn.functional.mse_loss(tr_img2, F.convert_image_dtype(F.pil_to_tensor(img))) @pytest.mark.parametrize("seed", range(10)) @pytest.mark.parametrize("mode", ["L", "RGB", "F"]) def test_randomperspective_fill(mode, seed): torch.manual_seed(seed) # assert fill being either a Sequence or a Number with pytest.raises(TypeError): transforms.RandomPerspective(fill={}) t = transforms.RandomPerspective(fill=None) assert t.fill == 0 height = 100 width = 100 img = torch.ones(3, height, width) to_pil_image = transforms.ToPILImage() img = to_pil_image(img) fill = 127 num_bands = len(mode) img_conv = img.convert(mode) perspective = transforms.RandomPerspective(p=1, fill=fill) tr_img = perspective(img_conv) pixel = tr_img.getpixel((0, 0)) if not isinstance(pixel, tuple): pixel = (pixel,) assert pixel == tuple([fill] * num_bands) startpoints, endpoints = transforms.RandomPerspective.get_params(width, height, 0.5) tr_img = F.perspective(img_conv, startpoints, endpoints, fill=fill) pixel = tr_img.getpixel((0, 0)) if not isinstance(pixel, tuple): pixel = (pixel,) assert pixel == tuple([fill] * num_bands) wrong_num_bands = num_bands + 1 with pytest.raises(ValueError): F.perspective(img_conv, startpoints, endpoints, fill=tuple([fill] * wrong_num_bands)) @pytest.mark.skipif(stats is None, reason="scipy.stats not available") def test_normalize(): def samples_from_standard_normal(tensor): tensor_list = list(tensor.view(-1)) tensor_list = [i.numpy() for i in tensor_list] p_value = stats.kstest(tensor_list, "norm", args=(0, 1)).pvalue return p_value > 0.0001 random_state = random.getstate() random.seed(42) for channels in [1, 3]: img = torch.rand(channels, 10, 10) mean = [img[c].mean() for c in range(channels)] std = [img[c].std() for c in range(channels)] normalized = transforms.Normalize(mean, std)(img) assert samples_from_standard_normal(normalized) random.setstate(random_state) # Checking if Normalize can be printed as string transforms.Normalize(mean, std).__repr__() # Checking the optional in-place behaviour tensor = torch.rand((1, 16, 16)) tensor_inplace = transforms.Normalize((0.5,), (0.5,), inplace=True)(tensor) # assert_equal(tensor, tensor_inplace) assert np.allclose(tensor.numpy(), tensor_inplace.numpy()) @pytest.mark.parametrize("dtype1", [torch.float32, torch.float64]) @pytest.mark.parametrize("dtype2", [torch.int64, torch.float32, torch.float64]) def test_normalize_different_dtype(dtype1, dtype2): img = torch.rand(3, 10, 10, dtype=dtype1) mean = torch.tensor([1, 2, 3], dtype=dtype2) std = torch.tensor([1, 2, 1], dtype=dtype2) # checks that it doesn't crash transforms.functional.normalize(img, mean, std) def test_normalize_3d_tensor(): torch.manual_seed(28) n_channels = 3 img_size = 10 mean = torch.rand(n_channels) std = torch.rand(n_channels) img = torch.rand(n_channels, img_size, img_size) target = F.normalize(img, mean, std) mean_unsqueezed = mean.view(-1, 1, 1) std_unsqueezed = std.view(-1, 1, 1) result1 = F.normalize(img, mean_unsqueezed, std_unsqueezed) result2 = F.normalize( img, mean_unsqueezed.repeat(1, img_size, img_size), std_unsqueezed.repeat(1, img_size, img_size) ) # torch.testing.assert_close(target, result1) # torch.testing.assert_close(target, result2) assert np.allclose(target.numpy(), result1.numpy()) assert np.allclose(target.numpy(), result2.numpy()) class TestAffine: @pytest.fixture(scope="class") def input_img(self): input_img = np.zeros((40, 40, 3), dtype=np.uint8) for pt in [(16, 16), (20, 16), (20, 20)]: for i in range(-5, 5): for j in range(-5, 5): input_img[pt[0] + i, pt[1] + j, :] = [255, 155, 55] return input_img def test_affine_translate_seq(self, input_img): with pytest.raises(TypeError, match=r"Argument translate should be a sequence"): F.affine(input_img, 10, translate=0, scale=1, shear=1) @pytest.fixture(scope="class") def pil_image(self, input_img): return F.to_pil_image(input_img) def _to_3x3_inv(self, inv_result_matrix): result_matrix = np.zeros((3, 3)) result_matrix[:2, :] = np.array(inv_result_matrix).reshape((2, 3)) result_matrix[2, 2] = 1 return np.linalg.inv(result_matrix) def _test_transformation(self, angle, translate, scale, shear, pil_image, input_img, center=None): a_rad = math.radians(angle) s_rad = [math.radians(sh_) for sh_ in shear] cnt = [20, 20] if center is None else center cx, cy = cnt tx, ty = translate sx, sy = s_rad rot = a_rad # 1) Check transformation matrix: C = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]]) T = np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]]) Cinv = np.linalg.inv(C) RS = np.array( [ [scale * math.cos(rot), -scale * math.sin(rot), 0], [scale * math.sin(rot), scale * math.cos(rot), 0], [0, 0, 1], ] ) SHx = np.array([[1, -math.tan(sx), 0], [0, 1, 0], [0, 0, 1]]) SHy = np.array([[1, 0, 0], [-math.tan(sy), 1, 0], [0, 0, 1]]) RSS = np.matmul(RS, np.matmul(SHy, SHx)) true_matrix = np.matmul(T, np.matmul(C, np.matmul(RSS, Cinv))) result_matrix = self._to_3x3_inv( F._get_inverse_affine_matrix(center=cnt, angle=angle, translate=translate, scale=scale, shear=shear) ) assert np.sum(np.abs(true_matrix - result_matrix)) < 1e-10 # 2) Perform inverse mapping: true_result = np.zeros((40, 40, 3), dtype=np.uint8) inv_true_matrix = np.linalg.inv(true_matrix) for y in range(true_result.shape[0]): for x in range(true_result.shape[1]): # Same as for PIL: # https://github.com/python-pillow/Pillow/blob/71f8ec6a0cfc1008076a023c0756542539d057ab/ # src/libImaging/Geometry.c#L1060 input_pt = np.array([x + 0.5, y + 0.5, 1.0]) res = np.floor(np.dot(inv_true_matrix, input_pt)).astype(np.int) _x, _y = res[:2] if 0 <= _x < input_img.shape[1] and 0 <= _y < input_img.shape[0]: true_result[y, x, :] = input_img[_y, _x, :] result = F.affine(pil_image, angle=angle, translate=translate, scale=scale, shear=shear, center=center) assert result.size == pil_image.size # Compute number of different pixels: np_result = np.array(result) n_diff_pixels = np.sum(np_result != true_result) / 3 # Accept 3 wrong pixels error_msg = ( f"angle={angle}, translate={translate}, scale={scale}, shear={shear}\nn diff pixels={n_diff_pixels}\n" ) assert n_diff_pixels < 3, error_msg def test_transformation_discrete(self, pil_image, input_img): # Test rotation angle = 45 self._test_transformation( angle=angle, translate=(0, 0), scale=1.0, shear=(0.0, 0.0), pil_image=pil_image, input_img=input_img ) # Test rotation angle = 45 self._test_transformation( angle=angle, translate=(0, 0), scale=1.0, shear=(0.0, 0.0), pil_image=pil_image, input_img=input_img, center=[0, 0], ) # Test translation translate = [10, 15] self._test_transformation( angle=0.0, translate=translate, scale=1.0, shear=(0.0, 0.0), pil_image=pil_image, input_img=input_img ) # Test scale scale = 1.2 self._test_transformation( angle=0.0, translate=(0.0, 0.0), scale=scale, shear=(0.0, 0.0), pil_image=pil_image, input_img=input_img ) # Test shear shear = [45.0, 25.0] self._test_transformation( angle=0.0, translate=(0.0, 0.0), scale=1.0, shear=shear, pil_image=pil_image, input_img=input_img ) # Test shear with top-left as center shear = [45.0, 25.0] self._test_transformation( angle=0.0, translate=(0.0, 0.0), scale=1.0, shear=shear, pil_image=pil_image, input_img=input_img, center=[0, 0], ) @pytest.mark.parametrize("angle", range(-90, 90, 36)) @pytest.mark.parametrize("translate", range(-10, 10, 5)) @pytest.mark.parametrize("scale", [0.77, 1.0, 1.27]) @pytest.mark.parametrize("shear", range(-15, 15, 5)) def test_transformation_range(self, angle, translate, scale, shear, pil_image, input_img): self._test_transformation( angle=angle, translate=(translate, translate), scale=scale, shear=(shear, shear), pil_image=pil_image, input_img=input_img, ) def test_random_affine(): with pytest.raises(ValueError): transforms.RandomAffine(-0.7) with pytest.raises(ValueError): transforms.RandomAffine([-0.7]) with pytest.raises(ValueError): transforms.RandomAffine([-0.7, 0, 0.7]) with pytest.raises(TypeError): transforms.RandomAffine([-90, 90], translate=2.0) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[-1.0, 1.0]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[-1.0, 0.0, 1.0]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[0.0]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[-1.0, 1.0]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[0.5, -0.5]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[0.5, 3.0, -0.5]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[0.5, 0.5], shear=-7) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[0.5, 0.5], shear=[-10]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[0.5, 0.5], shear=[-10, 0, 10]) with pytest.raises(ValueError): transforms.RandomAffine([-90, 90], translate=[0.2, 0.2], scale=[0.5, 0.5], shear=[-10, 0, 10, 0, 10]) # assert fill being either a Sequence or a Number with pytest.raises(TypeError): transforms.RandomAffine(0, fill={}) t = transforms.RandomAffine(0, fill=None) assert t.fill == 0 x = np.zeros((100, 100, 3), dtype=np.uint8) img = F.to_pil_image(x) t = transforms.RandomAffine(10, translate=[0.5, 0.3], scale=[0.7, 1.3], shear=[-10, 10, 20, 40]) for _ in range(100): angle, translations, scale, shear = t.get_params(t.degrees, t.translate, t.scale, t.shear, img_size=img.size) assert -10 < angle < 10 assert -img.size[0] * 0.5 <= translations[0] <= img.size[0] * 0.5 assert -img.size[1] * 0.5 <= translations[1] <= img.size[1] * 0.5 assert 0.7 < scale < 1.3 assert -10 < shear[0] < 10 assert -20 < shear[1] < 40 # Checking if RandomAffine can be printed as string t.__repr__() t = transforms.RandomAffine(10, interpolation=transforms.InterpolationMode.BILINEAR) assert "bilinear" in t.__repr__() # assert deprecation warning and non-BC with pytest.warns( UserWarning, match=re.escape( "The parameter 'resample' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'interpolation' instead." ), ): t = transforms.RandomAffine(10, resample=2) assert t.interpolation == transforms.InterpolationMode.BILINEAR with pytest.warns( UserWarning, match=re.escape( "The parameter 'fillcolor' is deprecated since 0.12 and will be removed in 0.14. " "Please use 'fill' instead." ), ): t = transforms.RandomAffine(10, fillcolor=10) assert t.fill == 10 # assert changed type warning with pytest.warns( UserWarning, match=re.escape( "Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. " "Please use InterpolationMode enum." ), ): t = transforms.RandomAffine(10, interpolation=2) assert t.interpolation == transforms.InterpolationMode.BILINEAR if __name__ == "__main__": pytest.main([__file__]) msadapter/testing/ut/torchvision/test_transforms_video.py000066400000000000000000000155161436713554100245300ustar00rootroot00000000000000import random import warnings import numpy as np import pytest # import torch # from torchvision.transforms import Compose import ms_adapter.pytorch as torch from ms_adapter.torchvision.transforms import Compose try: from scipy import stats except ImportError: stats = None with warnings.catch_warnings(record=True): warnings.simplefilter("always") import ms_adapter.torchvision.transforms._transforms_video as transforms class TestVideoTransforms: def test_random_crop_video(self): numFrames = random.randint(4, 128) height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 oheight = random.randint(5, (height - 2) / 2) * 2 owidth = random.randint(5, (width - 2) / 2) * 2 clip = torch.randint(0, 256, (numFrames, height, width, 3), dtype=torch.uint8) result = Compose( [ transforms.ToTensorVideo(), transforms.RandomCropVideo((oheight, owidth)), ] )(clip) assert result.size(2) == oheight assert result.size(3) == owidth transforms.RandomCropVideo((oheight, owidth)).__repr__() def test_random_resized_crop_video(self): numFrames = random.randint(4, 128) height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 oheight = random.randint(5, (height - 2) / 2) * 2 owidth = random.randint(5, (width - 2) / 2) * 2 clip = torch.randint(0, 256, (numFrames, height, width, 3), dtype=torch.uint8) result = Compose( [ transforms.ToTensorVideo(), transforms.RandomResizedCropVideo((oheight, owidth)), ] )(clip) assert result.size(2) == oheight assert result.size(3) == owidth transforms.RandomResizedCropVideo((oheight, owidth)).__repr__() def test_center_crop_video(self): numFrames = random.randint(4, 128) height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 oheight = random.randint(5, (height - 2) / 2) * 2 owidth = random.randint(5, (width - 2) / 2) * 2 clip = torch.ones((numFrames, height, width, 3), dtype=torch.uint8) * 255 oh1 = (height - oheight) // 2 ow1 = (width - owidth) // 2 # clipNarrow = clip[:, oh1 : oh1 + oheight, ow1 : ow1 + owidth, :] # clipNarrow.fill_(0) clip = clip.to(torch.float32) clip[:, oh1 : oh1 + oheight, ow1 : ow1 + owidth, :] = 0. clip = clip.to(torch.uint8) result = Compose( [ transforms.ToTensorVideo(), transforms.CenterCropVideo((oheight, owidth)), ] )(clip) msg = ( "height: " + str(height) + " width: " + str(width) + " oheight: " + str(oheight) + " owidth: " + str(owidth) ) assert result.sum().item() == 0, msg oheight += 1 owidth += 1 result = Compose( [ transforms.ToTensorVideo(), transforms.CenterCropVideo((oheight, owidth)), ] )(clip) sum1 = result.sum() msg = ( "height: " + str(height) + " width: " + str(width) + " oheight: " + str(oheight) + " owidth: " + str(owidth) ) assert sum1.item() > 1, msg oheight += 1 owidth += 1 result = Compose( [ transforms.ToTensorVideo(), transforms.CenterCropVideo((oheight, owidth)), ] )(clip) sum2 = result.sum() msg = ( "height: " + str(height) + " width: " + str(width) + " oheight: " + str(oheight) + " owidth: " + str(owidth) ) assert sum2.item() > 1, msg assert sum2.item() > sum1.item(), msg @pytest.mark.skipif(stats is None, reason="scipy.stats is not available") @pytest.mark.parametrize("channels", [1, 3]) def test_normalize_video(self, channels): def samples_from_standard_normal(tensor): # p_value = stats.kstest(list(tensor.view(-1)), "norm", args=(0, 1)).pvalue tensor_list = list(tensor.view(-1)) tensor_list = [i.numpy() for i in tensor_list] p_value = stats.kstest(tensor_list, "norm", args=(0, 1)).pvalue return p_value > 0.0001 random_state = random.getstate() random.seed(42) numFrames = random.randint(4, 128) height = random.randint(32, 256) width = random.randint(32, 256) mean = random.random() std = random.random() # clip = torch.normal(mean, std, size=(channels, numFrames, height, width)) clip = torch.Tensor(channels, numFrames, height, width).normal_adapter(mean, std) mean = [clip[c].mean().item() for c in range(channels)] std = [clip[c].std().item() for c in range(channels)] normalized = transforms.NormalizeVideo(mean, std)(clip) assert samples_from_standard_normal(normalized) random.setstate(random_state) # # Checking the optional in-place behaviour # tensor = torch.rand((3, 128, 16, 16)) # tensor_inplace = transforms.NormalizeVideo((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)(tensor) # assert_equal(tensor, tensor_inplace) # # assert np.allclose(tensor.numpy(), tensor_inplace.numpy()) # transforms.NormalizeVideo((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True).__repr__() def test_to_tensor_video(self): numFrames, height, width = 64, 4, 4 trans = transforms.ToTensorVideo() with pytest.raises(TypeError): np_rng = np.random.RandomState(0) trans(np_rng.rand(numFrames, height, width, 1).tolist()) with pytest.raises(TypeError): trans(torch.rand((numFrames, height, width, 1), dtype=torch.float)) with pytest.raises(ValueError): trans(torch.ones((3, numFrames, height, width, 3), dtype=torch.uint8)) with pytest.raises(ValueError): trans(torch.ones((height, width, 3), dtype=torch.uint8)) with pytest.raises(ValueError): trans(torch.ones((width, 3), dtype=torch.uint8)) with pytest.raises(ValueError): trans(torch.ones((3), dtype=torch.uint8)) trans.__repr__() @pytest.mark.parametrize("p", (0, 1)) def test_random_horizontal_flip_video(self, p): clip = torch.rand((3, 4, 112, 112), dtype=torch.float) hclip = clip.flip(-1) out = transforms.RandomHorizontalFlipVideo(p=p)(clip) if p == 0: # torch.testing.assert_close(out, clip) assert np.allclose(out.numpy(), clip.numpy()) elif p == 1: # torch.testing.assert_close(out, hclip) assert np.allclose(out.numpy(), hclip.numpy()) transforms.RandomHorizontalFlipVideo().__repr__() if __name__ == "__main__": pytest.main([__file__]) msadapter/third_party/000077500000000000000000000000001436713554100154005ustar00rootroot00000000000000msadapter/third_party/einops/000077500000000000000000000000001436713554100166755ustar00rootroot00000000000000msadapter/third_party/einops/patch/000077500000000000000000000000001436713554100177745ustar00rootroot00000000000000msadapter/third_party/einops/patch/0001-adapt-to-ms_adapter.pytorch.patch000066400000000000000000000050041436713554100267070ustar00rootroot00000000000000From ed543c49ec315e1e8cb9c6f8cb0fbc9a08deaad1 Mon Sep 17 00:00:00 2001 From: lvhaoyu Date: Tue, 22 Nov 2022 15:03:16 +0800 Subject: [PATCH] adapt to ms_adapter.pytorch --- einops/_backends.py | 4 ++-- einops/_torch_specific.py | 2 +- einops/layers/torch.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/einops/_backends.py b/einops/_backends.py index ba6fca8..80cce17 100644 --- a/einops/_backends.py +++ b/einops/_backends.py @@ -316,7 +316,7 @@ class TorchBackend(AbstractBackend): framework_name = 'torch' def __init__(self): - import torch + import ms_adapter.pytorch as torch self.torch = torch def is_appropriate_type(self, tensor): @@ -677,4 +677,4 @@ class OneFlowBackend(AbstractBackend): return oneflow def einsum(self, pattern, *x): - return self.flow.einsum(pattern, *x) \ No newline at end of file + return self.flow.einsum(pattern, *x) diff --git a/einops/_torch_specific.py b/einops/_torch_specific.py index 204d935..5ea9253 100644 --- a/einops/_torch_specific.py +++ b/einops/_torch_specific.py @@ -10,7 +10,7 @@ Importantly, whole lib is designed so that you can't use it from typing import Dict, List -import torch +import ms_adapter.pytorch as torch from einops.einops import TransformRecipe, _reconstruct_from_shape_uncached diff --git a/einops/layers/torch.py b/einops/layers/torch.py index 3199241..cb9af63 100644 --- a/einops/layers/torch.py +++ b/einops/layers/torch.py @@ -1,6 +1,6 @@ from typing import Optional, Dict, cast -import torch +import ms_adapter.pytorch as torch from . import RearrangeMixin, ReduceMixin from ._einmix import _EinmixMixin @@ -29,10 +29,10 @@ class Reduce(ReduceMixin, torch.nn.Module): class EinMix(_EinmixMixin, torch.nn.Module): def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound): - self.weight = torch.nn.Parameter(torch.zeros(weight_shape).uniform_(-weight_bound, weight_bound), + self.weight = torch.nn.Parameter(torch.zeros(weight_shape).uniform(-weight_bound, weight_bound), requires_grad=True) if bias_shape is not None: - self.bias = torch.nn.Parameter(torch.zeros(bias_shape).uniform_(-bias_bound, bias_bound), + self.bias = torch.nn.Parameter(torch.zeros(bias_shape).uniform(-bias_bound, bias_bound), requires_grad=True) else: self.bias = None -- 2.25.1 msadapter/third_party/einops/setup.sh000066400000000000000000000005541436713554100203750ustar00rootroot00000000000000#! bin/bash if [ -d "einops" ];then echo "[einops setup]: 'einops' dir has exist, please check and remove it before installing einops" exit 1 fi git clone https://github.com/arogozhnikov/einops.git || exit 1 cd einops git checkout v0.6.0 -b v0.6.0 || exit 1 patch -p1 < ../patch/0001-adapt-to-ms_adapter.pytorch.patch || exit 1 pip install -e ./ || exit 1