split_cnn.py 1.45 KB
Newer Older
okiss's avatar
okiss committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from torch import nn
import numpy as np
cluster_split=[14,11,11,12,13,12,11,11,15]

c1=[12,5,4,11,19,18,16,16,15,9,14,17,21,22]
c2=[8,1,2,122,121,116,123,3,124,117,118]
c3=[114,115,109,108,102,98,103,110,111,104,93]
c4=[100,101,97,96,95,89,90,84,91,85,92,86]
c5=[82,74,70,75,83,76,71,67,72,77,78,62,61]
c6=[69,65,64,58,57,50,51,59,66,60,52,53]
c7=[44,39,40,45,46,47,41,35,29,36,42]
c8=[38,33,32,25,26,23,27,34,28,24,20]
c9=[13,6,112,105,106,7,30,37,31,129,80,87,79,55,54]
cluster={}
cluster['c1']=c1
cluster['c2']=c2
cluster['c3']=c3
cluster['c4']=c4
cluster['c5']=c5
cluster['c6']=c6
cluster['c7']=c7
cluster['c8']=c8
cluster['c9']=c9
cluster_index=[]
for i in range (1,10):
    cluster_index+=cluster['c'+str(i)]

 #reorganise the electrode according to the cluster
def cluster():
    return cluster_index,cluster_split

class SplitCNN(nn.Module):
    def __init__(self,cluster_split,out_channels,kernel_size,stride,padding):
        super(SplitCNN, self).__init__()

        self.channels = out_channels
        self.padding= padding
        self.kernel = kernel_size
        self.split = cluster_split
okiss's avatar
typo    
okiss committed
40
        self.stride = stride
okiss's avatar
okiss committed
41
42
43
44
45
46
47
48
49
        self.cnn = nn.ModuleList()
        for i in range(len(split)):
            self.cnn.append(nn.conv1d(cluster_split[i],self.channels[i],
            self.kernel,self.stride,self.padding))

    def forward(self, x):
        x = x.split(self.split,0)
        x = [self.cnn[j](x[j]) for j in range(len(self.split))]
        return torch.cat(x,0)