diff --git a/.gitignore b/.gitignore
index 7ca6e7e00f515f5981a7853d48f28d1848a19fa8..c0c8a49b92343851eac1b2395a1efd1e34be1914 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,8 @@ raw_sf7_cross_instance/
matlab/raw_sf7_cross_instance.zip
*.pkl
*.png
-*.pyc
\ No newline at end of file
+*.pyc
+*.log
+*.yaml
+*.json
+*wandb
\ No newline at end of file
diff --git a/neural_enhanced_demodulation/matlab/evaluation.m b/neural_enhanced_demodulation/matlab/evaluation.m
index 09901f6626f7a95361ad8c5f998c8614d9f4e77f..5e4bc242e56666915223d67ada43911552fb4dfe 100644
--- a/neural_enhanced_demodulation/matlab/evaluation.m
+++ b/neural_enhanced_demodulation/matlab/evaluation.m
@@ -8,40 +8,71 @@ set(fig,'DefaultAxesFontSize',20);
set(fig,'DefaultAxesFontWeight','bold');
data_root = '';
-color_list = linspecer(2);
+color_list = linspecer(4);
BW=125000;
SF=7;
-SNR_list=[-40:15];
% nelora_file='evaluation/sf7_v1_';
-nelora_file='evaluation/sf7_125k_new_';
+% nelora_file='evaluation/sf7_125k_new_';
+% improved_nelora_file='evaluation/sf7_traineval5e-4_';
+nelora_file='evaluation/sf7_v4_';
+improved_nelora_file='evaluation/sf7_v2.3_maskcnn_no_mask_';
+prune_nelora_file='evaluation/prune_sf7_v1_';
+baseline_file='evaluation/baseline_error_matrix_';
+
% nelora_file='evaluation/sf7_v3_';
+% nelora_file='evaluation/sf7_traineval5e-4_';
+SNR_list=[-40:15];
+SNR_improved_list=[-40:15];
+SNR_prone_list=[-40:15];
+SNR_list_baseline=[-30:0];
-SNR_list_baseline=-30:0;
-baseline_file='evaluation/baseline_error_matrix_';
+% improved nelora
+name_str=[improved_nelora_file,num2str(SF),'_',num2str(BW),'.mat'];
+error_path = [data_root,name_str];
+a2 = load(error_path);
+error_matrix = a2.error_matrix;
+error_rate = 1-error_matrix;
+plot(SNR_improved_list,error_rate,"-.*",'LineWidth',3,'color',color_list(3,:));
+hold on;
+% nelora
name_str=[nelora_file,num2str(SF),'_',num2str(BW),'.mat'];
error_path = [data_root,name_str];
a = load(error_path);
error_matrix = a.error_matrix;
-% error_matrix_info = a.error_matrix_info;
+error_rate = 1-error_matrix;
+plot(SNR_list,error_rate,"-.*",'LineWidth',3,'color',color_list(1,:));
+hold on;
-plot(SNR_list,1-error_matrix,"-.*",'LineWidth',3,'color',color_list(1,:));
+% prone
+name_str=[prune_nelora_file,num2str(SF),'_',num2str(BW),'.mat'];
+error_path = [data_root,name_str];
+a = load(error_path);
+error_matrix = a.error_matrix;
+error_rate = 1-error_matrix;
+plot(SNR_prone_list,error_rate,"-.*",'LineWidth',2,'color',color_list(4,:));
hold on;
-name_str=[baseline_file,num2str(SF),'_',num2str(BW),'.mat'];
+% baseline
+name_str=[baseline_file,num2str(SF),'_',num2str(BW),'V1.mat'];
error_path = [data_root,name_str];
a = load(error_path);
error_matrix = a.error_matrix;
plot(SNR_list_baseline,1-error_matrix,"-.*",'LineWidth',2,'color',color_list(2,:));
hold on;
-legend('NeLoRA','Baseline')
+% legend('Imroved NeLoRA', 'NeLoRA', 'Prone', 'Baseline')
+% legend('NeLoRA', 'Baseline')
+% legend('Imroved NeLoRA', 'NeLoRA', 'Baseline')
+% legend('Imroved NeLoRA', 'NeLoRA', 'Baseline')
+legend('Imroved NeLoRA', 'NeLoRA', 'Prone', 'Baseline')
% legend('abs_baselineNELoRa','Baseline')
xlabel('SNR (dB)'); % x label
ylabel('SER'); % y label
title('Decode SER for SF=7')
-xlim([-30,-10]);
+xlim([-20,-10]);
set(gcf,'WindowStyle','normal','Position', [200,200,640,360]);
+grid on;
saveas(gcf,[data_root,'res/',num2str(SF),'_',num2str(BW),'.pdf'])
diff --git a/neural_enhanced_demodulation/matlab/evaluation/sf7_v2.3_maskcnn_no_mask_7_125000.mat b/neural_enhanced_demodulation/matlab/evaluation/sf7_v2.3_maskcnn_no_mask_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..b9500550d89e86ea0512af9e8ff6815d2b8d935d
Binary files /dev/null and b/neural_enhanced_demodulation/matlab/evaluation/sf7_v2.3_maskcnn_no_mask_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/None_7_125000.mat b/neural_enhanced_demodulation/pytorch/None_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..4fb75fbec0c34a24a5e73980f1b6cf270e33dcd6
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/None_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/config.py b/neural_enhanced_demodulation/pytorch/config.py
index d4abb27fac5d559e5a8ba25e54efad2fff72b9bb..f387d928511de0d75e4bf3782c9ed6d73dcd8ff5 100644
--- a/neural_enhanced_demodulation/pytorch/config.py
+++ b/neural_enhanced_demodulation/pytorch/config.py
@@ -144,7 +144,7 @@ def create_parser():
help='Choose the root path to rf signals.',
)
- parser.add_argument('--network', type=str, default='end2end', choices=['end2end', 'end2end_fig4', 'end2end_real'])
+ parser.add_argument('--network', type=str, default='end2end_debug', choices=['end2end', 'end2end_debug', 'end2end_fig4', 'end2end_real'])
parser.add_argument(
'--feature_name',
@@ -196,6 +196,11 @@ def create_parser():
default='checkpoints'
)
parser.add_argument('--dir_comment', type=str, default='None')
+ parser.add_argument('--checkpoint_dir',
+ type=str,
+ default='checkpoints')
+ parser.add_argument('--dir_comment', type=str, default='add cos anealing')
+ parser.add_argument('--wandb_entity', type=str, default='jackthewizard')
parser.add_argument('--sample_dir', type=str, default='samples')
parser.add_argument('--testing_dir', type=str, default='testing')
# parser.add_argument('--load', type=str, default='pre_trained')
diff --git a/neural_enhanced_demodulation/pytorch/end2end.py b/neural_enhanced_demodulation/pytorch/end2end.py
index 936cc2ce1da5145d8db455815b98bde8ba2ccf37..7d954b25e52d7675fe0a1fc03d6a7abaae782d58 100644
--- a/neural_enhanced_demodulation/pytorch/end2end.py
+++ b/neural_enhanced_demodulation/pytorch/end2end.py
@@ -194,6 +194,9 @@ def training_loop(
* Saves checkpoint every opts.checkpoint_every iterations
* Saves generated samples every opts.sample_every iterations
"""
+ import wandb;
+ wandb.init(project="CSE891", entity=opts.wandb_entity)
+
loss_spec = torch.nn.MSELoss(reduction='mean')
loss_class = nn.CrossEntropyLoss()
# Create generators and discriminators
@@ -285,6 +288,7 @@ def training_loop(
G_Image_loss = opts.scaling_for_imaging_loss * g_y_pix_loss
G_Class_loss = opts.scaling_for_classification_loss * g_y_class_loss
G_Y_loss = G_Image_loss + G_Class_loss
+ wandb.log({"G_Y_loss": G_Y_loss, "G_Image_loss": G_Image_loss, "G_class_loss": G_Class_loss})
G_Y_loss.backward()
g_optimizer.step()
@@ -312,6 +316,7 @@ def training_loop(
if iteration % opts.checkpoint_every == 0:
checkpoint(iteration, mask_CNN, C_XtoY, opts)
+ wandb.finish()
test_iter_X = iter(testing_dataloader_X)
test_iter_Y = iter(testing_dataloader_Y)
iter_per_epoch_test = min(len(test_iter_X), len(test_iter_Y))
diff --git a/neural_enhanced_demodulation/pytorch/end2end_improve.py b/neural_enhanced_demodulation/pytorch/end2end_improve.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc63ef27f7952c1e8a73789e86ab9f03775127b5
--- /dev/null
+++ b/neural_enhanced_demodulation/pytorch/end2end_improve.py
@@ -0,0 +1,408 @@
+# end2end_improve.py
+
+from __future__ import division
+import os
+
+import warnings
+
+warnings.filterwarnings("ignore")
+
+# Torch imports
+import torch
+import torch.fft
+import torch.nn as nn
+import torch.optim as optim
+
+# Numpy & Scipy imports
+import numpy as np
+import scipy.io
+
+import cv2
+# Local imports
+from utils import to_var, to_data, spec_to_network_input
+from models.model_components import maskCNNModel, DenoisingNet, Autoencoder, classificationHybridModel
+import torch.autograd.profiler as profiler
+import time
+
+SEED = 11
+
+# Set the random seed manually for reproducibility.
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+if torch.cuda.is_available():
+ torch.cuda.manual_seed(SEED)
+
+
+def print_models(Model):
+ """Prints model information for the generators and discriminators.
+ """
+ print(" Model ")
+ print("---------------------------------------")
+ print(Model)
+ print("---------------------------------------")
+
+
+def create_model(opts):
+ """Builds the generators and discriminators.
+ """
+
+ maskCNN = maskCNNModel(opts)
+
+ C_XtoY = classificationHybridModel(conv_dim_in=opts.y_image_channel,
+ conv_dim_out=opts.n_classes,
+ conv_dim_lstm=opts.conv_dim_lstm)
+
+ if torch.cuda.is_available():
+ maskCNN.cuda()
+ C_XtoY.cuda()
+ print('Models moved to GPU.')
+
+ return maskCNN, C_XtoY
+
+
+def create_model2(opts):
+ """Builds the generators and discriminators.
+ """
+
+ denoisenet = Autoencoder()
+
+ C_XtoY = classificationHybridModel(conv_dim_in=opts.y_image_channel,
+ conv_dim_out=opts.n_classes,
+ conv_dim_lstm=opts.conv_dim_lstm)
+
+ if torch.cuda.is_available():
+ denoisenet.cuda()
+ C_XtoY.cuda()
+ print('Models moved to GPU.')
+
+ return denoisenet, C_XtoY
+
+
+def checkpoint(iteration, mask_CNN, C_XtoY, opts):
+ """Saves the parameters of both generators G_YtoX, G_XtoY and discriminators D_X, D_Y.
+ """
+
+ mask_CNN_path = os.path.join(opts.checkpoint_dir, str(iteration) + '_maskCNN.pkl')
+ torch.save(mask_CNN.state_dict(), mask_CNN_path)
+
+ C_XtoY_path = os.path.join(opts.checkpoint_dir, str(iteration) + '_C_XtoY.pkl')
+ torch.save(C_XtoY.state_dict(), C_XtoY_path)
+
+
+def load_checkpoint(opts):
+ """Loads the generator and discriminator models from checkpoints.
+ """
+
+ maskCNN_path = os.path.join(opts.checkpoint_dir, str(opts.load_iters) + '_maskCNN.pkl')
+
+ maskCNN = maskCNNModel(opts)
+
+ maskCNN.load_state_dict(torch.load(
+ maskCNN_path, map_location=lambda storage, loc: storage),
+ strict=False)
+
+ C_XtoY_path = os.path.join(opts.checkpoint_dir, str(opts.load_iters) + '_C_XtoY.pkl')
+
+ C_XtoY = classificationHybridModel(conv_dim_in=opts.x_image_channel,
+ conv_dim_out=opts.n_classes,
+ conv_dim_lstm=opts.conv_dim_lstm)
+
+ C_XtoY.load_state_dict(torch.load(
+ C_XtoY_path, map_location=lambda storage, loc: storage),
+ strict=False)
+
+ if torch.cuda.is_available():
+ maskCNN.cuda()
+ C_XtoY.cuda()
+ print('Models moved to GPU.')
+
+ return maskCNN, C_XtoY
+
+
+def merge_images(sources, targets, batch_size, image_channel):
+ """Creates a grid consisting of pairs of columns, where the first column in
+ each pair contains images source images and the second column in each pair
+ contains images generated by the CycleGAN from the corresponding images in
+ the first column.
+ """
+ _, _, h, w = sources.shape
+ row = int(np.sqrt(batch_size))
+ column = int(batch_size / row)
+ merged = np.zeros([image_channel, row * h, column * w * 2])
+ for idx, (s, t) in enumerate(zip(sources, targets)):
+ i = idx // column
+ j = idx % column
+ merged[:, i * h:(i + 1) * h, (j * 2) * w:(j * 2 + 1) * w] = s
+ merged[:, i * h:(i + 1) * h, (j * 2 + 1) * w:(j * 2 + 2) * w] = t
+ return merged.transpose(1, 2, 0)
+
+
+def save_samples(iteration, fixed_Y, fixed_X, mask_CNN, opts):
+ """Saves samples from both generators X->Y and Y->X.
+ """
+ fake_Y = mask_CNN(fixed_X)
+ fixed_X = to_data(fixed_X)
+
+ Y, fake_Y = to_data(fixed_Y), to_data(fake_Y)
+
+ merged = merge_images(fixed_X, fake_Y, opts.batch_size, opts.y_image_channel)
+
+ path = os.path.join(opts.sample_dir,
+ 'sample-{:06d}-Y.png'.format(iteration))
+ merged = np.abs(merged[:, :, 0] + 1j * merged[:, :, 1])
+ merged = (merged - np.amin(merged)) / (np.amax(merged) - np.amin(merged)) * 255
+ merged = cv2.flip(merged, 0)
+ cv2.imwrite(path, merged)
+ print('Saved {}'.format(path))
+
+
+def save_samples_separate(iteration, fixed_Y, fixed_X, mask_CNN, opts,
+ name_X_test, labels_Y_test, saved_dir):
+ """Saves samples from both generators X->Y and Y->X.
+ """
+ fake_Y = mask_CNN(fixed_X)
+
+ fixed_Y, fake_Y, fixed_X = to_data(fixed_Y), to_data(fake_Y), to_data(fixed_X)
+
+ for batch_index in range(opts.batch_size):
+ if batch_index < len(name_X_test):
+ path_src = os.path.join(saved_dir, name_X_test[batch_index])
+ groundtruth_image = (
+ np.squeeze(fixed_Y[batch_index, :, :, :]).transpose(1, 2, 0))
+
+ groundtruth_image = np.abs(groundtruth_image[:, :, 0] + 1j * groundtruth_image[:, :, 1])
+ groundtruth_image = (groundtruth_image - np.amin(groundtruth_image)) / (
+ np.amax(groundtruth_image) - np.amin(groundtruth_image)) * 255
+ cv2.imwrite(path_src + '_groundtruth_' + str(iteration) + '.png', groundtruth_image)
+
+ fake_image = (
+ np.squeeze(fake_Y[batch_index, :, :, :]).transpose(1, 2, 0))
+ fake_image = np.abs(fake_image[:, :, 0] + 1j * fake_image[:, :, 1])
+ fake_image = (fake_image - np.amin(fake_image)) / (np.amax(fake_image) - np.amin(fake_image)) * 255
+ cv2.imwrite(path_src + '_fake_' + str(iteration) + '.png', fake_image)
+
+ raw_image = (
+ np.squeeze(fixed_X[batch_index, :, :, :]).transpose(1, 2, 0))
+ raw_image = np.abs(raw_image[:, :, 0] + 1j * raw_image[:, :, 1])
+ raw_image = (raw_image - np.amin(raw_image)) / (np.amax(raw_image) - np.amin(raw_image)) * 255
+ cv2.imwrite(path_src + '_raw_' + str(iteration) + '.png', raw_image)
+
+ # print('Saved {}'.format(path))
+
+
+def training_loop(training_dataloader_X, training_dataloader_Y, testing_dataloader_X,
+ testing_dataloader_Y, opts):
+ """Runs the training loop.
+ * Saves checkpoint every opts.checkpoint_every iterations
+ * Saves generated samples every opts.sample_every iterations
+ """
+ import wandb;
+ wandb.init(project="CSE891")
+
+ loss_spec = torch.nn.MSELoss(reduction='mean')
+ loss_class = nn.CrossEntropyLoss()
+ # Create generators and discriminators
+ if opts.load:
+ mask_CNN, C_XtoY = load_checkpoint(opts)
+ else:
+ # mask_CNN, C_XtoY = create_model(opts)
+ mask_CNN, C_XtoY = create_model2(opts)
+
+ g_params = list(mask_CNN.parameters()) + list(C_XtoY.parameters())
+ g_optimizer = optim.Adam(g_params, opts.lr, [opts.beta1, opts.beta2])
+
+ iter_X = iter(training_dataloader_X)
+ iter_Y = iter(training_dataloader_Y)
+
+ test_iter_X = iter(testing_dataloader_X)
+ test_iter_Y = iter(testing_dataloader_Y)
+
+ # Get some fixed data from domains X and Y for sampling. These are images that are held
+ # constant throughout training, that allow us to inspect the model's performance.
+ fixed_X, name_X_fixed = next(test_iter_X)
+ fixed_X = to_var(fixed_X)
+
+ fixed_Y, name_Y_fixed = next(test_iter_Y)
+ fixed_Y = to_var(fixed_Y)
+ # print("Fixed_X {}".format(fixed_X.shape))
+ fixed_X_spectrum_raw = torch.stft(input=fixed_X, n_fft=opts.stft_nfft, hop_length=opts.stft_overlap,
+ win_length=opts.stft_window, pad_mode='constant');
+ fixed_X_spectrum = spec_to_network_input(fixed_X_spectrum_raw, opts)
+ # print("Fixed {}".format(fixed_X_spectrum.shape))
+
+ fixed_Y_spectrum_raw = torch.stft(input=fixed_Y, n_fft=opts.stft_nfft, hop_length=opts.stft_overlap,
+ win_length=opts.stft_window, pad_mode='constant');
+ fixed_Y_spectrum = spec_to_network_input(fixed_Y_spectrum_raw, opts)
+
+ iter_per_epoch = min(len(iter_X), len(iter_Y))
+
+ for iteration in range(1, opts.train_iters + 1):
+ if iteration % iter_per_epoch == 0:
+ iter_X = iter(training_dataloader_X)
+ iter_Y = iter(training_dataloader_Y)
+
+ # images_X, name_X = iter_X.next()
+ images_X, name_X = next(iter_X)
+ labels_X_mapping = list(
+ map(lambda x: int(x.split('_')[5]), name_X))
+ images_X, labels_X = to_var(images_X), to_var(
+ torch.tensor(labels_X_mapping))
+ # images_Y, name_Y = iter_Y.next()
+ images_Y, name_Y = next(iter_Y)
+ labels_Y_mapping = list(
+ map(lambda x: int(x.split('_')[5]), name_Y))
+ images_Y, labels_Y = to_var(images_Y), to_var(
+ torch.tensor(labels_Y_mapping))
+
+ # ============================================
+ # TRAIN THE GENERATOR
+ # ============================================
+
+ images_X_spectrum_raw = torch.stft(input=images_X, n_fft=opts.stft_nfft, hop_length=opts.stft_overlap,
+
+ win_length=opts.stft_window, pad_mode='constant');
+ images_X_spectrum = spec_to_network_input(images_X_spectrum_raw, opts)
+
+ images_Y_spectrum_raw = torch.stft(input=images_Y, n_fft=opts.stft_nfft, hop_length=opts.stft_overlap,
+ win_length=opts.stft_window, pad_mode='constant');
+ images_Y_spectrum = spec_to_network_input(images_Y_spectrum_raw, opts)
+
+ #########################################
+ ## FILL THIS IN: X--Y ##
+ #########################################
+ if iteration % 50 == 0:
+ print("Iteration: {}/{}".format(iteration, opts.train_iters))
+ mask_CNN.train()
+ C_XtoY.train()
+ fake_Y_spectrum = mask_CNN(images_X_spectrum) # B, 2, 128, 33
+ # 2. Compute the generator loss based on domain Y
+ g_y_pix_loss = loss_spec(fake_Y_spectrum, images_Y_spectrum) # Try other loss functions
+ labels_X_estimated = C_XtoY(fake_Y_spectrum) # B, 128
+ g_y_class_loss = loss_class(labels_X_estimated, labels_X) # Try other classification loss functions
+ g_optimizer.zero_grad()
+ G_Image_loss = opts.scaling_for_imaging_loss * g_y_pix_loss
+ G_Class_loss = opts.scaling_for_classification_loss * g_y_class_loss
+ G_Y_loss = G_Image_loss + G_Class_loss
+ G_Y_loss.backward()
+ g_optimizer.step()
+ current_lr = g_optimizer.param_groups[0]['lr']
+ wandb.log({"G_Y_loss": G_Y_loss, "G_Image_loss": G_Image_loss, "G_class_loss": G_Class_loss, 'lr': current_lr})
+
+ # Print the log info
+ if iteration % opts.log_step == 0:
+ print(
+ 'Iteration [{:5d}/{:5d}] | G_Y_loss: {:6.4f}| G_Image_loss: {:6.4f}| G_Class_loss: {:6.4f}'
+ .format(iteration, opts.train_iters,
+ G_Y_loss.item(),
+ G_Image_loss.item(),
+ G_Class_loss.item()))
+
+ # Save the generated samples
+ if (iteration % opts.sample_every == 0) and (not opts.server):
+ # save_samples(iteration, fixed_Y_spectrum, fixed_X_spectrum, mask_CNN, opts)
+ save_samples_separate(iteration, fixed_Y_spectrum, fixed_X_spectrum,
+ mask_CNN, opts, name_X_fixed, name_Y_fixed, opts.sample_dir)
+
+ # Save the model parameters
+ if iteration % opts.checkpoint_every == 0:
+ checkpoint(iteration, mask_CNN, C_XtoY, opts)
+
+ ###########################################
+ ################## Testing ################
+ ###########################################
+ if iteration % opts.log_step == 0:
+ error_matrix, error_matrix_count, error_matrix_info, saved_data = testing(mask_CNN, C_XtoY, testing_dataloader_X, testing_dataloader_Y, opts)
+ wandb.log({"error_matrix": error_matrix[11:41].mean(), 'error_matrix_count': error_matrix_count.mean()})
+ wandb.finish()
+
+ scipy.io.savemat(
+ opts.root_path + '/' + opts.dir_comment + '_' + str(opts.sf) + '_' + str(opts.bw) + '.mat',
+ dict(error_matrix=error_matrix,
+ error_matrix_count=error_matrix_count,
+ error_matrix_info=error_matrix_info))
+
+
+ with open('test.npy', 'wb') as f:
+ np.save(f, saved_data)
+ f.close()
+
+
+def testing(mask_CNN, C_XtoY, testing_dataloader_X,
+ testing_dataloader_Y, opts):
+ mask_CNN.eval()
+ C_XtoY.eval()
+ test_iter_X = iter(testing_dataloader_X)
+ test_iter_Y = iter(testing_dataloader_Y)
+ # iter_per_epoch_test = min(len(test_iter_X), len(test_iter_Y))
+ iter_per_epoch_test = 3000
+
+ error_matrix = np.zeros([len(opts.snr_list), 1], dtype=float)
+ error_matrix_count = np.zeros([len(opts.snr_list), 1], dtype=int)
+
+ error_matrix_info = []
+
+ saved_data = {}
+ for iteration in range(iter_per_epoch_test):
+ # images_X_test, name_X_test = test_iter_X.next()
+ images_X_test, name_X_test = next(test_iter_X)
+
+ code_X_test_mapping = list(
+ map(lambda x: float(x.split('_')[0]), name_X_test))
+
+ snr_X_test_mapping = list(
+ map(lambda x: int(x.split('_')[1]), name_X_test))
+
+ instance_X_test_mapping = list(
+ map(lambda x: int(x.split('_')[4]), name_X_test))
+
+ labels_X_test_mapping = list(
+ map(lambda x: int(x.split('_')[5]), name_X_test))
+
+ images_X_test, labels_X_test = to_var(images_X_test), to_var(
+ torch.tensor(labels_X_test_mapping))
+
+ # images_Y_test, labels_Y_test = test_iter_Y.next()
+ images_Y_test, labels_Y_test = next(test_iter_Y)
+ images_Y_test = to_var(images_Y_test)
+
+ images_X_test_spectrum_raw = torch.stft(input=images_X_test, n_fft=opts.stft_nfft,
+ hop_length=opts.stft_overlap, win_length=opts.stft_window,
+ pad_mode='constant');
+ images_X_test_spectrum = spec_to_network_input(images_X_test_spectrum_raw, opts)
+
+ images_Y_test_spectrum_raw = torch.stft(input=images_Y_test, n_fft=opts.stft_nfft,
+ hop_length=opts.stft_overlap, win_length=opts.stft_window,
+ pad_mode='constant');
+ images_Y_test_spectrum = spec_to_network_input(images_Y_test_spectrum_raw, opts)
+ fake_Y_test_spectrum = mask_CNN(images_X_test_spectrum)
+ labels_X_estimated = C_XtoY(fake_Y_test_spectrum)
+ saved_sample = to_data(labels_X_estimated)
+
+ for i, label in enumerate(to_data(labels_X_test)):
+ if label not in saved_data.keys():
+ saved_data[label] = []
+ saved_data[label].append(saved_sample[i])
+ else:
+ saved_data[label].append(saved_sample[i])
+ _, labels_X_test_estimated = torch.max(labels_X_estimated, 1)
+
+ test_right_case = (labels_X_test_estimated == labels_X_test)
+ test_right_case = to_data(test_right_case)
+
+ for batch_index in range(opts.batch_size):
+ try:
+ snr_index = opts.snr_list.index(snr_X_test_mapping[batch_index])
+ error_matrix[snr_index] += test_right_case[batch_index]
+ error_matrix_count[snr_index] += 1
+ error_matrix_info.append([instance_X_test_mapping[batch_index], code_X_test_mapping[batch_index],
+ snr_X_test_mapping[batch_index],
+ labels_X_test_estimated[batch_index].cpu().data.int(),
+ labels_X_test[batch_index].cpu().data.int()])
+ except:
+ print("Something else went wrong")
+ if iteration % opts.log_step == 0:
+ print('Testing Iteration [{:5d}/{:5d}]'
+ .format(iteration, iter_per_epoch_test))
+ error_matrix = np.divide(error_matrix, error_matrix_count+0.0001)
+ error_matrix_info = np.array(error_matrix_info)
+ return error_matrix, error_matrix_count, error_matrix_info, saved_data
diff --git a/neural_enhanced_demodulation/pytorch/main.py b/neural_enhanced_demodulation/pytorch/main.py
index 42528b85271740cb6f8a87da5e9e376ab05b63ed..d2608fc0d3df12d784f7115840ca960b7dae3a87 100644
--- a/neural_enhanced_demodulation/pytorch/main.py
+++ b/neural_enhanced_demodulation/pytorch/main.py
@@ -1,32 +1,26 @@
"""Main script for project."""
from __future__ import print_function
-
-import os
-
+from utils import generate_dataset, create_dir, set_gpu, print_opts
import config
import datasets.data_loader as data_loader
-import end2end
-from prune import go_prune
-from utils import create_dir, generate_dataset, print_opts, set_gpu
+import end2end, end2end_improve
+import end2end, end2end_improve
+import os
def main(opts):
"""Loads the data, creates checkpoint and sample directories, and starts the training loop.
"""
[files_train, files_test
- ] = generate_dataset(
- opts.root_path, opts.data_dir, opts.ratio_bt_train_and_test,
- opts.code_list, opts.snr_list, opts.bw_list, opts.sf_list,
- opts.instance_list, opts.sorting_type
- )
+ ] = generate_dataset(opts.root_path, opts.data_dir, opts.ratio_bt_train_and_test,
+ opts.code_list, opts.snr_list, opts.bw_list, opts.sf_list,
+ opts.instance_list, opts.sorting_type)
# Create train and test dataloaders for images from the two domains X and Y
training_dataloader_X, testing_dataloader_X = data_loader.lora_loader(
- opts, files_train, files_test, False
- )
+ opts, files_train, files_test, False)
training_dataloader_Y, testing_dataloader_Y = data_loader.lora_loader(
- opts, files_train, files_test, True
- )
+ opts, files_train, files_test, True)
# Create checkpoint and sample directories
create_dir(opts.checkpoint_dir)
@@ -37,18 +31,13 @@ def main(opts):
# Start training
set_gpu(opts.free_gpu_id)
- # if prune mode
- if opts.prune:
- go_prune(opts, testing_dataloader_X, testing_dataloader_Y)
- return
-
# select the model
if opts.network == 'end2end':
- end2end.training_loop(
- training_dataloader_X, training_dataloader_Y, testing_dataloader_X,
- testing_dataloader_Y, opts
- )
-
+ end2end.training_loop(training_dataloader_X, training_dataloader_Y, testing_dataloader_X,
+ testing_dataloader_Y, opts)
+ elif opts.network == 'end2end_improve':
+ end2end_debug.training_loop(training_dataloader_X, training_dataloader_Y, testing_dataloader_X,
+ testing_dataloader_Y, opts)
if __name__ == "__main__":
parser = config.create_parser()
diff --git a/neural_enhanced_demodulation/pytorch/models/model_components.py b/neural_enhanced_demodulation/pytorch/models/model_components.py
index b1a963d17417a9e85981e00911d772300daffca0..9e9e4882f9c0ef0afbc0a0e819c51f04dcc89969 100644
--- a/neural_enhanced_demodulation/pytorch/models/model_components.py
+++ b/neural_enhanced_demodulation/pytorch/models/model_components.py
@@ -24,11 +24,20 @@ class classificationHybridModel(nn.Module):
self.drop1 = nn.Dropout(0.2)
self.drop2 = nn.Dropout(0.5)
self.act = nn.ReLU()
-
- def forward(self, x):
- out = self.act(self.conv1(x))
- out = self.pool1(out)
- out = out.view(out.size(0), -1)
+ self._initialize_weights()
+
+ def _initialize_weights(self):
+ # import ipdb; ipdb.set_trace()
+ for module in self.modules():
+ if isinstance(module, nn.Linear):
+ nn.init.kaiming_normal_(module.weight, mode='fan_in')
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+
+ def forward(self, x): # 8, 2, 128, 33
+ out = self.act(self.conv1(x)) # 8, 16, 64, 17
+ out = self.pool1(out) # 8, 16, 32, 8
+ out = out.reshape(out.size(0), -1) #
out = self.act(self.dense(out))
out = self.drop2(out)
@@ -85,7 +94,7 @@ class maskCNNModel(nn.Module):
nn.BatchNorm2d(8), nn.ReLU(),
)
-
+
self.lstm = nn.LSTM(
opts.conv_dim_lstm,
opts.lstm_dim,
@@ -95,21 +104,136 @@ class maskCNNModel(nn.Module):
self.fc1 = nn.Linear(2 * opts.lstm_dim, opts.fc1_dim)
self.fc2 = nn.Linear(opts.fc1_dim, opts.freq_size * opts.y_image_channel)
-
- def forward(self, x):
+ self._initialize_weights() #initialize with kaiming initialization for all the linear layers
+
+ def _initialize_weights(self):
+ # import ipdb; ipdb.set_trace()
+ for module in self.modules():
+ if isinstance(module, nn.Linear):
+ nn.init.kaiming_normal_(module.weight, mode='fan_in')
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+
+ def forward(self, x): # x: 8, 2, 128, 33
out = x.transpose(2, 3).contiguous()
- out = self.conv(out)
- out = out.transpose(1, 2).contiguous()
- out = out.view(out.size(0), out.size(1), -1)
- out, _ = self.lstm(out)
+ out = self.conv(out) # out: 8, 8, 33, 128
+ out = out.transpose(1, 2).contiguous() # out: 8, 33, 8, 128
+ out = out.view(out.size(0), out.size(1), -1) # out: 8, 33, 1024
+ out, _ = self.lstm(out) # out: 8, 33, 800
out = F.relu(out)
- out = self.fc1(out)
+ out = self.fc1(out) # out: 8, 33, 600
out = F.relu(out)
- out = self.fc2(out)
-
- out = out.view(out.size(0), out.size(1), self.opts.y_image_channel, -1)
- out = torch.sigmoid(out)
- out = out.transpose(1, 2).contiguous()
- out = out.transpose(2, 3).contiguous()
- masked = out * x # out is mask, masked is denoised
+ out = self.fc2(out) # out: 8, 33, 256
+
+ out = out.view(out.size(0), out.size(1), self.opts.y_image_channel, -1) # out: 8, 33, 2, 128
+ # out = torch.sigmoid(out) #
+ out = out.transpose(1, 2).contiguous() # out: 8, 2, 33, 128
+ out = out.transpose(2, 3).contiguous() # out: 8, 2, 128, 33
+ masked = out
+ # masked = out * x # out is mask, masked is denoised
return masked
+
+
+class ResidualBlock(nn.Module):
+ def __init__(self, channels):
+ super(ResidualBlock, self).__init__()
+ self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+ self.bn1 = nn.BatchNorm2d(channels)
+ self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+ self.bn2 = nn.BatchNorm2d(channels)
+
+ def forward(self, x):
+ identity = x
+
+ out = F.relu(self.bn1(self.conv1(x)))
+ out = self.bn2(self.conv2(out))
+
+ out += identity
+ out = F.relu(out)
+
+ return out
+
+
+class AEResidualBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
+ super(AEResidualBlock, self).__init__()
+ self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
+ self.bn1 = nn.BatchNorm2d(out_channels)
+ self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size, stride, padding)
+ self.bn2 = nn.BatchNorm2d(out_channels)
+ self.relu = nn.ReLU(inplace=True)
+
+ self.skip = nn.Sequential()
+ if stride != 1 or in_channels != out_channels:
+ self.skip = nn.Sequential(
+ nn.Conv2d(in_channels, out_channels, 1, stride),
+ nn.BatchNorm2d(out_channels)
+ )
+
+ def forward(self, x):
+ identity = self.skip(x)
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+class DenoisingNet(nn.Module):
+ def __init__(self, opts=None):
+ super(DenoisingNet, self).__init__()
+ self.initial_conv = nn.Conv2d(2, 64, kernel_size=3, padding=1)
+ self.res_block1 = ResidualBlock(64)
+ self.res_block2 = ResidualBlock(64)
+ # self.res_block3 = ResidualBlock(64)
+ # self.res_block4 = ResidualBlock(64)
+ self.final_conv = nn.Conv2d(64, 2, kernel_size=3, padding=1)
+
+ def forward(self, x):
+ output = F.relu(self.initial_conv(x))
+ output = self.res_block1(output)
+ output = self.res_block2(output)
+ # output = self.res_block3(output)
+ # output = self.res_block4(output)
+ output = self.final_conv(output)
+ mask = output.sigmoid()
+ output = mask * x
+ return output
+
+class Autoencoder(nn.Module):
+ def __init__(self):
+ super(Autoencoder, self).__init__()
+
+ self.encoder = nn.Sequential(
+ AEResidualBlock(2, 32),
+ AEResidualBlock(64, 64),
+ AEResidualBlock(64, 128)
+ )
+ self.decoder = nn.Sequential(
+ AEResidualBlock(128, 64),
+ AEResidualBlock(64, 32),
+ AEResidualBlock(32, 2)
+ )
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.decoder(x)
+ return x
+
+if __name__ == '__main__':
+ # create mo
+ # model = DenoisingNet()
+ model = Autoencoder()
+
+ # Input data
+ batch_size = 8
+ dummy_input = torch.randn(batch_size, 2, 128, 33)
+
+ output = model(dummy_input)
+ print(output.shape) # [5, 2, 128, 33]
diff --git a/neural_enhanced_demodulation/pytorch/prune_sf7_v1_7_125000.mat b/neural_enhanced_demodulation/pytorch/prune_sf7_v1_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..2d4f0afd4582a668945848e1c3e8889f2398d549
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/prune_sf7_v1_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/sf7_v2.1_7_125000.mat b/neural_enhanced_demodulation/pytorch/sf7_v2.1_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..383f170471ddfacbe21e109752c3461597ed4d4a
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/sf7_v2.1_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/sf7_v2.2_2resblks_mask_7_125000.mat b/neural_enhanced_demodulation/pytorch/sf7_v2.2_2resblks_mask_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..2be0f9c41c4615a9775b289a27283ddb0ab1fd70
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/sf7_v2.2_2resblks_mask_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/sf7_v2.2_7_125000.mat b/neural_enhanced_demodulation/pytorch/sf7_v2.2_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..99ae6d9830f9486e65b7ad7dff225a52b032e054
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/sf7_v2.2_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/sf7_v2.3_maskcnn_no_mask_7_125000.mat b/neural_enhanced_demodulation/pytorch/sf7_v2.3_maskcnn_no_mask_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..b9500550d89e86ea0512af9e8ff6815d2b8d935d
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/sf7_v2.3_maskcnn_no_mask_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/sf7_v2.4_AE_7_125000.mat b/neural_enhanced_demodulation/pytorch/sf7_v2.4_AE_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..6ed4cae766f5d0c058a05c209536a046272f6f50
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/sf7_v2.4_AE_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/sf7_v2_7_125000.mat b/neural_enhanced_demodulation/pytorch/sf7_v2_7_125000.mat
new file mode 100644
index 0000000000000000000000000000000000000000..6512a078d03897feb2f11bd76fe14a2c2bb1a661
Binary files /dev/null and b/neural_enhanced_demodulation/pytorch/sf7_v2_7_125000.mat differ
diff --git a/neural_enhanced_demodulation/pytorch/shs/train.sh b/neural_enhanced_demodulation/pytorch/shs/train.sh
index 9a5f5dba05d537fde9703c50934eccf7a1d4619b..de27d711883d449a2913d265dd4b1ee36424bf1a 100644
--- a/neural_enhanced_demodulation/pytorch/shs/train.sh
+++ b/neural_enhanced_demodulation/pytorch/shs/train.sh
@@ -1,10 +1,14 @@
-python main.py --dir_comment sf7_v1 \
+
+python main.py --dir_comment sf7_v2.4_AE \
--batch_size 16 \
+ --lr 0.0005 \
--root_path . \
--data_dir /home/liyifa11/MyCodes/AIoT891/project1/sf7_125k \
--groundtruth_code 35 \
--normalization \
--train_iter 100000 \
+ --log_step 2000 \
--ratio_bt_train_and_test 0.8 \
- --network end2end
+ --network end2end_debug \
+ --free_gpu_id 1
diff --git a/neural_enhanced_demodulation/pytorch/test.npy b/neural_enhanced_demodulation/pytorch/test.npy
index e29e468f9d2d89f14c5067ea8fd694fcc3733207..304b1bf5603868822917c944e05b8e2a6d068391 100644
Binary files a/neural_enhanced_demodulation/pytorch/test.npy and b/neural_enhanced_demodulation/pytorch/test.npy differ