Training the networks

To train the 3D-GAN, perform the following steps:

1. Start by specifying the values for the different hyperparameters required for the training, shown as follows:

gen_learning_rate = 0.0025
dis_learning_rate = 0.00001
beta = 0.5
batch_size = 32
z_size = 200
DIR_PATH = 'Path to the 3DShapenets dataset directory'
generated_volumes_dir = 'generated_volumes'
log_dir = 'logs'

1. Next, create and compile both of the networks, shown as follows:

# Create instances
generator = build_generator()
discriminator = build_discriminator()

# Specify optimizer 
gen_optimizer = Adam(lr=gen_learning_rate, beta_1=beta)
dis_optimizer = Adam(lr=dis_learning_rate, beta_1=0.9)

# Compile networks
generator.compile(loss="binary_crossentropy", optimizer="adam")
discriminator.compile(loss='binary_crossentropy', optimizer=dis_optimizer)

We are using the Adam optimizer as the optimization algorithm and binary_crossentropy as the loss function. The hyperparameter values for the Adam optimizer are specified in the first step.

1. Then, create and compile the adversarial model:

discriminator.trainable = False
adversarial_model = Sequential()
adversarial_model.add(generator)
adversarial_model.add(discriminator)
adversarial_model.compile(loss="binary_crossentropy", optimizer=Adam(lr=gen_learning_rate, beta_1=beta))

1. After that, extract and load all the airplane images for training:

def getVoxelsFromMat(path, cube_len=64):
    voxels = io.loadmat(path)['instance']
    voxels = np.pad(voxels, (1, 1), 'constant', constant_values=(0, 0))
    if cube_len != 32 and cube_len == 64:
        voxels = nd.zoom(voxels, (2, 2, 2), mode='constant', order=0)
    return voxels

def get3ImagesForACategory(obj='airplane', train=True, cube_len=64, obj_ratio=1.0):
    obj_path = DIR_PATH + obj + '/30/'
    obj_path += 'train/' if train else 'test/'
    fileList = [f for f in os.listdir(obj_path) if f.endswith('.mat')]
    fileList = fileList[0:int(obj_ratio * len(fileList))]
    volumeBatch = np.asarray([getVoxelsFromMat(obj_path + f, cube_len) for f in fileList], dtype=np.bool)
    return volumeBatch

volumes = get3ImagesForACategory(obj='airplane', train=True, obj_ratio=1.0)
volumes = volumes[..., np.newaxis].astype(np.float)

1. Next, add a TensorBoard callback and add the generator and discriminator networks:

tensorboard = TensorBoard(log_dir="{}/{}".format(log_dir, time.time()))
tensorboard.set_model(generator)
tensorboard.set_model(discriminator)

1. Add a loop, which will run for a specified number of epochs:

for epoch in range(epochs):
    print("Epoch:", epoch)

    # Create two lists to store losses
    gen_losses = []
    dis_losses = []

1. Add another loop inside the first loop to run for a specified number of batches:

    number_of_batches = int(volumes.shape[0] / batch_size)
    print("Number of batches:", number_of_batches)
    for index in range(number_of_batches):
        print("Batch:", index + 1)

1. Next, sample a batch of images from the set of real images and a batch of noise vectors from the Gaussian Normal distribution. The shape of a noise vector should be (1, 1, 1, 200):

z_sample = np.random.normal(0, 0.33, size=[batch_size, 1, 1, 1, 
                            z_size]).astype(np.float32)
volumes_batch = volumes[index * batch_size:(index + 1) * batch_size, 
                        :, :, :]

1. Generate fake images using the generator network. Pass it a batch of noise vectors from z_sample and generate a batch of fake images:

gen_volumes = generator.predict(z_sample,verbose=3)

1. Next, train the discriminator network on the fake images generated by the generator and a batch of real images from the set of real images. Also, make the discriminator trainable:

# Make the discriminator network trainable
discriminator.trainable = True
           
# Create fake and real labels
labels_real = np.reshape([1] * batch_size, (-1, 1, 1, 1, 1))
labels_fake = np.reshape([0] * batch_size, (-1, 1, 1, 1, 1))
        
# Train the discriminator network
loss_real = discriminator.train_on_batch(volumes_batch, 
                                         labels_real)
loss_fake = discriminator.train_on_batch(gen_volumes, 
                                         labels_fake)
        
# Calculate total discriminator loss
d_loss = 0.5 * (loss_real + loss_fake)

The preceding piece of code trains the discriminator network and calculates the total discriminator loss.

1. Train the adversarial model containing both the generator and the discriminator network:

z = np.random.normal(0, 0.33, size=[batch_size, 1, 1, 1, z_size]).astype(np.float32)

        # Train the adversarial model
        g_loss = adversarial_model.train_on_batch(z, np.reshape([1] * batch_size, (-1, 1, 1, 1, 1)))
        

Also, add losses to their respective lists as follows:

        gen_losses.append(g_loss)
        dis_losses.append(d_loss)

1. Generate and save the 3D images after every second epoch:

        if index % 10 == 0:
            z_sample2 = np.random.normal(0, 0.33, size=[batch_size, 1, 1, 1, z_size]).astype(np.float32)
            generated_volumes = generator.predict(z_sample2, verbose=3)
            for i, generated_volume in enumerate(generated_volumes[:5]):
                voxels = np.squeeze(generated_volume)
                voxels[voxels < 0.5] = 0.
                voxels[voxels >= 0.5] = 1.
                saveFromVoxels(voxels, "results/img_{}_{}_{}".format(epoch, index, i))

1. After each epoch, save the average losses to tensorboard:

# Save losses to Tensorboard
write_log(tensorboard, 'g_loss', np.mean(gen_losses), epoch)
write_log(tensorboard, 'd_loss', np.mean(dis_losses), epoch)