From 05efbbd9218ff034d8bc9bf730c8ca7a10dc6a2f Mon Sep 17 00:00:00 2001 From: namsaraeva Date: Thu, 16 May 2024 16:13:12 +0200 Subject: [PATCH] change in_features in MLP --- src/sparcscore/ml/models.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/sparcscore/ml/models.py b/src/sparcscore/ml/models.py index f4dfb36..f53626a 100644 --- a/src/sparcscore/ml/models.py +++ b/src/sparcscore/ml/models.py @@ -52,7 +52,7 @@ def make_layers(self, cfg, in_channels, batch_norm = True): else: layers += [(f"conv{i}", conv2d), (f"relu{i}", nn.ReLU(inplace=True))] in_channels = v - i +=1 + i += 1 return nn.Sequential(OrderedDict(layers)) def make_layers_MLP(self, cfg_MLP, cfg, regression = False): @@ -70,7 +70,10 @@ def make_layers_MLP(self, cfg_MLP, cfg, regression = False): nn.Sequential: A sequential model representing the MLP architecture. """ # get output feature size of CNN with chosen configuration - in_features = int(cfg[-2]) * 2 * 2 + if regression: + in_features = int(cfg[-2]) * 4 * 4 + else: + in_features = int(cfg[-2]) * 2 * 2 layers = [] i = 0 @@ -173,6 +176,7 @@ def __init__(self, super(VGG2_regression, self).__init__() self.norm = nn.BatchNorm2d(in_channels) + self.features = self.make_layers(self.cfgs[cfg], in_channels) self.classifier = self.make_layers_MLP(self.cfgs_MLP[cfg_MLP], self.cfgs[cfg], regression=True) # regression is set to True to make the final layer a single output @@ -189,7 +193,7 @@ def forward(self, x): x = torch.flatten(x, 1) print("x.shape after flatten", x.shape) - + x = self.classifier(x) return x