Coverage for src / models / churn_mlp.py: 100%

1import torch 

2import torch.nn as nn 

3 

4 

5class ChurnMLP(nn.Module): 

6 """ 

7 Arquitetura de Rede Neural Densa (Feed-Forward) para classificação binária de Churn. 

8 A rede possui Input -> 64 -> 32 -> 1, com funções de ativação ReLU e Dropout de 30% 

9 para prevenção de overfitting. 

10 """ 

11 

12 def __init__(self, input_dim: int): 

13 super(ChurnMLP, self).__init__() 

14 self.fc1 = nn.Linear(input_dim, 64) 

15 self.relu1 = nn.ReLU() 

16 self.dropout1 = nn.Dropout(0.3) 

17 

18 self.fc2 = nn.Linear(64, 32) 

19 self.relu2 = nn.ReLU() 

20 self.dropout2 = nn.Dropout(0.3) 

21 

22 # A última camada não tem ativação Sigmoid no forward pois usaremos BCEWithLogitsLoss 

23 self.fc3 = nn.Linear(32, 1) 

24 

25 def forward(self, x: torch.Tensor) -> torch.Tensor: 

26 x = self.fc1(x) 

27 x = self.relu1(x) 

28 x = self.dropout1(x) 

29 

30 x = self.fc2(x) 

31 x = self.relu2(x) 

32 x = self.dropout2(x) 

33 

34 x = self.fc3(x) 

35 return x