mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
85 lines
3.6 KiB
85 lines
3.6 KiB
from typing import Any, Callable, Union
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from colossalai.fx._compatibility import is_compatible_with_meta
|
|
from colossalai.testing import clear_cache_before_run
|
|
|
|
if is_compatible_with_meta():
|
|
from colossalai.fx.profiler import MetaTensor
|
|
|
|
aten = torch.ops.aten
|
|
|
|
registered_meta = {
|
|
('aten.convolution.default', True): [ # (aten ops, requires_backward)
|
|
(nn.Conv1d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4)),
|
|
(nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4, 4)),
|
|
(nn.Conv3d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4, 4, 4)),
|
|
(nn.ConvTranspose1d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4)),
|
|
(nn.ConvTranspose2d(in_channels=3, out_channels=4, kernel_size=2, padding=1,
|
|
dilation=2), torch.rand(2, 3, 4, 4)),
|
|
(nn.ConvTranspose3d(in_channels=3, out_channels=4, kernel_size=2, padding=1,
|
|
dilation=2), torch.rand(2, 3, 4, 4, 4)),
|
|
],
|
|
('aten.native_batch_norm.default', True): [
|
|
(nn.BatchNorm1d(4), torch.rand(2, 4)),
|
|
(nn.BatchNorm2d(4), torch.rand(1, 4, 4, 4)),
|
|
(nn.BatchNorm3d(4), torch.rand(1, 4, 4, 4, 4)),
|
|
],
|
|
('aten.native_layer_norm.default', True): [(nn.LayerNorm(4), torch.rand(1, 2, 3, 4)),],
|
|
('aten.avg_pool1d.default', True): [
|
|
(nn.MaxPool1d(3, stride=2), torch.rand(4, 5, 5)),
|
|
(nn.AvgPool1d(3, stride=2), torch.rand(4, 5, 5)),
|
|
(nn.AdaptiveMaxPool1d(3), torch.rand(4, 5, 5)),
|
|
(nn.AdaptiveAvgPool1d(3), torch.rand(4, 5, 5)),
|
|
],
|
|
('aten.avg_pool2d.default', True): [
|
|
(nn.MaxPool2d((3, 2), stride=(2, 1)), torch.rand(2, 4, 5, 5)),
|
|
(nn.AvgPool2d((3, 2), stride=(2, 1)), torch.rand(2, 4, 5, 5)),
|
|
(nn.AdaptiveMaxPool2d((3, 2)), torch.rand(2, 4, 5, 5)),
|
|
(nn.AdaptiveAvgPool2d((3, 2)), torch.rand(2, 4, 5, 5)),
|
|
],
|
|
('aten.relu.default', True): [
|
|
(nn.ReLU(), torch.rand(4, 3, 1, 2)),
|
|
(nn.LeakyReLU(), torch.rand(4, 3, 1, 2)),
|
|
(nn.SiLU(), torch.rand(4, 3, 1, 2)),
|
|
(nn.GELU(), torch.rand(4, 3, 1, 2)),
|
|
(nn.ELU(), torch.rand(4, 3, 1, 2)),
|
|
(nn.Sigmoid(), torch.rand(4, 3, 1, 2)),
|
|
(nn.Tanh(), torch.rand(4, 3, 1, 2)),
|
|
(nn.Hardswish(), torch.rand(4, 3, 1, 2)),
|
|
]
|
|
}
|
|
|
|
|
|
def compare_all(tensor: torch.Tensor, meta_tensor: torch.Tensor) -> Any:
|
|
assert tensor.shape == meta_tensor.shape, f'the shape of tensor ({tensor.shape}) and meta tensor ({meta_tensor.shape}) does not match.'
|
|
assert tensor.dtype == meta_tensor.dtype, f'the dtype of tensor ({tensor.dtype}) and meta tensor ({meta_tensor.dtype}) does not match.'
|
|
assert tensor.stride() == meta_tensor.stride(
|
|
), f'the stride of tensor ({tensor.stride()}) and meta tensor ({meta_tensor.stride()}) does not match.'
|
|
|
|
|
|
def run_and_compare(f: Union[nn.Module, Callable], x: torch.Tensor, requires_backward=False) -> Any:
|
|
x.requires_grad = requires_backward
|
|
meta_x = MetaTensor(x)
|
|
x_out, meta_out = f(x), f(meta_x)
|
|
compare_all(x_out, meta_out)
|
|
if requires_backward:
|
|
x_out.sum().backward()
|
|
meta_out.sum().backward()
|
|
compare_all(x.grad, meta_x.grad)
|
|
|
|
|
|
@pytest.mark.skipif(not is_compatible_with_meta(), reason='torch version is lower than 1.12.0')
|
|
@clear_cache_before_run()
|
|
def test_meta_aten():
|
|
for (aten_op, requires_backward), v in registered_meta.items():
|
|
for f, x in v:
|
|
run_and_compare(f, x, requires_backward)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_meta_aten()
|