code

class NN(nn.Module):
    def __init__(self , inp_len):
        super(NN ,self).__init__()
        self.fc1 = nn.Linear(inp_len ,256)
        self.fc2 = nn.Linear(256 , 128)
        self.fc3 = nn.Linear(128 , 64)
        self.fc4 = nn.Linear(64,1)
        self.norm1 = nn.BatchNorm1d(256)
        self.norm2 = nn.BatchNorm1d(128)
        self.norm3 = nn.BatchNorm1d(64)
        
        
    def forward(self , X):
        out = F.relu(self.norm1(self.fc1(X)))
        out = F.dropout(out , training =self.training , p =0.5)
        out = F.relu(self.norm2(self.fc2(out)))
        out = F.dropout(out , training =self.training , p =0.3)
        out = F.relu(self.norm3(self.fc3(out))) 
        out = F.dropout(out , training =self.training , p =0.2)
        out = F.relu(self.fc4(out))
        return out 
    
inp_len = X_train.shape[1]
model = NN(inp_len)
model.train()


n_epochs = 20 
lr = 0.07 
optimizer = optim.Adam(params = model.parameters() , lr = lr)
ceriterion = nn.MSELoss()
device = 'cuda' if torch.cuda.is_available() else 'cpu'



## train model 


for epoch in range(n_epochs):
    acc_total  = 0
    loss_total = 0
    for train_idx , out_idx in train_loader:
        train_idx = train_idx.view(-1 ,inp_len)
        
        output = model(train_idx).to(device)
        
        out_idx = out_idx.view(-1,1).to(device)
        loss = ceriterion(output , out_idx)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        accuracy = ((output.argmax(dim=1)==out_idx).float().mean())
        acc_total += accuracy / len(train_loader)
        loss_total += loss / len(train_loader)
    print(f'epoch: {epoch} , accuracy: {acc_total} , loss: {loss_total}')