-
Notifications
You must be signed in to change notification settings - Fork 41
/
fbcnet.py
210 lines (169 loc) · 8.03 KB
/
fbcnet.py
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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
import torch
import torch.nn as nn
class Conv2dWithConstraint(nn.Conv2d):
def __init__(self, *args, weight_norm=True, max_norm=1, **kwargs):
self.max_norm = max_norm
self.weight_norm = weight_norm
super(Conv2dWithConstraint, self).__init__(*args, **kwargs)
def forward(self, x):
if self.weight_norm:
self.weight.data = torch.renorm(self.weight.data,
p=2,
dim=0,
maxnorm=self.max_norm)
return super(Conv2dWithConstraint, self).forward(x)
class LinearWithConstraint(nn.Linear):
def __init__(self, *args, weight_norm=True, max_norm=1, **kwargs):
self.max_norm = max_norm
self.weight_norm = weight_norm
super(LinearWithConstraint, self).__init__(*args, **kwargs)
def forward(self, x):
if self.weight_norm:
self.weight.data = torch.renorm(self.weight.data,
p=2,
dim=0,
maxnorm=self.max_norm)
return super(LinearWithConstraint, self).forward(x)
class VarLayer(nn.Module):
def __init__(self, dim):
super(VarLayer, self).__init__()
self.dim = dim
def forward(self, x):
return x.var(dim=self.dim, keepdim=True)
class StdLayer(nn.Module):
def __init__(self, dim):
super(StdLayer, self).__init__()
self.dim = dim
def forward(self, x):
return x.std(dim=self.dim, keepdim=True)
class LogVarLayer(nn.Module):
def __init__(self, dim):
super(LogVarLayer, self).__init__()
self.dim = dim
def forward(self, x):
return torch.log(
torch.clamp(x.var(dim=self.dim, keepdim=True), 1e-6, 1e6))
class MeanLayer(nn.Module):
def __init__(self, dim):
super(MeanLayer, self).__init__()
self.dim = dim
def forward(self, x):
return x.mean(dim=self.dim, keepdim=True)
class MaxLayer(nn.Module):
def __init__(self, dim):
super(MaxLayer, self).__init__()
self.dim = dim
def forward(self, x):
ma, ima = x.max(dim=self.dim, keepdim=True)
return ma
class swish(nn.Module):
def __init__(self):
super(swish, self).__init__()
def forward(self, x):
return x * torch.sigmoid(x)
class FBCNet(nn.Module):
r'''
An Efficient Multi-view Convolutional Neural Network for Brain-Computer Interface. For more details, please refer to the following information.
- Paper: Mane R, Chew E, Chua K, et al. FBCNet: A multi-view convolutional neural network for brain-computer interface[J]. arXiv preprint arXiv:2104.01233, 2021.
- URL: https://arxiv.org/abs/2104.01233
- Related Project: https://github.com/ravikiran-mane/FBCNet
Below is a recommended suite for use in emotion recognition tasks:
.. code-block:: python
dataset = DEAPDataset(io_path=f'./deap',
root_path='./data_preprocessed_python',
chunk_size=512,
num_baseline=1,
baseline_chunk_size=512,
offline_transform=transforms.BandSignal(),
online_transform=transforms.ToTensor(),
label_transform=transforms.Compose([
transforms.Select('valence'),
transforms.Binary(5.0),
]))
model = FBCNet(num_classes=2,
num_electrodes=32,
chunk_size=512,
in_channels=4,
num_S=32)
Args:
num_electrodes (int): The number of electrodes. (default: :obj:`28`)
chunk_size (int): Number of data points included in each EEG chunk. (default: :obj:`1000`)
in_channels (int): The number of channels of the signal corresponding to each electrode. If the original signal is used as input, in_channels is set to 1; if the original signal is split into multiple sub-bands, in_channels is set to the number of bands. (default: :obj:`9`)
num_S (int): The number of spatial convolution block. (default: :obj:`32`)
num_classes (int): The number of classes to predict. (default: :obj:`2`)
temporal (str): The temporal layer used, with options including VarLayer, StdLayer, LogVarLayer, MeanLayer, and MaxLayer, used to compute statistics using different techniques in the temporal dimension. (default: :obj:`LogVarLayer`)
stride_factor (int): The stride factor. (default: :obj:`4`)
weight_norm (bool): Whether to use weight renormalization technique in Conv2dWithConstraint. (default: :obj:`True`)
'''
def __init__(self,
num_electrodes: int = 20,
chunk_size: int = 1000,
in_channels: int = 9,
num_S: int = 32,
num_classes: int = 2,
temporal: str = 'LogVarLayer',
stride_factor: int = 4,
weight_norm: bool = True):
super(FBCNet, self).__init__()
self.num_electrodes = num_electrodes
self.chunk_size = chunk_size
self.num_classes = num_classes
self.in_channels = in_channels
self.num_S = num_S
self.temporal = temporal
self.stride_factor = stride_factor
self.weight_norm = weight_norm
assert chunk_size % stride_factor == 0, f'chunk_size should be divisible by stride_factor, chunk_size={chunk_size},stride_factor={stride_factor} does not meet the requirements.'
self.scb = self.SCB(num_S,
num_electrodes,
self.in_channels,
weight_norm=weight_norm)
if temporal == 'VarLayer':
self.temporal_layer = VarLayer(dim=3)
elif temporal == 'StdLayer':
self.temporal_layer = StdLayer(dim=3)
elif temporal == 'LogVarLayer':
self.temporal_layer = LogVarLayer(dim=3)
elif temporal == 'MeanLayer':
self.temporal_layer = MeanLayer(dim=3)
elif temporal == 'MaxLayer':
self.temporal_layer = MaxLayer(dim=3)
else:
raise NotImplementedError
self.last_layer = self.last_block(self.num_S * self.in_channels *
self.stride_factor,
num_classes,
weight_norm=weight_norm)
def SCB(self, num_S, num_electrodes, in_channels, weight_norm=True):
return nn.Sequential(
Conv2dWithConstraint(in_channels,
num_S * in_channels, (num_electrodes, 1),
groups=in_channels,
max_norm=2,
weight_norm=weight_norm,
padding=0),
nn.BatchNorm2d(num_S * in_channels), swish())
def last_block(self, in_channels, out_channels, weight_norm=True):
return nn.Sequential(
LinearWithConstraint(in_channels,
out_channels,
max_norm=0.5,
weight_norm=weight_norm), nn.LogSoftmax(dim=1))
def forward(self, x: torch.Tensor) -> torch.Tensor:
r'''
Args:
x (torch.Tensor): EEG signal representation, the ideal input shape is :obj:`[n, in_channel, num_electrodes, chunk_size]`. Here, :obj:`n` corresponds to the batch size
Returns:
torch.Tensor[number of sample, number of classes]: the predicted probability that the samples belong to the classes.
'''
x = self.scb(x)
x = x.reshape([
*x.shape[0:2], self.stride_factor,
int(x.shape[3] / self.stride_factor)
])
x = self.temporal_layer(x)
x = torch.flatten(x, start_dim=1)
x = self.last_layer(x)
return x
def feature_dim(self):
return self.num_S * self.in_channels * self.stride_factor