123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990 |
- # Copyright (c) Microsoft Corporation.
- # SPDX-License-Identifier: Apache-2.0
- # DeepSpeed Team
- import pytest
- import torch
- import deepspeed
- from deepspeed.ops.op_builder import SpatialInferenceBuilder
- from deepspeed.ops.transformer.inference.bias_add import nhwc_bias_add
- from deepspeed.accelerator import get_accelerator
- if not deepspeed.ops.__compatible_ops__[SpatialInferenceBuilder.NAME]:
- pytest.skip("Inference ops are not available on this system", allow_module_level=True)
- def allclose(x, y):
- assert x.dtype == y.dtype
- rtol, atol = {torch.float32: (5e-3, 5e-4), torch.float16: (3e-2, 2e-3), torch.int8: (1, 1)}[x.dtype]
- return torch.allclose(x, y, rtol=rtol, atol=atol)
- def ref_bias_add(activations, bias):
- return activations + bias.reshape(1, -1, 1, 1)
- channels_list = [192, 384, 320, 576, 640, 768, 960, 1152, 1280, 1536, 1600, 1920, 2240, 2560]
- @pytest.mark.inference_ops
- @pytest.mark.parametrize("batch", [1, 2, 10])
- @pytest.mark.parametrize("image_size", [16, 32, 64])
- @pytest.mark.parametrize("channels", channels_list)
- def test_bias_add(batch, image_size, channels):
- activations = torch.randn((batch, channels, image_size, image_size),
- dtype=torch.float16,
- device=get_accelerator().device_name()).to(memory_format=torch.channels_last)
- bias = torch.randn((channels), dtype=torch.float16, device=get_accelerator().device_name())
- ref_vals = ref_bias_add(activations.clone().detach(), bias)
- ds_vals = nhwc_bias_add(activations, bias)
- assert allclose(ds_vals, ref_vals)
- def ref_bias_add_add(activations, bias, other):
- return (activations + bias.reshape(1, -1, 1, 1)) + other
- @pytest.mark.inference_ops
- @pytest.mark.parametrize("batch", [1, 2, 10])
- @pytest.mark.parametrize("image_size", [16, 32, 64])
- @pytest.mark.parametrize("channels", channels_list)
- def test_bias_add_add(batch, image_size, channels):
- activations = torch.randn((batch, channels, image_size, image_size),
- dtype=torch.float16,
- device=get_accelerator().device_name()).to(memory_format=torch.channels_last)
- other = torch.randn((batch, channels, image_size, image_size),
- dtype=torch.float16,
- device=get_accelerator().device_name()).to(memory_format=torch.channels_last)
- bias = torch.randn((channels), dtype=torch.float16, device=get_accelerator().device_name())
- ref_vals = ref_bias_add_add(activations.clone().detach(), bias, other)
- ds_vals = nhwc_bias_add(activations, bias, other=other)
- assert allclose(ds_vals, ref_vals)
- def ref_bias_add_bias_add(activations, bias, other, other_bias):
- return (activations + bias.reshape(1, -1, 1, 1)) + (other + other_bias.reshape(1, -1, 1, 1))
- @pytest.mark.inference_ops
- @pytest.mark.parametrize("batch", [1, 2, 10])
- @pytest.mark.parametrize("image_size", [16, 32, 64])
- @pytest.mark.parametrize("channels", channels_list)
- def test_bias_add_bias_add(batch, image_size, channels):
- activations = torch.randn((batch, channels, image_size, image_size),
- dtype=torch.float16,
- device=get_accelerator().device_name()).to(memory_format=torch.channels_last)
- other = torch.randn((batch, channels, image_size, image_size),
- dtype=torch.float16,
- device=get_accelerator().device_name()).to(memory_format=torch.channels_last)
- bias = torch.randn((channels), dtype=torch.float16, device=get_accelerator().device_name())
- other_bias = torch.randn((channels), dtype=torch.float16, device=get_accelerator().device_name())
- ref_vals = ref_bias_add_bias_add(activations.clone().detach(), bias, other, other_bias)
- ds_vals = nhwc_bias_add(activations, bias, other=other, other_bias=other_bias)
- assert allclose(ds_vals, ref_vals)
|