169 lines
4.5 KiB
Python
169 lines
4.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def complex_mse_loss(input, target):
|
|
"""
|
|
Compute the mean squared error between two complex tensors.
|
|
"""
|
|
if input.is_complex():
|
|
return torch.mean(torch.square(input.real - target.real) + torch.square(input.imag - target.imag))
|
|
else:
|
|
return F.mse_loss(input, target)
|
|
|
|
|
|
def complex_sse_loss(input, target):
|
|
"""
|
|
Compute the sum squared error between two complex tensors.
|
|
"""
|
|
if input.is_complex():
|
|
return torch.sum(torch.square(input.real - target.real) + torch.square(input.imag - target.imag))
|
|
else:
|
|
return torch.sum(torch.square(input - target))
|
|
|
|
|
|
class UnitaryLayer(nn.Module):
|
|
def __init__(self, in_features, out_features, dtype=None):
|
|
assert in_features >= out_features
|
|
super(UnitaryLayer, self).__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.weight = nn.Parameter(torch.randn(in_features, out_features, dtype=dtype))
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self):
|
|
q, _ = torch.linalg.qr(self.weight)
|
|
self.weight.data = q
|
|
|
|
def forward(self, x):
|
|
return torch.matmul(x, self.weight)
|
|
|
|
def __repr__(self):
|
|
return f"UnitaryLayer({self.in_features}, {self.out_features})"
|
|
|
|
|
|
class SemiUnitaryLayer(nn.Module):
|
|
def __init__(self, input_dim, output_dim, dtype=None):
|
|
super(SemiUnitaryLayer, self).__init__()
|
|
self.input_dim = input_dim
|
|
self.output_dim = output_dim
|
|
|
|
# Create a larger square matrix for QR decomposition
|
|
self.weight = nn.Parameter(torch.randn(max(input_dim, output_dim), max(input_dim, output_dim), dtype=dtype))
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self):
|
|
# Ensure the weights are semi-unitary by QR decomposition
|
|
q, _ = torch.linalg.qr(self.weight)
|
|
if self.input_dim > self.output_dim:
|
|
self.weight.data = q[: self.input_dim, : self.output_dim]
|
|
else:
|
|
self.weight.data = q[: self.output_dim, : self.input_dim].t()
|
|
|
|
def forward(self, x):
|
|
out = torch.matmul(x, self.weight)
|
|
return out
|
|
|
|
def __repr__(self):
|
|
return f"SemiUnitaryLayer({self.input_dim}, {self.output_dim})"
|
|
|
|
|
|
# class SpreadLayer(nn.Module):
|
|
# def __init__(self, in_features, out_features, dtype=None):
|
|
# super(SpreadLayer, self).__init__()
|
|
# self.in_features = in_features
|
|
# self.out_features = out_features
|
|
# self.mat = torch.ones(in_features, out_features, dtype=dtype)*torch.sqrt(torch.tensor(in_features/out_features))
|
|
|
|
# def forward(self, x):
|
|
# # N in_features -> M out_features, Enery is preserved (P = abs(x)^2)
|
|
# out = torch.matmul(x, self.mat)
|
|
# return out
|
|
|
|
|
|
#### as defined by zhang et al
|
|
|
|
|
|
class Identity(nn.Module):
|
|
"""
|
|
implements the "activation" function
|
|
M(z) = z
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(Identity, self).__init__()
|
|
|
|
def forward(self, x):
|
|
return x
|
|
|
|
|
|
class Mag(nn.Module):
|
|
"""
|
|
implements the activation function
|
|
M(z) = ||z||
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(Mag, self).__init__()
|
|
|
|
def forward(self, x):
|
|
return torch.abs(x).to(dtype=x.dtype)
|
|
|
|
|
|
class ModReLU(nn.Module):
|
|
"""
|
|
implements the activation function
|
|
M(z) = ReLU(||z|| + b)*exp(j*theta_z)
|
|
= ReLU(||z|| + b)*z/||z||
|
|
"""
|
|
|
|
def __init__(self, b=0):
|
|
super(ModReLU, self).__init__()
|
|
self.b = torch.tensor(b)
|
|
|
|
def forward(self, x):
|
|
if x.is_complex():
|
|
mod = torch.abs(x.real**2 + x.imag**2)
|
|
return torch.relu(mod + self.b) * x / mod
|
|
|
|
else:
|
|
return torch.relu(x + self.b)
|
|
|
|
def __repr__(self):
|
|
return f"ModReLU(b={self.b})"
|
|
|
|
|
|
class CReLU(nn.Module):
|
|
"""
|
|
implements the activation function
|
|
M(z) = ReLU(Re(z)) + j*ReLU(Im(z))
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(CReLU, self).__init__()
|
|
|
|
def forward(self, x):
|
|
if x.is_complex():
|
|
return torch.relu(x.real) + 1j * torch.relu(x.imag)
|
|
else:
|
|
return torch.relu(x)
|
|
|
|
|
|
class ZReLU(nn.Module):
|
|
"""
|
|
implements the activation function
|
|
|
|
M(z) = z if 0 <= angle(z) <= pi/2
|
|
= 0 otherwise
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(ZReLU, self).__init__()
|
|
|
|
def forward(self, x):
|
|
if x.is_complex():
|
|
return x * (torch.angle(x) >= 0) * (torch.angle(x) <= torch.pi / 2)
|
|
else:
|
|
return torch.relu(x)
|