|
- import matplotlib.pyplot as plt
- import torch
- import numpy as np
- import argparse
- from dataset_tools.dataset_bpsk_labelReceive import BPSK2Send01_ChunkPairsDataset
- from torch.utils.data import DataLoader
- from end2end_transmitter_channel_receiver_model_0926 import Transmitter_simple_bpsk,\
- Transmitter_seq2seq_useall_bpsk
- import torch.optim as optim
- import torch.nn as nn
- from torch.optim import lr_scheduler
-
-
- def EuclideanDistances(a,b):
- sq_a = a**2
- sum_sq_a = torch.sum(sq_a,dim=1).unsqueeze(1) # m->[m, 1]
- sq_b = b**2
- sum_sq_b = torch.sum(sq_b,dim=1).unsqueeze(0) # n->[1, n]
- bt = b.t()
- return torch.sqrt(sum_sq_a+sum_sq_b-2*a.mm(bt))
-
-
- def count_parameters(model):
- return sum(p.numel() for p in model.parameters() if p.requires_grad)
-
- def get_args():
- """Text generation arguments."""
- parser = argparse.ArgumentParser()
- parser.add_argument('--feature_dim', default=128, type=int)
- parser.add_argument('--hidden_dim', default=256, type=int)
- parser.add_argument('--n_layers', default=2, type=int)
-
- parser.add_argument('--dec_emb_dim', default=1024, type=int)
- parser.add_argument('--dec_dropout', default=0.5, type=float)
-
- parser.add_argument('--lr', default=0.002, type=float)
- parser.add_argument('--weight_decay', default=0.01, type=float)
- parser.add_argument('--batch_size', default=2, type=int)
- parser.add_argument('--do_chunk', action='store_true')
- parser.add_argument('--chunk_size', default=128, type=int)
-
- parser.add_argument('--do_chunk_withlinear', action='store_true')
-
- args = parser.parse_args()
- # print(args)
- return args
-
-
- def do_eval(model, test_dataloader, args, step=None, is_test=True):
- model.eval()
- test_distances = []
- with torch.no_grad():
- for i, data in enumerate(test_dataloader):
- bpsk_singal, send01 = data
- data = send01.float().cuda()
- labels = bpsk_singal.float().cuda()
-
- if args.do_chunk is False:
- data = data.reshape(args.batch_size, -1, 1)
- labels = labels.reshape(args.batch_size, -1, args.chunk_size)
- else:
- data = data.view(-1, 1)
-
- output = model(data) # [1000, 1, 2]
- # prediction = F.log_softmax(output, dim=-1)
- # gt = F.softmax(labels, dim=-1)
-
- test_avg_dist = torch.mean(torch.pairwise_distance(output, labels)).item()
- test_distances.append(test_avg_dist)
-
- plot_chunks_num = 5
-
- if i==0:
- signal0type = bpsk_singal.detach().cpu().numpy()
- pred0type = output.detach().cpu().numpy()
- if args.do_chunk_withlinear:
- np.save('./gifs/bpsk_learning-npys/step_{}_signal.npy'.format(step), signal0type)
- np.save('./gifs/bpsk_learning-npys/step_{}_pred.npy'.format(step), signal0type)
- else:
- np.save('./gifs/bpsk_all_learning-npys/step_{}_signal.npy'.format(step), signal0type)
- np.save('./gifs/bpsk_all_learning-npys/step_{}_pred.npy'.format(step), signal0type)
-
- signal0type = signal0type[i, :args.chunk_size * plot_chunks_num]
- pred0type = pred0type[i, :, :]
- pred0type = pred0type[:plot_chunks_num, :]
- signal0type = signal0type.reshape(-1)
- pred0type = pred0type.reshape(-1)
-
- xdata = np.linspace(0, len(signal0type) - 1, len(signal0type))
-
- plt.plot(signal0type)
- plt.scatter(xdata, signal0type)
- plt.plot(pred0type, 'r_', lw=5)
- plt.ylim([-1.05, 1.05])
- plt.tight_layout()
- if args.do_chunk_withlinear:
- plt.savefig('./gifs/bpsk_learning/step_{}.png'.format(step), dpi=300)
- else:
- plt.savefig('./gifs/bpsk_all_learning/step_{}.png'.format(step), dpi=300)
- plt.close()
- # plt.pause(0.3)
-
-
- print('> test avg dist: {}\n'.format(np.mean(test_distances)))
-
- # test_loss = F.
- # if is_test:
- # print('Eval on test_dataset, acc is: [{}%]\n'.format(np.mean(test_acc_list) * 100))
- # else:
- # print('Eval on val_dataset, acc is: [{}%]'.format(np.mean(test_acc_list) * 100))
-
-
-
- def setup_model_with_datasets(raw_train_data_folder, raw_test_data_folder, args):
- # dataset01: [bs, 128], with [bs, 1] --do_chunk=False
- # dataset02: [bs, 1000, 128], with [bs, 1000, 1] --do_chunk=False
- dataset = BPSK2Send01_ChunkPairsDataset(raw_train_data_folder, do_chunk=args.do_chunk)
- train_dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) # , num_workers=20, drop_last=True)
- test_dataset = BPSK2Send01_ChunkPairsDataset(raw_test_data_folder, do_chunk=args.do_chunk)
- test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False) # , num_workers=20, drop_last=True)
-
- if args.do_chunk:
- my_model = Transmitter_simple_bpsk(do_chunk=args.do_chunk)
- else:
- if args.do_chunk_withlinear:
- my_model = Transmitter_simple_bpsk(do_chunk=args.do_chunk)
- else:
- my_model = Transmitter_seq2seq_useall_bpsk(args.feature_dim,
- args.hidden_dim,
- args.n_layers,
- 1)
- my_model.cuda()
-
- return my_model, train_dataloader, test_dataloader
-
- def train(raw_train_data_folder, raw_test_data_folder, args):
-
-
- my_model, train_dataloader, test_dataloader = \
- setup_model_with_datasets(raw_train_data_folder, raw_test_data_folder, args)
-
- print(f'The model has {count_parameters(my_model):,} trainable parameters')
- optimizer = optim.Adam(my_model.parameters(), lr=args.lr)#, weight_decay=args.weight_decay, betas=(0.9, 0.98), eps=1e-8)
- scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=2e-5, verbose=False)
-
- # criterion = nn.KLDivLoss()
- criterion = nn.MSELoss()
-
- epoch_loss = 0
- tmp_idx = 0
-
- # plt.ion()
- # plt.show()
-
- for epoch in range(200):
- for (j, data) in enumerate(train_dataloader):
- my_model.train()
- optimizer.zero_grad()
-
- bpsk_signal, send01 = data
- data = send01.float().cuda()
- labels = bpsk_signal.float().cuda()
-
- if args.do_chunk is False:
- data = data.reshape(args.batch_size, -1, 1)
- labels = labels.reshape(args.batch_size, -1, args.chunk_size)
- else:
- data = data.view(-1, 1)
-
- output = my_model(data) # [1000, 1, 2]
- # prediction = F.log_softmax(output, dim=-1)
- # gt = F.softmax(labels, dim=-1)
- # loss = criterion(prediction, gt)
-
- loss = criterion(output, labels)
- loss.backward()
-
- optimizer.step()
- scheduler.step()
- epoch_loss += loss.item()
- tmp_idx += 1
-
- if (j+1) % 10 == 0:
- print('[{}]/[{}], loss is: {}, lr {}'.format(j, len(train_dataloader),
- epoch_loss / tmp_idx,
- optimizer.state_dict()['param_groups'][0]['lr']
- )
- )
- train_avg_dist = torch.mean(torch.pairwise_distance(output, labels)).item()
- print('> train avg dist: {}'.format(train_avg_dist))
- do_eval(my_model, test_dataloader, args, step=j * (epoch+1), is_test=False)
-
- # if j==499:
- # exit()
- #
- # for i, each01 in enumerate(send01):
- # if (each01 == 0).item():
- # signal0type = bpsk_singal[i, :].detach().cpu().numpy()
- # pred0type = output[i, :].detach().cpu().numpy()
- # xdata = np.linspace(0, len(signal0type) - 1, len(signal0type))
- #
- # plt.plot(signal0type)
- # plt.scatter(xdata, signal0type)
- # plt.plot(pred0type, 'r_', lw=5)
- # plt.ylim([-1.05, 1.05])
- # plt.tight_layout()
- # plt.savefig('./gifs/bpsk_learning/step_{}.png'.format(j), dpi=300)
- # plt.close()
- # # plt.pause(0.3)
- # break
-
-
- if __name__ == '__main__':
- raw_train_data_folder = '/userhome/wave_training_old/bpsk_sending_signal/train1500_send_wave/'
- raw_test_data_folder = '/userhome/wave_training_old/bpsk_sending_signal/testA100_send_wave/'
- testB_data_dir = '/userhome/wave_training_old/bpsk_sending_signal/testB100_send_wave/'
- testC_data_dir = '/userhome/wave_training_old/bpsk_sending_signal/testC50_send_wave/'
-
- args = get_args()
-
- train(raw_train_data_folder, raw_test_data_folder, args)
-
-
|