simpler version with adaptiveavgpool2d suggested by alamie
@@ -147,16 +147,12 @@ class EEGResNet(nn.Sequential):
self.eval()
if self.final_pool_length == 'auto':
- out = self(np_to_var(np.ones(
- (1, self.in_chans, self.input_window_samples, 1),
- dtype=np.float32)))
- n_out_time = out.cpu().data.numpy().shape[2]
- self.final_pool_length = n_out_time
- pool_dilation = 1
+ self.add_module('mean_pool', nn.AdaptiveAvgPool2d((1, 1)))
else:
pool_dilation = int(cur_dilation[0]), int(cur_dilation[1])
self.add_module('mean_pool', AvgPool2dWithConv(
- (self.final_pool_length, 1), (1, 1), dilation=pool_dilation))
+ (self.final_pool_length, 1), (1, 1),
+ dilation=pool_dilation))
self.add_module('conv_classifier',
nn.Conv2d(n_cur_filters, self.n_classes,
(1, 1), bias=True))
