Add files via upload

Author: RusselZHANG

CC_NBDF_Net.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Apr 17 15:23:55 2020
+
+@author: admin
+"""
+import torch
+import torch.nn as nn
+
+
+class CC_NBDF(nn.Module):
+    def __init__(feature_map_num = 64, hidden1_dim= 256, hidden2_dim = 128, num_direction = 2, num_layers = 1 ,biFlag = True):
+
+        super(CRNN_NBDFNet,self).__init__()
+        
+        self.feature_map_num = feature_map_num
+        
+        self.hidden1_dim=hidden1_dim
+        self.hidden2_dim=hidden2_dim
+        
+        self.num_layers= num_layers
+        self.num_direction = num_direction
+        
+        self.output1_dim=self.hidden1_dim*self.num_direction
+        self.output2_dim=self.hidden2_dim*self.num_direction
+        
+        
+
+        self.biFlag=biFlag
+        
+        self.BN = nn.BatchNorm2d(self.feature_map_num)
+        
+        self.relu = torch.nn.ReLU()
+        
+        self.cnn1 = nn.Conv2d(2,self.feature_map_num,(2,1))
+        
+        self.cnn2 = nn.Conv2d(self.feature_map_num,self.feature_map_num,(2,1))
+
+        self.rnn1 = nn.LSTM(input_size=self.feature_map_num, hidden_size = self.hidden1_dim, \
+                        num_layers=self.num_layers,batch_first=True, \
+                        bidirectional=biFlag)
+        
+        self.rnn2 = nn.LSTM(input_size=self.output1_dim,hidden_size = self.hidden2_dim, \
+                        num_layers=self.num_layers,batch_first=True, \
+                        bidirectional=biFlag)
+        
+        
+        self.linearTimeDistributed = nn.Linear(self.output2_dim, 1)
+        
+    def forward(self,inputsignal):
+        
+        
+        cnn1out = self.relu(self.cnn1(inputsignal))           # (512,2,4,192) -> (512,64,3,192)
+
+        
+        while cnn1out.shape[2] != 1:                          # recursive convolution (512,64,3,192) -> (512,64,192)
+            cnn1out = self.relu(self.cnn2(cnn1out))
+            
+        cnn2out = torch.squeeze(cnn1out)
+        cnn2out = torch.transpose(cnn2out,1,2)
+        
+        rnn1out,_ = self.rnn1(cnn2out)
+        
+        rnn2out,_ = self.rnn2(rnn1out)
+            
+        outsignal = torch.sigmoid(self.linearTimeDistributed(rnn2out))  # linear1out.dim =  (1024, 192, 8)  -> (1024, 192, 1)
+            
+        
+        
+        return outsignal

NBDF_Net.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Apr 17 15:23:55 2020
+
+@author: admin
+"""
+import torch
+import torch.nn as nn
+
+
+
+
+
+class NBDF_Net(nn.Module):
+    def __init__(self, n_chanels, hidden1_dim= 256, hidden2_dim = 128, num_direction = 2, num_layers = 1 ,biFlag = True):
+
+        super(NBDFNet,self).__init__()
+        
+        self.n_chanels=n_chanels
+        
+        self.input_dim= 2*self.n_chanels
+        
+        self.hidden1_dim=hidden1_dim
+        self.hidden2_dim=hidden2_dim
+        
+        self.output1_dim=self.hidden1_dim*num_direction
+        self.output2_dim=self.hidden2_dim*num_direction
+        
+        self.num_layers= num_layers
+
+        self.target = target
+        self.biFlag=biFlag
+
+        
+        self.rnn1 = nn.LSTM(input_size=self.input_dim, hidden_size = self.hidden1_dim, \
+                        num_layers=self.num_layers,batch_first=True, \
+                        bidirectional=biFlag)
+        
+        self.rnn2 = nn.LSTM(input_size=self.output1_dim,hidden_size = self.hidden2_dim, \
+                        num_layers=self.num_layers,batch_first=True, \
+                        bidirectional=biFlag)
+
+        
+        self.linearTimeDistributed = nn.Linear(self.output2_dim, 1)
+        
+
+    def forward(self,inputsignal):
+        
+        
+        rnn1out,_ = self.rnn1(inputsignal)
+        
+        rnn2out,_ = self.rnn2(rnn1out)
+
+        outsignal = torch.sigmoid(self.linearTimeDistributed(rnn2out))  # linear1out.dim =  (1024, 192, 512)  -> (1024, 192, 1)
+            
+        
+        
+        return outsignal

NB_Dataset.py

@@ -0,0 +1,91 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Apr 17 15:13:05 2020
+
+@author: admin
+"""
+
+import torch
+import torch.nn as nn
+# import torchvision
+# import torch.nn.functional as F
+from torch import optim
+# import torchvision.transforms as transforms
+from torch.utils.data import DataLoader, Dataset
+
+
+import os,fnmatch
+import numpy as np
+
+
+class NB_Dataset(Dataset):
+
+    def __init__(self, data_path, batchsize = 512, time_steps = 192, shuffle=True):
+        self.data_path = data_path
+        self.time_steps = time_steps
+        self.batchsize = batchsize
+        self.shuffle = shuffle
+        self.on_epoch_end()
+        
+
+    def __getitem__(self, index):
+        
+        batchname = fnmatch.filter(os.listdir(self.data_path),'batch{}*'.format(self.indexes[index]))[0]
+        sample = np.load(self.data_path+ batchname)
+        X = sample['X'][:self.batchsize,:self.time_steps,:].astype('float32')
+        y = sample['mrm'][:self.batchsize,:self.time_steps].reshape(self.batchsize,self.time_steps,1).astype('float32')
+        
+
+        X = torch.from_numpy(X)
+        y = torch.from_numpy(y)
+        
+        return X, y
+
+    def __len__(self):
+        
+        return len(fnmatch.filter(os.listdir(self.data_path),'batch*.npz'))
+    
+    def on_epoch_end(self):
+    #    'Updates indexes after each epoch'
+        self.indexes = np.arange(self.__len__())
+        if self.shuffle == True:
+            np.random.shuffle(self.indexes)
+            
+            
+"""           
+if __name__ == "__main__":
+    
+    train_path = '../Array_position/train_val_batch/train_batch/'
+    val_path = '../Array_position/train_val_batch/validation_batch/'
+    
+    # wavFiles = fnmatch.filter(os.listdir(train_path),'batch10606*.npz')
+    
+    
+    btz = 1
+    
+    train_NBDataset = NBDataset(data_path = train_path, time_steps=192, shuffle = True)
+    val_NBDataset = NBDataset(data_path = val_path,  time_steps=192, shuffle = True)
+    
+
+    train_DataLoader = DataLoader(
+        dataset=train_NBDataset,      # torch TensorDataset format
+        batch_size=btz,      # mini batch size
+        shuffle=True,       # random shuffle for training
+        drop_last=True,
+        num_workers=0,      # subprocesses for loading data
+    )
+
+    val_DataLoader = DataLoader(
+        dataset=val_NBDataset,      # torch TensorDataset format
+        batch_size=btz,      # mini batch size
+        shuffle=True,       # random shuffle for training
+        drop_last=True,
+        num_workers=0,      # subprocesses for loading data
+    )
+    
+    n,m = train_NBDataset[2]
+    print(n.shape,m.shape)
+    
+    for a,b in train_DataLoader:
+        print(a.shape,b.shape)
+""" 

PW_NBDF_Net.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Apr 17 15:23:55 2020
+
+@author: admin
+"""
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+
+
+
+class PW_NBDF(nn.Module):
+    def __init__(self, input_dim = 4, hidden1_dim = 256, hidden2_dim = 128, num_direction = 2, device = 'cuda', num_layers = 1, biFlag = True):
+
+        super(PW_NBDF,self).__init__()
+        
+        self.input_dim=input_dim
+        
+        self.hidden1_dim=hidden1_dim
+        self.hidden2_dim=hidden2_dim
+        
+        self.num_direction = num_direction
+        
+        self.output1_dim=self.hidden1_dim*num_direction
+        self.output2_dim=self.hidden2_dim*num_direction
+        
+        self.num_layers= num_layers
+        self.device = device
+
+
+        self.biFlag=biFlag
+        
+        
+        self.rnn1 = nn.LSTM(input_size=self.input_dim, hidden_size = self.hidden1_dim, \
+                        num_layers=self.num_layers,batch_first=True, \
+                        bidirectional=self.biFlag)
+        
+        self.rnn2 = nn.LSTM(input_size=self.output1_dim,hidden_size = self.hidden2_dim, \
+                        num_layers=self.num_layers,batch_first=True, \
+                        bidirectional=self.biFlag)
+        
+        self.linearTimeDistributed = nn.Linear(self.output2_dim, 1)
+        
+        
+
+
+    def forward(self,inputsignal):
+        
+        B,T,C = inputsignal.shape # (B,T,C)
+        
+        n_pairs = C//2-1  # number of channel of pairs
+        x = torch.zeros(B, n_pairs, T, 4, device = self.device)  # (B,N,T,4)
+        
+        for i in range(n_pairs):
+            x[:, i, :, :2] = inputsignal[:,:,:2]
+            x[:, i, :, 2:] = inputsignal[:,:,(i+1)*2:(i+2)*2]
+        x = x.view(B*n_pairs, T, 4)                             # (B*N , T,  4)
+        rnn1out, _ = self.rnn1(x)
+        rnn1out = rnn1out.view(B, n_pairs, T, self.output1_dim) # (B, N, T, Dim1*2)
+        rnn1out_combined = torch.mean(rnn1out,dim = 1)          # (B, T, Dim1*2)
+
+        rnn2out,_ = self.rnn2(rnn1out_combined)                 # (B, T, Dim2*2)
+        outsignal = torch.sigmoid(self.linearTimeDistributed(rnn2out)).squeeze()  # (B, T) 
+            
+
+        
+        return outsignal       # 1D mask output