How do i save images for each class using Pytorch? (no grid)

Question

I am using acgan to do image augmentation. At present, sample images are generated in grid format. But I want to save images for each class separately. (eg 1.png; 2.png ...) How should I modify this code? Or is there an answer that I would like to refer to?

class Generator(nn.Module):
def __init__(self):
    super(Generator, self).__init__()

    self.label_emb = nn.Embedding(opt.n_classes, opt.latent_dim)

    self.init_size = opt.img_size // 4  # Initial size before upsampling
    self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))

    self.conv_blocks = nn.Sequential(
        nn.BatchNorm2d(128),
        nn.Upsample(scale_factor=2),
        nn.Conv2d(128, 128, 3, stride=1, padding=1),
        nn.BatchNorm2d(128, 0.8),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Upsample(scale_factor=2),
        nn.Conv2d(128, 64, 3, stride=1, padding=1),
        nn.BatchNorm2d(64, 0.8),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),
        nn.Tanh(),
    )

..

generator = Generator()

..

def sample_image(n_row, batches_done):
   """Saves a grid of generated digits ranging from 0 to n_classes"""
   # Sample noise
   z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))
   # Get labels ranging from 0 to n_classes for n rows
   labels = np.array([num for _ in range(n_row) for num in range(n_row)])
   labels = Variable(LongTensor(labels))
   gen_imgs = generator(z, labels)
   save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=n_row, normalize=True)

Answer 1

in the def sample_image , you have line that defines target labels for the generator: labels = np.array([num for _ in range(n_row) for num in range(n_row)]) .

Instead of using num which changes due to being sampled from range, use constant number you pass as an argument ( class_id below):

def sample_image(n_row, batches_done, class_id):
   """Saves a grid of generated digits ranging from 0 to n_classes"""
   # Sample noise
   z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))
   # Get labels ranging from 0 to n_classes for n rows
   labels = np.array([class_id for _ in range(n_row) for __ in range(n_row)])
   labels = Variable(LongTensor(labels))
   gen_imgs = generator(z, labels)
   save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=n_row, normalize=True)

This way you will get a rectangular array full of images of the class you have requested.

Furthermore, to have just one image you can set n_row to 1. Note that you didn't provide code for save_image function, there may be some tricks to it.