| """ |
| EchMo Model Implementation |
| End-to-End Radar Echo Segmentation Model |
| """ |
|
|
| import os |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torch.nn.init as init |
| import scipy.io as scio |
|
|
|
|
| |
| def complex_std(x, epsilon=1e-5): |
| mu_R = x.real.mean(dim=3, keepdim=True).detach() |
| mu_I = x.imag.mean(dim=3, keepdim=True).detach() |
| x_R = x.real - mu_R |
| x_I = x.imag - mu_I |
| sigma_R = torch.sqrt(torch.var(x.real, dim=3, keepdim=True, unbiased=False) + epsilon) |
| sigma_I = torch.sqrt(torch.var(x.imag, dim=3, keepdim=True, unbiased=False) + epsilon) |
| x_R /= sigma_R |
| x_I /= sigma_I |
| return x_R + 1j * x_I |
|
|
|
|
| def complex_double2float(x): |
| x_r = x.real.float() |
| x_i = x.imag.float() |
| return x_r + 1j * x_i |
|
|
|
|
| def filter_coef_loader(path, coef_type=None): |
| coef = scio.loadmat(os.path.join(path, 's_tx.mat'))['s_tx'] |
| coef = torch.tensor(coef) |
| |
| if coef_type == 'ori_mf': |
| coef = complex_double2float(torch.conj(coef)) |
| zero = torch.complex(torch.zeros([1, 1], dtype=torch.float32), torch.zeros([1, 1], dtype=torch.float32)) |
| coef = torch.cat((zero, coef), dim=1) |
| return coef |
| elif coef_type == 'tay_mf': |
| win = scio.loadmat(os.path.join(path, 'w_taylor.mat'))['w_taylor'] |
| elif coef_type == 'che_mf': |
| win = scio.loadmat(os.path.join(path, 'w_chebyshev.mat'))['w_chebyshev'] |
| elif coef_type == 'ham_mf': |
| win = scio.loadmat(os.path.join(path, 'w_hamming.mat'))['w_hamming'] |
| elif coef_type == 'han_mf': |
| win = scio.loadmat(os.path.join(path, 'w_hanning.mat'))['w_hanning'] |
| elif coef_type == 'black_mf': |
| win = scio.loadmat(os.path.join(path, 'w_blackman.mat'))['w_blackman'] |
| elif coef_type == 'cos_mf': |
| win = scio.loadmat(os.path.join(path, 'w_cosine.mat'))['w_cosine'] |
| else: |
| win = scio.loadmat(os.path.join(path, 'w_taylor.mat'))['w_taylor'] |
|
|
| win = torch.permute(torch.tensor(win), [1, 0]) |
| coef = complex_double2float(torch.conj(coef * win)) |
| zero = torch.complex(torch.zeros([1, 1], dtype=torch.float32), torch.zeros([1, 1], dtype=torch.float32)) |
| coef = torch.cat((zero, coef), dim=1) |
| return coef |
|
|
|
|
| |
| class ComplexBatchNorm2d(nn.Module): |
| def __init__(self, num_features): |
| super().__init__() |
| self.bn_real = nn.BatchNorm2d(num_features) |
| self.bn_imag = nn.BatchNorm2d(num_features) |
| |
| def forward(self, x): |
| return torch.complex(self.bn_real(x.real), self.bn_imag(x.imag)) |
|
|
|
|
| class ComplexReLU(nn.Module): |
| def __init__(self, activation='leaky_relu'): |
| super(ComplexReLU, self).__init__() |
| self.activation = activation |
| |
| def forward(self, input): |
| if self.activation == 'relu': |
| return nn.ReLU()(input.real) + 1j * nn.ReLU()(input.imag) |
| elif self.activation == 'leaky_relu': |
| return F.leaky_relu(input.real) + 1j * F.leaky_relu(input.imag) |
| elif self.activation == 'celu': |
| return F.celu(input.real) + 1j * F.celu(input.imag) |
| else: |
| return input |
|
|
|
|
| class DoubleConvBlock(nn.Module): |
| def __init__(self, in_ch, out_ch, k_size, pad, dil): |
| super().__init__() |
| self.block = nn.Sequential( |
| nn.Conv2d(in_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), |
| nn.BatchNorm2d(out_ch), |
| nn.LeakyReLU(inplace=True), |
| nn.Conv2d(out_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), |
| nn.BatchNorm2d(out_ch), |
| nn.LeakyReLU(inplace=True) |
| ) |
|
|
| def forward(self, x): |
| return self.block(x) |
|
|
|
|
| class Double3DConvBlock(nn.Module): |
| def __init__(self, in_ch, out_ch, k_size, pad, dil): |
| super().__init__() |
| self.block = nn.Sequential( |
| nn.Conv3d(in_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), |
| nn.BatchNorm3d(out_ch), |
| nn.LeakyReLU(inplace=True), |
| nn.Conv3d(out_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), |
| nn.BatchNorm3d(out_ch), |
| nn.LeakyReLU(inplace=True) |
| ) |
|
|
| def forward(self, x): |
| return self.block(x) |
|
|
|
|
| class ConvBlock(nn.Module): |
| def __init__(self, in_ch, out_ch, k_size, pad, dil): |
| super().__init__() |
| self.block = nn.Sequential( |
| nn.Conv2d(in_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), |
| nn.BatchNorm2d(out_ch), |
| nn.LeakyReLU(inplace=True) |
| ) |
|
|
| def forward(self, x): |
| return self.block(x) |
|
|
|
|
| class PeakConv2D(nn.Module): |
| def __init__(self, in_ch, out_ch, bias, refer_band, guard_band_dop, guard_band, signal_type): |
| super(PeakConv2D, self).__init__() |
| self.signal_type = signal_type |
| if self.signal_type in ('range_doppler', 'angle_doppler'): |
| self.guard_band = (guard_band_dop, guard_band_dop) |
| self.padding = (self.guard_band[1] + refer_band, |
| self.guard_band[1] + refer_band, |
| self.guard_band[0] + refer_band, |
| self.guard_band[0] + refer_band) |
| else: |
| self.guard_band = (guard_band, guard_band) |
| self.padding = (self.guard_band[1] + refer_band, |
| self.guard_band[1] + refer_band, |
| self.guard_band[0] + refer_band, |
| self.guard_band[0] + refer_band) |
| self.refer_band = refer_band |
| self.kernel_size = (self.guard_band[0] + self.guard_band[1] + self.refer_band * 2), 4 |
| self.stride = self.kernel_size |
| self.zero_padding = nn.ZeroPad2d(self.padding) |
| self.bias = bias |
| self.pc_strd = 1 |
| self.peak_conv = nn.Conv2d(in_ch, out_ch, |
| kernel_size=self.kernel_size, |
| padding=0, |
| stride=self.stride, |
| bias=self.bias) |
|
|
| def forward(self, x): |
| h, w = x.shape[2], x.shape[3] |
| rb = self.refer_band |
| gb = self.guard_band |
| k_pk = self.kernel_size |
| N = k_pk[0] * k_pk[1] |
| dtype = x.data.type() |
| b = x.size(0) |
|
|
| x = self.zero_padding(x) |
| p_r = self._get_p_r(b, h, w, N, dtype) |
| x_r = self._sample_x(x, p_r, N) |
| p_c = self._get_p_c(b, h, w, N, dtype) |
| x_c = self._sample_x(x, p_c, N) |
| x_prf = x_c - x_r |
| x_prf = self._reshape_x_prf(k_pk[0], k_pk[1], x_prf) |
| out = self.peak_conv(x_prf) |
| return out |
|
|
| def _get_p_c(self, b, h, w, N, dtype): |
| p_c_x, p_c_y = torch.meshgrid( |
| torch.arange(self.padding[2], h * self.pc_strd + self.padding[2], self.pc_strd), |
| torch.arange(self.padding[0], w * self.pc_strd + self.padding[0], self.pc_strd), indexing='ij') |
| p_c_x = torch.flatten(p_c_x).view(1, 1, h, w).repeat(1, N, 1, 1) |
| p_c_y = torch.flatten(p_c_y).view(1, 1, h, w).repeat(1, N, 1, 1) |
| p_c = torch.cat([p_c_x, p_c_y], 1).type(dtype) |
| p_c = p_c.repeat(b, 1, 1, 1) |
| return p_c |
|
|
| def _gen_prf_grid(self, rb, gb, N, dtype): |
| h_t = -(rb + gb[0]) |
| h_d = rb + gb[0] |
| w_l = -(rb + gb[1]) |
| w_r = rb + gb[1] |
| w_prf = (rb + gb[1]) * 2 + 1 |
| h_prf = (rb + gb[0]) * 2 + 1 |
|
|
| prf_x_idx, prf_y_idx = torch.meshgrid( |
| torch.arange(h_t, h_d + 1), |
| torch.arange(w_l, w_r + 1), indexing='ij') |
|
|
| prf_xt = prf_x_idx[0:rb] |
| prf_xr = prf_x_idx[rb:(h_prf - 2 * rb + 1), (w_prf - rb):w_prf] |
| prf_xd = prf_x_idx[(h_prf - rb):h_prf] |
| prf_xl = prf_x_idx[rb:(h_prf - 2 * rb + 1), 0:rb] |
| prf_x = torch.cat([torch.flatten(prf_xt), |
| torch.flatten(prf_xr), |
| torch.flatten(prf_xd), |
| torch.flatten(prf_xl)], 0) |
| prf_yt = prf_y_idx[0:rb] |
| prf_yr = prf_y_idx[rb:(h_prf - 2 * rb + 1), (w_prf - rb):w_prf] |
| prf_yd = prf_y_idx[(h_prf - rb):h_prf] |
| prf_yl = prf_y_idx[rb:(h_prf - 2 * rb + 1), 0:rb] |
| prf_y = torch.cat([torch.flatten(prf_yt), |
| torch.flatten(prf_yr), |
| torch.flatten(prf_yd), |
| torch.flatten(prf_yl)], 0) |
|
|
| prf = torch.cat([prf_x, prf_y], 0) |
| prf = prf.view(1, 2 * N, 1, 1).type(dtype) |
| return prf |
|
|
| def _get_p_r(self, b, h, w, N, dtype): |
| rb = self.refer_band |
| gb = self.guard_band |
| prf = self._gen_prf_grid(rb, gb, N, dtype) |
| p_c = self._get_p_c(b, h, w, N, dtype) |
| p_r = p_c + prf |
| return p_r |
|
|
| def _sample_x(self, x_pad, p, N): |
| p = p.contiguous().permute(0, 2, 3, 1).floor() |
| b, h, w, _ = p.size() |
| h_pad = x_pad.size(2) |
| w_pad = x_pad.size(3) |
| c = x_pad.size(1) |
| x_pad = x_pad.contiguous().view(b, c, -1) |
| index = p[..., :N] * w_pad + p[..., N:] |
| index = index.contiguous().unsqueeze(dim=1).expand(-1, c, -1, -1, -1).contiguous().view(b, c, -1) |
| index = index.long().to(x_pad.device) |
| x_r = x_pad.gather(dim=-1, index=index).contiguous().view(b, c, h, w, N) |
| return x_r |
|
|
| @staticmethod |
| def _reshape_x_prf(k_h, k_w, x_prf): |
| b, c, h, w, N = x_prf.size() |
| x_prf = torch.cat([x_prf[..., s:s + k_w].contiguous().view(b, c, h, w * k_w) for s in range(0, N, k_w)], dim=-1) |
| x_prf = x_prf.contiguous().view(b, c, h * k_h, w * k_w) |
| return x_prf |
|
|
|
|
| class DoublePeakConv2D(nn.Module): |
| def __init__(self, in_ch, out_ch, bias, rb, gb_dop, gb, sig): |
| super(DoublePeakConv2D, self).__init__() |
| self.bias = bias |
| self.refer_band = rb |
| self.guard_band_dop = gb_dop |
| self.guard_band = gb |
| self.signal_type = sig |
| self.pk_conv1 = PeakConv2D(in_ch, out_ch, bias=self.bias, |
| refer_band=self.refer_band, |
| guard_band_dop=self.guard_band_dop, |
| guard_band=self.guard_band, |
| signal_type=self.signal_type) |
| self.bn1 = nn.BatchNorm2d(out_ch) |
| self.act1 = nn.LeakyReLU(inplace=True) |
| self.pk_conv2 = PeakConv2D(out_ch, out_ch, bias=self.bias, |
| refer_band=self.refer_band, |
| guard_band_dop=self.guard_band_dop, |
| guard_band=self.guard_band, |
| signal_type=self.signal_type) |
| self.bn2 = nn.BatchNorm2d(out_ch) |
| self.act2 = nn.LeakyReLU(inplace=True) |
|
|
| def forward(self, x): |
| x = self.pk_conv1(x) |
| x = self.bn1(x) |
| x = self.act1(x) |
| x = self.pk_conv2(x) |
| x = self.bn2(x) |
| x = self.act2(x) |
| return x |
|
|
|
|
| class ASPPBlock(nn.Module): |
| def __init__(self, in_ch, out_ch, shape): |
| super().__init__() |
| self.shape = shape |
| self.global_avg_pool = nn.AvgPool2d(self.shape) |
| self.conv1_1x1 = nn.Conv2d(in_ch, out_ch, kernel_size=1, padding=0, dilation=1) |
| self.single_conv_block1_1x1 = ConvBlock(in_ch, out_ch, k_size=1, pad=0, dil=1) |
| self.single_conv_block1_3x3 = ConvBlock(in_ch, out_ch, k_size=3, pad=6, dil=6) |
| self.single_conv_block2_3x3 = ConvBlock(in_ch, out_ch, k_size=3, pad=12, dil=12) |
| self.single_conv_block3_3x3 = ConvBlock(in_ch, out_ch, k_size=3, pad=18, dil=18) |
|
|
| def forward(self, x): |
| x1 = F.interpolate(self.global_avg_pool(x), size=self.shape, align_corners=False, mode='bilinear') |
| x1 = self.conv1_1x1(x1) |
| x2 = self.single_conv_block1_1x1(x) |
| x3 = self.single_conv_block1_3x3(x) |
| x4 = self.single_conv_block2_3x3(x) |
| x5 = self.single_conv_block3_3x3(x) |
| x_cat = torch.cat((x2, x3, x4, x5, x1), 1) |
| return x_cat |
|
|
|
|
| class EncodingBranch(nn.Module): |
| def __init__(self, signal_type, n_frames, in_ch=1): |
| super().__init__() |
| self.signal_type = signal_type |
| self.n_frames = n_frames |
| self.in_ch = in_ch |
| self.double_3dconv_block1 = Double3DConvBlock(in_ch=self.in_ch, out_ch=32, |
| k_size=(1+n_frames//2, 3, 3), |
| pad=(0, 1, 1), dil=1) |
| self.doppler_max_pool = nn.MaxPool2d(2, stride=(1, 2)) |
| self.max_pool = nn.MaxPool2d(2, stride=2) |
| self.double_pkc_block = DoublePeakConv2D(in_ch=32, out_ch=64, |
| bias=True, rb=1, gb_dop=1, gb=1, sig=self.signal_type) |
| self.single_conv_block1_1x1 = ConvBlock(in_ch=64, out_ch=128, k_size=1, pad=0, dil=1) |
|
|
| def forward(self, x): |
| x1 = self.double_3dconv_block1(x) |
| x1 = torch.squeeze(x1, 2) |
| x1_pad = F.pad(x1, (0, 0, 0, 0), "constant", 0) |
| x1_down = self.max_pool(x1_pad) |
| x2 = self.double_pkc_block(x1_down) |
| x2_pad = F.pad(x2, (0, 0, 0, 1), "constant", 0) |
| x2_down = self.doppler_max_pool(x2_pad) |
| x3 = self.single_conv_block1_1x1(x2_down) |
| return x2_down, x3 |
|
|
|
|
| class PKCIn4ANY(nn.Module): |
| def __init__(self, n_classes, n_frames, in_ch, w, h): |
| super().__init__() |
| self.n_classes = n_classes |
| self.n_frames = n_frames |
| self.in_ch = in_ch |
| self.rd_shape = (int(w//2), int(h//4)) |
|
|
| self.rd_encoding_branch = EncodingBranch('range_doppler', self.n_frames, self.in_ch) |
| self.rd_aspp_block = ASPPBlock(in_ch=64, out_ch=128, shape=self.rd_shape) |
| self.rd_single_conv_block1_1x1 = ConvBlock(in_ch=640, out_ch=128, k_size=1, pad=0, dil=1) |
|
|
| self.rd_upconv1 = nn.ConvTranspose2d(256, 128, (2, 2), stride=(2, 2)) |
| self.rd_double_conv_block1 = DoubleConvBlock(in_ch=128, out_ch=128, k_size=3, pad=1, dil=1) |
| self.rd_upconv2 = nn.ConvTranspose2d(128, 128, (1, 2), stride=(1, 2)) |
| self.rd_double_conv_block2 = DoubleConvBlock(in_ch=128, out_ch=128, k_size=3, pad=1, dil=1) |
| self.rd_final = nn.Conv2d(in_channels=128, out_channels=n_classes, kernel_size=1) |
|
|
| def forward(self, x_rd): |
| rd_features, rd_latent = self.rd_encoding_branch(x_rd) |
| x2_rd = self.rd_aspp_block(rd_features) |
| x2_rd = self.rd_single_conv_block1_1x1(x2_rd) |
| x4_rd = torch.cat((x2_rd, rd_latent), 1) |
| x4_rd = self.rd_upconv1(x4_rd) |
| x4_rd = self.rd_double_conv_block1(x4_rd) |
| x4_rd = self.rd_upconv2(x4_rd) |
| x4_rd = self.rd_double_conv_block2(x4_rd) |
| x4_rd = self.rd_final(x4_rd) |
| return x4_rd |
|
|
|
|
| def _init_translator_weights(m): |
| if isinstance(m, nn.Linear): |
| torch.nn.init.xavier_uniform_(m.weight) |
| nn.init.constant_(m.bias, 0.) |
| elif isinstance(m, nn.Conv2d): |
| torch.nn.init.xavier_uniform_(m.weight) |
| if m.bias is not None: |
| nn.init.constant_(m.bias, 0.) |
| elif isinstance(m, nn.BatchNorm2d): |
| nn.init.uniform_(m.weight, 0., 1.) |
| nn.init.constant_(m.bias, 0.) |
|
|
|
|
| class Leanable_Match_FilterV5(nn.Module): |
| def __init__(self, k_size, in_ch, out_ch, lmf_coef, activation='relu', bias=True): |
| super(Leanable_Match_FilterV5, self).__init__() |
| self.k_size = k_size |
| self.stride = 1 |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.activation = activation |
| self.bias = bias |
| self.lmf_real = nn.Conv2d(self.in_ch, self.out_ch, |
| kernel_size=(1, k_size), stride=(1, 1), |
| bias=self.bias) |
| self.lmf_imag = nn.Conv2d(self.in_ch, self.out_ch, |
| kernel_size=(1, k_size), stride=(1, 1), |
| bias=self.bias) |
| init.kaiming_uniform_(self.lmf_real.weight) |
| init.kaiming_uniform_(self.lmf_imag.weight) |
| if self.bias: |
| init.zeros_(self.lmf_real.bias) |
| init.zeros_(self.lmf_imag.bias) |
| self.cause_padding = torch.nn.ZeroPad2d(padding=(0, self.k_size - 1, 0, 0)) |
| self.cmplx_relu = ComplexReLU(self.activation) |
|
|
| def forward(self, x): |
| inputR = self.cause_padding(x.real) |
| inputI = self.cause_padding(x.imag) |
| x_RR = self.lmf_real(inputR) |
| x_II = self.lmf_imag(inputI) |
| x_RI = self.lmf_imag(inputR) |
| x_IR = self.lmf_real(inputI) |
| x = x_RR - x_II + 1j * (x_RI + x_IR) |
| x = self.cmplx_relu(x) |
| return x |
|
|
|
|
| class LMF_Chain(nn.Module): |
| def __init__(self, configs): |
| super(LMF_Chain, self).__init__() |
| self.lmf_coef = filter_coef_loader(os.path.join(configs.dataset_path, configs.coef_path), configs.lmf_coefs) |
| self.in_ch = 1 |
| self.Chain = nn.Sequential() |
| for l in range(configs.lmf_layers): |
| self.Chain.add_module('lmfV5-{0}'.format(l), |
| Leanable_Match_FilterV5(k_size=self.lmf_coef.shape[1], |
| in_ch=self.in_ch, |
| out_ch=configs.lmf_channels, |
| lmf_coef=self.lmf_coef, |
| activation=configs.activation, |
| bias=configs.bias)) |
| self.in_ch = configs.lmf_channels |
| |
| def forward(self, x): |
| x = self.Chain(x) |
| return x |
|
|
|
|
| class Multi_Order_LDC(nn.Module): |
| def __init__(self, in_ch, out_ch, orders, activation='relu', bias=True): |
| super(Multi_Order_LDC, self).__init__() |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.orders = orders |
| self.activation = activation |
| self.bias = bias |
| self.ldc_list = nn.ModuleList() |
| self.pad_list = nn.ModuleList() |
| for n in range(self.orders): |
| ldc = nn.Conv2d(self.in_ch, self.out_ch, kernel_size=(n+2, 1), bias=self.bias) |
| pad = torch.nn.ZeroPad2d(padding=(0, 0, n+1, 0)) |
| self.in_ch = self.out_ch |
| self.ldc_list.append(ldc) |
| self.pad_list.append(pad) |
| self.cmplx_relu = ComplexReLU(self.activation) |
| self._initialize_MLDC_weights() |
|
|
| def _initialize_MLDC_weights(self): |
| for n, conv in enumerate(self.ldc_list): |
| nn.init.kaiming_normal_(conv.weight) |
| if self.bias: |
| nn.init.zeros_(conv.bias) |
|
|
| def forward(self, x): |
| B, P, R = x.shape[0], x.shape[2], x.shape[3] |
| x = x.unsqueeze(dim=1).broadcast_to(B, self.orders, x.shape[1], P, R) |
| y_r = torch.zeros(B, self.orders, self.out_ch, P, R).to(x.device) |
| y_i = torch.zeros(B, self.orders, self.out_ch, P, R).to(x.device) |
| for n in range(self.orders): |
| x_n_real = self.ldc_list[n](self.pad_list[n](x[:,n,...].real)) |
| x_n_imag = self.ldc_list[n](self.pad_list[n](x[:,n,...].imag)) |
| y_r[:,n,...] = x_n_real |
| y_i[:,n,...] = x_n_imag |
| y_r = y_r.view(B, self.orders*self.out_ch, P, R) |
| y_i = y_i.view(B, self.orders*self.out_ch, P, R) |
| y = torch.complex(y_r, y_i) |
| y = self.cmplx_relu(y) |
| return y |
|
|
|
|
| class Group_Purturble_DF(nn.Module): |
| def __init__(self, in_ch, out_ch, k_size, groups, activation='relu', bias=True): |
| super(Group_Purturble_DF, self).__init__() |
| self.k_size = k_size |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.groups = groups |
| self.activation = activation |
| self.bias = bias |
| self.gdf_cos = nn.Conv2d(self.in_ch*self.groups, self.out_ch*self.groups, |
| kernel_size=(k_size, 1), stride=(k_size, 1), |
| padding=0, groups=self.groups, bias=self.bias) |
| self.gdf_sin = nn.Conv2d(self.in_ch*self.groups, self.out_ch*self.groups, |
| kernel_size=(k_size, 1), stride=(k_size, 1), |
| padding=0, groups=self.groups, bias=self.bias) |
| self._initialize_GDF_weights() |
| self.cmplx_relu = ComplexReLU(self.activation) |
|
|
| def _initialize_GDF_weights(self): |
| init.kaiming_uniform_(self.gdf_cos.weight) |
| init.kaiming_uniform_(self.gdf_sin.weight) |
| if self.bias: |
| init.zeros_(self.gdf_cos.bias) |
| init.zeros_(self.gdf_sin.bias) |
|
|
| def forward(self, x): |
| B, R = x.shape[0], x.shape[3] |
| x_real, x_imag = x.real.to(x.device), x.imag.to(x.device) |
| x_rr = self.gdf_cos(x_real) |
| x_ii = self.gdf_sin(x_imag) |
| x_ri = self.gdf_sin(x_real) |
| x_ir = self.gdf_cos(x_imag) |
| x = x_rr + x_ii - 1j * x_ri + 1j * x_ir |
| x = self.cmplx_relu(x) |
| x = x.view(B, self.groups, self.out_ch, 1, R).squeeze(dim=3) |
| return x |
|
|
|
|
| class EchMo(nn.Module): |
| def __init__(self, configs): |
| super(EchMo, self).__init__() |
| self.cfg = configs |
| self.task_name = configs.task_name |
| self.n_range = configs.n_range |
| self.n_pulse = configs.n_pulse |
| self.n_classes = configs.n_classes |
| self.n_frames = configs.n_frames |
| self.scale = configs.scale |
| self.activation = configs.activation |
| self.bias = configs.bias |
| self.lmf_chains = configs.lmf_chains |
| self.lmf_channels = configs.lmf_channels |
| self.lmf_layers = configs.lmf_layers |
| self.lmf_coefs = configs.lmf_coefs |
| self.ldc_orders = configs.ldc_orders |
| self.ldc_channels = configs.ldc_channels |
| self.df_groups = configs.ldc_orders |
| self.df_channels = configs.df_channels |
| |
| self.LMFChain = nn.ModuleList([LMF_Chain(configs=self.cfg) for _ in range(self.lmf_chains)]) |
| self.MLDC = Multi_Order_LDC(in_ch=self.lmf_channels*self.lmf_chains, |
| out_ch=self.ldc_channels, |
| orders=self.ldc_orders, |
| activation=self.activation, |
| bias=self.bias) |
| self.GDF = Group_Purturble_DF(in_ch=self.ldc_channels, |
| out_ch=self.df_channels, |
| k_size=self.n_pulse, |
| groups=self.df_groups, |
| activation=self.activation, |
| bias=self.bias) |
| self.RDTranslator = PKCIn4ANY(n_classes=self.n_classes, |
| n_frames=self.n_frames, |
| in_ch=self.df_groups, |
| w=self.df_channels, |
| h=self.n_range) |
| self.RDTranslator.apply(_init_translator_weights) |
|
|
| def forward(self, x): |
| if self.scale == 'instance': |
| x = complex_double2float(complex_std(x, epsilon=1e-8)) |
| elif self.scale == 'global': |
| x = complex_double2float(x) |
|
|
| for i in range(self.lmf_chains): |
| if i == 0: |
| x_k = self.LMFChain[i](x) |
| else: |
| x_ki = self.LMFChain[i](x) |
| x_k = torch.concatenate([x_k, x_ki], dim=1) |
| x_h = self.MLDC(x_k) |
| x_d = self.GDF(x_h) |
| gdf_out = x_d |
| x_d = torch.abs(x_d).unsqueeze(dim=2) |
| y = self.RDTranslator(x_d) |
| return y, x_k, x_h, gdf_out |
|
|
