Genetic Algorithm for a Neural Network training

I am trying to set up a basic CNN that is going to be used later for GA training. I want to use XOR algorithm eventually, thought I am very far from understanding how to set it up.

Here I want to share the prototype what is going to be used for generating individuals. The NN will evolve only in terms of weights and biases, not structure. It takes the structure of a CNN and outputs the fitness function and the chromosome (an array of weights). A question to those who have experience with GA, am I on a right track? From here on I want to proceed with writing up the GA part. Any advise on where to start is appreciated. Right now I am reading this paper, and also looking into PyGAD library. Thank you in advance! My code is below:

import torch
from pylab import *

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

from torchvision import datasets
from torchvision.transforms import ToTensor
train_data = datasets.MNIST(
    root = 'data',
    train = True,                         
    transform = ToTensor(), 
    download = True,            
)
test_data = datasets.MNIST(
    root = 'data', 
    train = False, 
    transform = ToTensor()
)

from torch.utils.data import DataLoader
batch_size = 32

loaders = {
    'train' : torch.utils.data.DataLoader(train_data, 
                                          batch_size=batch_size, 
                                          shuffle=True, 
                                          num_workers=1),
    
    'test'  : torch.utils.data.DataLoader(test_data, 
                                          batch_size=batch_size, 
                                          shuffle=True, 
                                          num_workers=1),
}
loaders

import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(         
            nn.Conv2d(
                in_channels=1,              
                out_channels=16,            
                kernel_size=5,              
                stride=1,                   
                padding=2,                  
            ),                                                    
            nn.MaxPool2d(2,2),
            nn.BatchNorm2d(16),
            nn.Flatten()  
        )
        self.conv2 = nn.Sequential(           
            nn.Linear(16 * 14 * 14, batch_size * 8),  
            nn.ReLU(),                      
            nn.Flatten()               
        )
        self.out = nn.Sequential(nn.Linear(batch_size*8, 10), nn.Softmax(dim=1))
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        output = self.out(x)
        return output

cnn = CNN()
print(cnn)

activation = {}
def getActivation(name):
  # the hook signature
  def hook(model, input, output):
    activation[name] = output.detach()
  return hook

h1 = cnn.conv1.register_forward_hook(getActivation('conv1'))
h2 = cnn.conv2.register_forward_hook(getActivation('conv2'))
h3 = cnn.out.register_forward_hook(getActivation('out'))

from torch.autograd import Variable
import numpy as np

images, labels = next(iter(loaders['train']))
num_epochs = 1

chromosome = []
fitness = 0

def single_run(images, labels, num_epochs, chromosome, fitness):

  def train(num_epochs, cnn, loaders):
    cnn.train()
  for i, (images, labels) in enumerate(loaders['train']): 

      b_x = Variable(images)   # batch x
      b_y = Variable(labels)   # batch y

      output = cnn(b_x)[0] 

      out_ten = activation['out']

      layer_weights_1 = activation['conv1'].numpy().flatten()
      layer_weights_2 = activation['conv1'].numpy().flatten()

      chromosome = np.append(layer_weights_1,layer_weights_2)

      #fitness
      for j in range(len(out_ten)): 
        if out_ten[j].argmax().item() == labels[j].item(): 
          fitness += 1

  return fitness, chromosome

fitness, chromosome = single_run(images, labels,num_epochs, chromosome, fitness)

print(fitness)
print('______')
print(chromosome)

Hi maria.solyanik. I am also learning PyGAD. I recommend an example in PyGAD. I hope that helps! I get to know how to use PyGAD for deep learning with the code.
example code