Skip to content

CPU full usage available if GPU not available #70

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jeremyfix merged 1 commit into from
Jan 14, 2026
Merged

CPU full usage available if GPU not available #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions LabsSolutions/pytorch/gan/config-sample.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ data:
rootdir: "/opt/datasets/"
debug: False
batch_size: 256
nthreads: 32 # You may need to adapt this in case no GPU is available

optim:
base_lr: 0.0002
Expand Down
37 changes: 31 additions & 6 deletions LabsSolutions/pytorch/gan/src/ganlab/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,20 @@ def train(config):

use_cuda = torch.cuda.is_available()
device = torch.device("cuda") if use_cuda else torch.device("cpu")

# Configure CPU parallelism settings
if not use_cuda:
# Get the number of CPU threads on the system
cpu_threads = os.cpu_count()
logger.info(f"CPU mode: counted {cpu_threads} threads")
logger.info(f"CPU mode: using {nthreads} workers for data loading")

# Set the number of threads for inter-op parallelism (parallel ops)
torch.set_num_threads(nthreads)

# Set the number of threads for intra-op parallelism (within ops like matrix mult)
torch.set_num_interop_threads(nthreads)


# Dataloaders
train_loader, valid_loader, img_shape = data.get_dataloaders(
Expand All @@ -121,10 +135,21 @@ def train(config):
generator_base_c,
)
model.to(device)

if not use_cuda:
# For CPU training with multiple cores
model = torch.nn.DataParallel(model)
logger.info(f"DataParallel enabled (processes threads internally)")

# Optimizers
critic = model.discriminator
generator = model.generator
# Handle both wrapped and unwrapped models
if isinstance(model, torch.nn.DataParallel):
critic = model.module.discriminator
generator = model.module.generator
else:
critic = model.discriminator
generator = model.generator

######################
# START CODING HERE ##
######################
Expand Down Expand Up @@ -190,7 +215,7 @@ def train(config):

# Generate few samples from the initial generator
model.eval()
fake_images = model.generator(X=fixed_noise)
fake_images = generator(X=fixed_noise)
fake_images = (fake_images * data._IMG_STD + data._IMG_MEAN).clamp(0, 1.0)
grid = torchvision.utils.make_grid(fake_images, nrow=sample_nrows, normalize=True)
tensorboard_writer.add_image("Generated", grid, 0)
Expand Down Expand Up @@ -402,7 +427,7 @@ def train(config):

# Generate few samples from the generator
model.eval()
fake_images = model.generator(X=fixed_noise)
fake_images = generator(X=fixed_noise)
# Unscale the images
fake_images = (fake_images * data._IMG_STD + data._IMG_MEAN).clamp(0, 1.0)
grid = torchvision.utils.make_grid(fake_images, nrow=sample_nrows)
Expand All @@ -418,14 +443,14 @@ def train(config):

# We save the generator
logger.info(f"Generator saved at {save_path}")
torch.save(model.generator, save_path)
torch.save(generator, save_path)

# Important: ensure the model is in eval mode before exporting !
# the graph in train/test mode is not the same
model.eval()
dummy_input = torch.zeros((1, latent_size), device=device)
torch.onnx.export(
model.generator,
generator,
dummy_input,
logdir + "/generator.onnx",
verbose=False,
Expand Down