Generating 3D structures from a 2D slice with GAN-based dimensionality expansion

In plain words

Engineers who design battery electrodes or other porous materials want to know what the inside looks like in three dimensions. Getting real 3D scans is slow and expensive, but a flat cross-section photo is cheap. This paper asks whether a neural network can learn the full 3D structure from that one flat picture.

The trick is to let a 3D generator make a block of material, then cut thin slices through it in every direction. A 2D critic then judges those slices against real pieces of the photo. The generator only ever sees 2D evidence, yet has to produce a volume whose slices all look right. The authors also looked closely at why generated images go wrong near their edges, and set rules for the layer sizes so that every voxel receives the same amount of learning. After training, the network can produce a large new 3D sample in seconds.

TL;DR

SliceGAN is a Wasserstein GAN whose generator makes 3D volumes and whose discriminator is 2D and sees slices of those volumes in all three axes. A single 2D image of an isotropic material is enough to train it; anisotropic materials need one image per axis. Transposed convolutions must satisfy three rules, or edge artefacts appear. The results are author-reported and mostly shown as figures, with statistics on one battery cathode dataset.

Problem

Simulations of transport, deformation or electrochemistry need 3D microstructures, but 3D imaging is limited in resolution and field of view, while 2D micrographs are easier to get §1. Earlier 2D-to-3D methods are restricted to specific structure types: packed ellipsoids for powders and grain-growth tools for polycrystals §2.2. Correlation-based methods can be slow; the paper reports hours for a 10^6-voxel sample §2.2. A GAN trained on 3D data needs 3D training sets, which is the thing that is hard to get §1 [abstract].

Method

Training data. Real samples are 2D images of the material, one-hot encoded per phase (see Pre-processing). For an isotropic material the same image set is used for the x, y and z discriminator roles §3 run_slicegan.py model.py.

Generator. A 3D transposed-convolution network maps a latent tensor of spatial size 4 to a volume of 64 voxels per edge, with a softmax over phases in the final layer §3 §4 tab. 1 §5.1.

Slicing. For a generated cube of edge l, there are 3l slices, taken at every voxel depth along x, y and z §3.

Discriminator. A 2D convolutional critic, not a 3D one. It is trained on the fake slices and on real images, each fake slice being paired with a sampled real image §3 alg. 1.

Loss. A Wasserstein loss with a gradient penalty on random real-fake interpolations §3 alg. 1. The generator is trained on the negative critic score of its slices alg. 1.

Batch and update schedule. The critic is updated n_D times per generator update. The generator batch m_G is twice the critic batch m_D §3 alg. 1.

Anisotropic materials. One critic per axis, each trained on its own real image set. The paper's worked example uses two perpendicular images (circular and rectangular features) §3. The run script takes one path for isotropic data and three for anisotropic data run_slicegan.py. The adjusted algorithm is in the supplement [SI §A].

# SliceGAN training, isotropic case  [alg. 1]
# Input:  real images, l, n_D, m_D, m_G, lam (penalty), Adam settings
# Output: generator G

def train_slicegan(G, D, real_images, l, n_D, m_D, m_G, lam):
    # repeat until the generator has converged  [alg. 1]
    while not converged(G):
        # critic phase: n_D updates per generator update  [alg. 1]
        for t in range(n_D):
            loss_D = 0
            for i in range(m_D):
                # fresh latent vector and a 3D volume  [alg. 1] [tab. 1]
                z = sample_normal(64, 4, 4, 4)       # (C_z, 4, 4, 4)
                volume = G(z)                        # (3, l, l, l)
                for axis in range(3):
                    for depth in range(l):
                        # 2D slice at this depth along this axis  [§3]
                        fake = take_slice(volume, axis, depth)
                        # one real image, sampled for this slice  [alg. 1]
                        real = sample_real_image(real_images, l)
                        # random point between real and fake  [alg. 1]
                        eps = uniform(0, 1)
                        mix = eps * fake + (1 - eps) * real
                        # Wasserstein term plus penalty on the gradient
                        # norm at mix  [alg. 1]
                        loss_D += (D(fake) - D(real)
                                   + lam * (grad_norm(D, mix) - 1) ** 2)
            # Adam step on the averaged critic loss  [alg. 1]
            D.params = adam_step(D.params, grad(loss_D / m_D))
        # generator phase with m_G samples  [alg. 1] [§3]
        loss_G = 0
        for j in range(m_G):
            volume = G(sample_normal(64, 4, 4, 4))
            for axis in range(3):
                for depth in range(l):
                    fake = take_slice(volume, axis, depth)
                    # the generator raises the critic score of its slices
                    loss_G += -D(fake)
        G.params = adam_step(G.params, grad(loss_G / m_G))
    return G

Uniform information density. A transposed convolution with kernel k, stride s and padding p gives a uniform density only when three rules hold: s < k; k mod s = 0; and p ≥ k − s §4. Prime k are discarded because they force s = 1, and k > 6 is treated as impractical because parameters grow with k³ §4. The three remaining practical settings are (k, s, p) = (4, 2, 2), (6, 3, 3) and (6, 2, 4); the (4, 2, 2) setting is used for most layers §4. Resize-convolution is named as the alternative, which the paper says needs more memory for 3D volumes §4 [SI §C].

Latent spatial size. The first generator layer is fed a latent tensor of spatial size 4 rather than 1, so that kernel overlap is learned at training time. The paper says this also makes every 32nd plane come from the same combination of kernel elements, so the generator cannot learn anisotropy within a plane §4 [SI §E].

Pre-processing. Segmented n-phase data are one-hot encoded with one channel per phase. Grey or colour images can be used without this step, but the paper reports the approach works best with one-hot data §5.1 [SI §F].

Generation. The volume edge length is chosen at inference time to show the features §5.2. The paper does not state the latent size used for the 10^8-voxel claim.

Hyperparameters

Values not in the paper are marked "not stated". Code values are from the official repository at code_commit.

Results

Author-reported. Setup: seven training sets (synthetic grains, ceramic, carbon fibre, battery separator, steel, synthetic interface, NMC cathode), each with a 2D-to-3D generator tab. 2 §5.2. The text gives no tabulated numbers for these examples; they are visual, shown in fig. 3.

Two-point correlation functions and triple-phase-boundary densities for the NMC case are in the supplement [SI fig. S6]. The metric tool is TauFactor, not in the repository.

Paper vs code

The official repository is MIT-licensed by Steve Kench LICENSE. It is linked from the paper's Code Availability section.

• Batch sizes. Paper: m_G = 2 m_D §3. Code uses 8 for both model.py.
• Generator architecture. Paper: transposed convolutions in every layer, with the layer table given §4 tab. 1. Code default run_slicegan.py builds slicegan_rc_nets: four transposed convolutions, then trilinear upsampling and a 3×3×3 convolution. The fifth transposed-convolution layer exists but is never called networks.py run_slicegan.py. A pure transposed-convolution generator is also in the repository but not used by the run script networks.py.
• Channel counts. Paper: latent 64 and generator filters 512, 256, 128, 64, 3 tab. 1. Code: latent 32 and filters 1024, 512, 128, 32, 3 run_slicegan.py, run_slicegan.py.
• Isotropic critic updates. Algorithm 1 sums the loss over axes and makes one update of the single critic per iteration alg. 1. Code takes three separate Adam steps on that one critic per iteration, one for each axis model.py.
• Real images paired with fake slices. Algorithm 1 samples one real image per fake slice alg. 1. Code takes one real batch of m_D images and averages the critic over all fake slices of the batch, so real and fake counts differ [code: model.py#L103-L108].
• Gradient-penalty inputs. The penalty is computed on the first eight fake slices and the real batch model.py. Because the reshape orders slices by sample, these eight slices all come from the first generated volume, not from eight different volumes model.py.
• Loss scale. Algorithm 1 sums over slices and divides by the batch size. Code averages over slices and axes model.py model.py. Adam is largely insensitive to gradient scale, so this is probably minor; the effect was not tested.
• Real data for an isotropic material. The paper uses a single 2D image §3. The run script uses a 3D volume (NMC) and samples random 2D sections from it run_slicegan.py preprocessing.py.
• Generation. The run script calls test_img with latent size 8 and wraps the latent on two axes to make the output periodic, then trims one voxel from those axes run_slicegan.py util.py. The paper does not describe this. The code requires CUDA at this step util.py.
• Supplement numbering. The main text cites "Supplementary Information S1" to "S7", but the supplement's own figures are labelled S1 to S6 and the section letters differ. For example, the main text cites S2 for the k = 4, s = 3 example, which the supplement labels Figure S1 [SI §B].

No discrepancies were found for the gradient-penalty formula, the Wasserstein term, the slice count (3l per sample) or the softmax output for n-phase data.

Pitfalls

• Unstated values. The text leaves the learning rate, Adam betas, λ, n_D, m_D, training length and initialisation open. The repository's values are the only concrete ones, and they differ from the paper in several places (see above).
• Two generators. An implementer who follows Table 1 gets a different network from the one the run script trains. Pick one and say which.
• Isotropic update rule. Three Adam steps per iteration on one critic is not the update in Algorithm 1. It may change stability; the paper does not discuss it.
• Sampling the last section. Real sections are drawn from randint(n − 1), so the last section along each axis is never drawn preprocessing.py.
• Edge artefacts are sensitive to layer settings. The information-density rules are a constraint on (k, s, p), not a tuning choice. Breaking s < k, k mod s = 0, or p ≥ k − s produces the checkerboard or edge gradients the paper describes §4 [SI §B] [SI §C].
• Volume size and the 10^8 claim. The latent size used for the 10^8-voxel claim is not stated. By our reading of the code, the default test call (latent size 8) gives about 192 × 191 × 191 voxels, roughly 7 × 10^6, not 10^8. This was derived by hand from networks.py and not run.
• Training cost. c. 4 hours on one Titan Xp §2.2. The default code runs 100 epochs; each epoch has 3,600 iterations of batches of 8 per axis (derived from model.py#L36 and preprocessing.py#L54). The paper does not state its epoch count, so the two cannot be compared; the run was not timed.
• Anisotropic is only described. Anisotropic training appears in the supplement [SI §A], and the paper's worked example for it is a fibre composite, not the default run.
• Discriminator gradient penalty cost. Each step computes a second-order gradient on eight interpolates per axis, and the critic runs on every slice (512 per axis per step in the default code). Memory and time scale with l.
• Known limit. The paper says the method cannot capture anisotropy inside one plane, and that continuous label spaces (colour) lose quality at the 12 GB memory budget used §4 §5.2.

Lineage

• Builds on: the Wasserstein GAN loss, cited as Arjovsky et al. §3 §2.1. The gradient-penalty term is not attributed in the text.
• Builds on: Adam, cited for the optimiser in Algorithm 1 alg. 1.
• Compared against: a 3D-to-3D GAN by Gayon-Lombardo et al. §2.1 §5.2; correlation-based reconstruction by Zhang et al. §2.2.
• Superseded by: not stated.