Hyperspectral Classification Based on Lightweight 3-D-CNN With Transfer Learning

In plain words

A hyperspectral camera records many narrow colour bands for every pixel, not just red, green and blue. This paper classifies what covers each pixel of such an image. It looks at a small square patch of the image around the pixel, with all of its bands, and feeds that block into a small 3-D neural network. The network is deep, but it is built from cheap blocks that use far fewer weights than the usual 3-D convolution designs.

Labelled hyperspectral data is expensive, so the authors also reuse a network trained elsewhere. They try two sources: another hyperspectral scene from a different sensor, and ordinary colour photos that have been stretched into the same 3-D shape. The model is first trained on the source, then its weights are copied over and tuned on the target scene.

TL;DR

• A 19-layer 3-D CNN (3-D-LWNet) classifies hyperspectral pixels from 27 by 27 pixel cubes spanning all bands §IV-B §IV-C.
• Its residual units use a pointwise, depthwise, then pointwise convolution, which keeps the parameter count low §III-B fig. 2.
• Transfer from another HSI scene or from inflated RGB images is meant to help when labels are scarce §III-C §IV-D §IV-F.
• Author-reported accuracy is competitive with a 3-D CNN baseline on Indian Pines and KSC and slightly below it on Pavia University tab. VI tab. VII tab. VIII.
• The shipped code does not contain the weight-transfer step, and several details differ from the paper (see Pitfalls).

Problem

Hyperspectral classification with deep networks overfits when there are few labelled samples and the network has many parameters §I. Deeper 2-D CNNs help natural images, but their structure does not carry over directly to 3-D cubes, and the existing 3-D CNN models for hyperspectral data are shallow (three convolution layers) §I. The paper states that its parameter-saving design is aimed at this small-sample setting §I §IV-C.

Method

Framework. A pixel's label is the label of the centre pixel of a 27 by 27 by L cube cut from the image §III-A §IV-B. Each cube goes through three parts: an input convolution with max pooling, six lightweight (LW) units in four groups, and an adaptive average pooling followed by a fully connected layer that produces class scores §III-A fig. 1.

1. Input. 32 kernels of size 8 x 3 x 3 (spectral x height x width), then batch normalisation, ReLU and a max-pooling layer with stride 2 §III-A fig. 1.
2. LW units. Units are arranged in groups of 32, 64, 128 and 256 output channels, with 1, 2, 2 and 1 units per group §III-A fig. 1. In the first unit of each of the last three groups, stride 2 is applied in the depthwise convolution, which halves the spatial and spectral size. The shortcut at a change of width or size is an average pool followed by a pointwise convolution and batch norm §III-A §III-B fig. 2.
3. Head. Adaptive average pooling reduces the final volume to one value per channel, giving 256 values regardless of input size. A fully connected layer then gives C class scores, passed through log-softmax §III-A eq. 1.

Lightweight unit. Each unit has four stages. A pointwise convolution widens the channels to t times the input width, with t = 4 chosen by the authors from the values 1/4 to 6 tested §III-B. A depthwise 3 x 3 x 3 convolution is applied to that widened tensor, followed by a second pointwise convolution back to the unit width. The shortcut is added before a final ReLU §III-B fig. 2.

Counting parameters, the authors give 2,304 for the input layer and 11,648, 71,168, 257,024 and 420,864 for the four groups, which sums to 763,008 in the convolution layers. They compare this with about 44.9 million for their 3-D-CNN-LR baseline on Indian Pines §IV-C.

Training. Loss is the negative log-likelihood of the correct class (computed from the log-softmax outputs) plus an L2 penalty on all weights with coefficient lambda = 1e-5 eq. 2 §III-A. Optimiser is SGD with momentum §III-A §IV-B.

Transfer. Two strategies, both with fine-tuning §III-C fig. 3:

1. Cross-sensor: pretrain on a source HSI set with more labelled samples, then transfer to the target HSI set. No restriction is placed on the sensor. Adaptive average pooling sets the feature size to a fixed length, so different band counts can share the head §III-C fig. 3.
2. Cross-modal: pretrain on an RGB image set (CIFAR-10 and CIFAR-100 are used). Each RGB image is repeated l = 12 times along the band axis to make a 3l-band cube, which is a 36-band input §III-C §IV-A.

In fine-tuning, the whole network except the fully connected layer is copied from the pretrained model. The fully connected layer is initialised at random, and all layers are trained with the same learning rate §III-C.

# 3-D-LWNet forward pass  [§III-A] [eq. 1]
# Input:  cubes (B, 1, L, 27, 27), L = number of bands
# Output: log-scores (B, C)

def forward(cubes):
    # Input convolution: 32 kernels of 8 x 3 x 3  [§III-A]
    # [code: training/models/models_lw_3D.py#L165-L169]
    features = conv3d(cubes, out_channels=32, kernel=(8, 3, 3))
    features = relu(batch_norm(features))
    # Max pooling with stride 2; kernel size is in Pitfalls  [§III-A]
    features = max_pool3d(features, stride=2)
    # Six LW units in four groups: 1, 2, 2, 1 units  [§III-A] [fig. 1]
    for width, count in [(32, 1), (64, 2), (128, 2), (256, 1)]:
        for index in range(count):
            # First unit of groups 2 to 4 uses stride 2  [§III-A]
            stride = 2 if index == 0 and width > 32 else 1
            features = lw_unit(features, width, stride)
    # Adaptive average pooling to one value per channel  [§III-A]
    pooled = adaptive_avg_pool3d(features)   # (B, 256)
    # Fully connected layer, then log-softmax  [§III-A] [eq. 1]
    return log_softmax(linear(pooled))       # (B, C)

# Lightweight (LW) unit  [§III-B] [fig. 2]
# Input: features (B, width_in, D, H, W); output (B, width, D', H', W')

def lw_unit(features, width, stride):
    # Shortcut: identity, or average pool + pointwise conv + BN when
    # the stride or the width changes  [§III-B] [fig. 2a]
    shortcut = features
    if stride != 1 or features.channels != width:
        shortcut = avg_pool(features, stride)
        shortcut = batch_norm(conv1x1(shortcut, width))
    # Pointwise conv widens the channels by t = 4  [§III-B] [§III-A]
    # Code uses 4 x width, not 4 x input width  [see Paper vs code]
    wide = relu(batch_norm(conv1x1(features, 4 * width)))
    # Depthwise 3 x 3 x 3 conv; stride 2 in groups 2 to 4 first unit
    # [§III-B] [§III-A]. Code uses a grouped conv, see Paper vs code.
    wide = relu(batch_norm(depthwise3d(wide, stride)))
    # Pointwise conv back to the unit width, then BN, no ReLU  [§III-B]
    out = batch_norm(conv1x1(wide, width))
    # Add the shortcut, then ReLU  [§III-B] [fig. 2]
    return relu(out + shortcut)

# Transfer then fine-tune  [§III-C] [fig. 3]
# The pretrained step is described in the paper; the copy step is not in
# the shipped code (see Paper vs code).

def transfer_and_finetune(pretrained, num_classes, lr):
    # Build a new network for the target scene  [§III-C]
    target = build_3d_lwnet(num_classes)
    # Copy all weights except the final fully connected layer  [§III-C]
    copy_weights(target, pretrained, exclude="fc")
    # Train all layers with one learning rate  [§III-C]
    return sgd_train(target, learning_rate=lr)

Hyperparameters

Results

Author-reported. Setup: three public HSI scenes captured by three sensors (Pavia University, Indian Pines, KSC) as targets, with Pavia Center and Salinas as source scenes; CIFAR-10 and CIFAR-100 as RGB sources §IV-A. The baselines are 2-D CNN-LR, 3-D CNN-LR (a 3-D CNN with logistic regression), and 3-D ResNet (ResNet-20, as the text says it was selected) §IV-B §IV-C. Overall accuracy (OA) is shown as mean ± spread as reported in the tables; the number of runs is not stated. Training-sample counts differ between columns, as listed.

The 2-D CNN-LR baseline scores 94.04 on Pavia University, 89.99 on Indian Pines and 94.11 on KSC tab. VI tab. VII tab. VIII. The text reports that 3-D-LWNet is 0.14 points below 3-D CNN-LR on Pavia University and 1.31 and 1.91 points above it on Indian Pines and KSC §IV-C.

Transfer, with 25 labelled samples per class, is reported as follows.

Gains from transfer shrink as the number of training samples grows. At 50 samples per class, the authors report that gains from RGB pretraining fall to 1.71 points for CIFAR-10 and 2.75 for CIFAR-100 §IV-F. Pretraining on a four-class Salinas subset reduced accuracy on Pavia University at 25 samples per class: 88.37 falls to 87.78 with the four-class subset §IV-E tab. XI.

Other reported results: the 3-D-ResNet candidates in tab. V were cross-validated, and ResNet-20 was chosen as the baseline §IV-B. Table V shows overall accuracies of 98.40 (Pavia), 98.68 (Indian Pines) and 95.35 (KSC) for ResNet-20 tab. V.

Paper vs code

Official code: hkzhang-git/LWNet, commit ca8361cd4a1c6418e3366fcc2ebc91fb72ac8c78. The arXiv comments link github.com/hkzhang91/LWNet, which GitHub redirects to this repository, and the repository's description names this paper. The repository has no LICENSE file, so it is cite-only. Its training data paths point to a ../../../data/HSI_datasets/ directory that is not in the repository [data_preprocess/data_list/PaviaU_train.txt#L1].

• Unit count. The text says "eight LW units" §III-A. The groups, the figure and the 19-layer count all give six units (1 + 2 + 2 + 1) fig. 1 §IV-C. The code has six models_lw_3D.py.
• Expansion base. The paper says the first pointwise layer widens to t times the input channels §III-B. The code widens to 4 times the unit's output width, planes * 4 models_lw_3D.py. The two agree for units whose input and output widths are equal, but not for the first unit of groups 2 to 4, where the code's widened width is 8 times the input width (for example 256 from 32 in group 2).
• Depthwise convolution. The paper says the depthwise layer has the same channel count as the first pointwise layer §III-B. The code sets groups=planes on planes * 4 channels, which gives groups of 4 channels, not depthwise models_lw_3D.py.
• Max-pool kernel. Text says 3 x 3 x 3 §III-A. Code uses kernel_size=(2, 3, 3) models_lw_3D.py.
• Max-pool rounding. The paper defines its bracketed output-size operation as a ceiling §III-A. PyTorch's MaxPool3d defaults to floor (ceil_mode=False), and the code does not set it models_lw_3D.py.
• Loss. Eq. 2 sums the loss over N samples eq. 2. The code uses PyTorch's default mean reduction functions_for_training.py. The L2 term is implemented as SGD weight_decay=1e-5 on all parameters main.py, which corresponds to lambda times the parameters in the gradient. Eq. 2's constant convention is not stated, so the two may differ by a factor of 2.
• Transfer is not wired up. The paper's transfer procedure copies all layers except the fully connected layer §III-C. The shipped fine-tune script trains a fresh 3-D-LWNet, and its restore call is commented out LWNet_Salinas-4.py. The only pretraining script is for CIFAR LWNet_pretrain_main_cifar.py. Its model_restore loads the full state dict, including the 10-class head, so it cannot load into a 13-class target as-is functions_for_pretrain.py. No code for the cross-sensor pretraining (Pavia Center or Salinas) is shipped.
• Fine-tuning script. LWNet_Salinas-4.py defaults to the KSC data, 13 classes, 10 epochs, batch 8 and CPU LWNet_Salinas-4.py LWNet_Salinas-4.py. The paper's batch size is 20 §IV-B.
• Baseline model in the main script. main.py defaults to resnet26 main.py, while the text chooses ResNet-20 as the baseline §IV-B.
• Learning-rate schedule. The paper says the rate is divided by 10 at epoch 50 §IV-B. The code keeps 0.01 through epoch 50 and uses 0.001 from epoch 51 functions_for_training.py. The CIFAR pretraining schedule drops at 0.8 of the epochs instead of 0.84 functions_for_pretrain.py.
• Evaluation set. main.py reads <dataset>_test_part.txt main.py. That file is not shipped. The generator takes kk[n:2n] per class, a slice about the size of the training set, not the full held-out remainder functions_for_samples_extraction.py functions_for_samples_extraction.py.
• Default test set for OA/AA/K. OA_AA_K_calculation_main.py defaults --test_data to Indian_train.txt, so running it unchanged reports accuracy on training samples OA_AA_K_calculation_main.py.
• Augmentation. HSI training uses no augmentation main.py. The CIFAR pretraining does use random flips and 90-degree rotations LWNet_pretrain_main_cifar.py functions_for_pretrain.py.
• Normalisation and padding. The HSI pipeline min-max normalises the whole image and zero-pads the borders functions_for_samples_extraction.py functions_for_samples_extraction.py. The paper does not state either step.
• RGB inflation order. The code repeats each colour channel 12 times in a row, giving R R ... G G ... B B function_for_cifar_preprocess.py. The paper says each RGB image is repeated along the third axis, without fixing the order §III-C.
• Split sizes. The split files and Table I give 160 Pavia training pixels per class [data_preprocess/data_list/PaviaU_split.txt]. The shipped PaviaU_train.txt has 120 lines for class 1 and 120 for class 9 (1080 total) [data_preprocess/data_list/PaviaU_train.txt]. The test totals in Tables II and III (6,223 and 4,297) do not match the shipped held-out lists (8,484 and 4,752). The paper does not say which pixels were excluded.

The ResNet training script matches the paper's optimiser settings: SGD, momentum 0.9, initial rate 0.01, batch 20, weight decay 1e-5 and 60 epochs §IV-B main.py main.py. The LW unit group layout and stride placement match §III-A models_lw_3D.py.

Pitfalls

• Unit count is inconsistent in the text. "Eight LW units" appears once §III-A, but the groups, figure and parameter count describe six fig. 1 §IV-C. Implement six.
• Expansion width is ambiguous. Read "t times the input channels" §III-B as the paper says. The code instead uses t times the unit's output width (see Paper vs code). Pick one and check the parameter count: 763,008 conv parameters §IV-C is the paper's figure.
• Max-pool details are contradictory. The text gives a 3 x 3 x 3 kernel, but the code uses 2 x 3 x 3. The text also defines the rounding as a ceiling, while the code floors. Spatial sizes after the first pool and after the stride-2 units differ between these readings.
• Transfer is not reproducible from the repository. The copy step, and the cross-sensor pretraining, are missing. The paper does not state the fine-tuning epoch count or the pretraining epoch count for HSI sources.
• Dataset files are external. The split lists point to a path that is not in the repository, so the .npy cubes must be regenerated from the raw HSI data with preprocess.py [data_preprocess/data_list/PaviaU_train.txt#L1].
• Evaluation scripts can report the wrong numbers. main.py prints accuracy on the _test_part slice, and the OA/AA/K script defaults to training pixels (see Paper vs code). Compare against the paper's test counts only after fixing both.
• Sample counts disagree. Pavia University training counts differ between the split file, the shipped list and Table I (see Paper vs code). Confirm the split before comparing numbers.
• Few runs, wide spreads. Several transfer cells have spreads larger than the gap between methods, for example 92.54 ± 11.38 for Salinas-16 at 25 samples per class tab. IX tab. XI. The number of runs is not stated.
• Numerical sensitivity (inferred from code). Min-max normalisation uses the image-wide maximum and minimum, so one outlier pixel rescales the whole image functions_for_samples_extraction.py.
• Batch size and memory. The authors state that a large batch takes more GPU memory, and they chose 20 as sufficient for at most 16 classes §IV-B.

Lineage

• Builds on: ResNet residual units §II-B §III-B, depthwise and pointwise convolutions from MobileNet §II-B, and 3-D CNN models for HSI, including the 3-D CNN-LR baseline §I §IV-C.
• Compared against: 2-D CNN-LR and 3-D CNN-LR §IV-B §IV-C, the two-CNN method of Yang et al. [46] §IV-D tab. IX, and the 3-D ResNet family in this paper tab. V.
• Superseded by: none recorded in this wiki.