File size: 3,588 Bytes
eca4864 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | # SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
import torch
from torch import nn
class ComplexReLU(nn.Module):
def __init__(self, negative_slope=0.0, mode="cartesian", bias_shape=None):
super(ComplexReLU, self).__init__()
# store parameters
self.mode = mode
if self.mode in ["modulus", "halfplane"]:
if bias_shape is not None:
self.bias = nn.Parameter(torch.zeros(bias_shape, dtype=torch.float32))
else:
self.bias = nn.Parameter(torch.zeros((1), dtype=torch.float32))
else:
bias = torch.zeros((1), dtype=torch.float32)
self.register_buffer("bias", bias)
self.negative_slope = negative_slope
self.act = nn.LeakyReLU(negative_slope=negative_slope)
def forward(self, z: torch.Tensor) -> torch.Tensor:
if self.mode == "cartesian":
zr = torch.view_as_real(z)
za = self.act(zr)
out = torch.view_as_complex(za)
elif self.mode == "modulus":
zabs = torch.sqrt(torch.square(z.real) + torch.square(z.imag))
out = self.act(zabs + self.bias) * torch.exp(1.0j * z.angle())
elif self.mode == "halfplane":
# bias is an angle parameter in this case
modified_angle = torch.angle(z) - self.bias
condition = torch.logical_and(
(0.0 <= modified_angle), (modified_angle < torch.pi / 2.0)
)
out = torch.where(condition, z, self.negative_slope * z)
elif self.mode == "real":
zr = torch.view_as_real(z)
outr = torch.stack((self.act(zr[..., 0]), zr[..., 1]), dim=-1)
out = torch.view_as_complex(outr)
else:
# identity
out = z
return out
class ComplexActivation(nn.Module):
def __init__(self, activation, mode="cartesian", bias_shape=None):
super(ComplexActivation, self).__init__()
# store parameters
self.mode = mode
if self.mode == "modulus":
if bias_shape is not None:
self.bias = nn.Parameter(torch.zeros(bias_shape, dtype=torch.float32))
else:
self.bias = nn.Parameter(torch.zeros((1), dtype=torch.float32))
else:
bias = torch.zeros((1), dtype=torch.float32)
self.register_buffer("bias", bias)
# real valued activation
self.act = activation
def forward(self, z: torch.Tensor) -> torch.Tensor:
if self.mode == "cartesian":
zr = torch.view_as_real(z)
za = self.act(zr)
out = torch.view_as_complex(za)
elif self.mode == "modulus":
zabs = torch.sqrt(torch.square(z.real) + torch.square(z.imag))
out = self.act(zabs + self.bias) * torch.exp(1.0j * z.angle())
else:
# identity
out = z
return out
|