CNN-BCI-cursor/train_cnn_decoder.m at main · XiaoYicong/CNN-BCI-cursor · GitHub

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
function model = train_cnn_decoder(neural_train, cursor_train, dt, win_ms, hop, offset, sgolay)
if nargin<7, sgolay = []; end

D = size(cursor_train,1);
schema = detect_schema(cursor_train, dt);          % auto detect pos/vel
if ~isempty(schema.vel_idx), out_idx = schema.vel_idx; else, out_idx = 1:D; end

win = max(5, round(win_ms/1000/dt));              % ms → sample point

% Training Set Z-Score Statistics
muY = mean(neural_train,2);
sdY = std(neural_train,0,2); sdY(sdY==0) = 1;

% Building the Training Window
[X4D, Ytr] = prepare_cnn_data(neural_train, cursor_train, muY, sdY, win, hop, offset, out_idx);

% CNN (Time Convolution Only) Regressor: Output dimension = length(out_idx) (if regressing vx, vy, then = 2)
out_dim = size(Ytr,1);
layers = [
    imageInputLayer([size(X4D,1) win 1], 'Normalization','none')
    convolution2dLayer([1 7], 32, 'Padding','same')
    reluLayer
    convolution2dLayer([1 7], 64, 'Padding','same')
    reluLayer
    globalAveragePooling2dLayer
    fullyConnectedLayer(64)
    reluLayer
    dropoutLayer(0.3)
    fullyConnectedLayer(out_dim)
    regressionLayer];

opts = trainingOptions('adam', ...
    'InitialLearnRate',1e-3, 'MiniBatchSize',128, 'MaxEpochs',30, ...
    'Shuffle','every-epoch', 'Verbose',false, 'Plots','none', ...
    'L2Regularization',1e-4);

net = trainNetwork(X4D, Ytr', layers, opts);

model = struct('net',net,'muY',muY,'sdY',sdY, ...
    'dt',dt,'win',win,'hop',hop,'offset',offset, ...
    'sgolay',sgolay,'cursor_dim',D,'schema',schema,'out_idx',out_idx);
end