How to train a proper StyleGan Model

Beginners
1

Hi HF community. I’m developing my first StyleGan model with a small dataset consisting of 200 Chest-X ray pneumonia images. I am not familiar with the implementation.

import torch
import torchvision
import os
from stylegan2_pytorch import Trainer

def save_sample_images(generator, step, latent_dim, device, num_images):
    # Create the styles vector (latent vector)
    styles = torch.randn(num_images, latent_dim).to(device)


    with torch.no_grad():
        # Generate images
        images = generator(styles)

    os.makedirs('/content/stylegan_pneumonia_generated_images', exist_ok=True)
    for i, image in enumerate(images):
        torchvision.utils.save_image(image, f'/content/stylegan_pneumonia_generated_images/generated_image_{step}_{i}.png')

# Assume the image size and latent dimension size
image_size = 256 # The image size used in your generator
latent_dim = 512 # The latent vector size used in your generator

# Set up the configuration for training
trainer = Trainer(
    name='pneumonia_stylegan',
    results_dir='results',
    models_dir='models',
    base_dir='.',  # Adjust as needed
    image_size=256,  # Match this with your dataset
    network_capacity=16,
    transparent=False,  # Set to True if your images are RGBA (transparent)
    batch_size=4,
    lr=2e-4,
    lr_mlp=0.1,
    ttur_mult=1.5,
    num_workers=2,
    save_every=1000,  # How often to save models and images
    evaluate_every=1000,  # How often to evaluate and save generated images
    num_image_tiles=8,
    trunc_psi=0.75,
    fp16=False,
    cl_reg=False
)


# Before entering the training loop, check if the GAN and its generator are initialized
if trainer.GAN is None or not hasattr(trainer.GAN, 'GE'):
    raise AttributeError("The GAN model or the GE attribute has not been initialized.")


# Modify your training loop to include the image size and save images periodically
for step in range(10000):  # Number of training steps
    # ... your training loop ...

    if step % 1000 == 0:
        save_sample_images(trainer.GAN.GE, step, latent_dim, device, trainer.batch_size)


---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-13-f1d977416ec3> in <cell line: 47>()
     46 # Before entering the training loop, check if the GAN and its generator are initialized
     47 if trainer.GAN is None or not hasattr(trainer.GAN, 'GE'):
---> 48     raise AttributeError("The GAN model or the GE attribute has not been initialized.")
     49 
     50 

AttributeError: The GAN model or the GE attribute has not been initialized.

I attempted to put in place a code that would display a progress bar to indicate whether the model was operating as well as to implement another code to save the generated images in another directory, but it kept returning errors. How can I solve this?

Also, there is an output error that reads, "ERROR: Could not open requirements file: [Errno 2] No such file or directory: ‘requirements.txt’ " when these few lines are run:

!git clone https://github.com/NVlabs/stylegan2.git
%cd stylegan2
!ls -l
!pip install -r requirements.txt
type or paste code here

How can I solve this?