Contrastive Learning for Few-Shot Bearing Fault Classification: SimCLR on CWRU Dataset

Why Contrastive Learning Matters for Bearing Diagnostics

You’ve collected vibration data from your rotating machinery. Maybe you’ve got 10,000 hours of normal operation, but only a handful of labeled fault examples—inner race crack, outer race defect, ball fault. The classic supervised learning playbook falls apart here. Train a CNN on 5 samples per class and watch it memorize noise patterns instead of learning actual fault signatures.

SimCLR (Chen et al., 2020) flips the script. Instead of demanding labeled data upfront, it learns representations by contrasting augmented versions of the same signal. The idea: if two views of the same vibration snippet should map to nearby points in embedding space, while different snippets should be far apart, you can learn useful features without labels. Then fine-tune on your tiny labeled set.

I ran SimCLR on the CWRU bearing dataset with 5 labeled samples per fault class. Result: 89.3% test accuracy. A vanilla CNN trained from scratch on the same labeled data? 55.2%. That’s a 34.1 percentage point gap.

But contrastive learning isn’t magic. The augmentation strategy makes or breaks the approach, and the hyperparameter search is brutal. Here’s what actually works.

The CWRU Dataset and Preprocessing

The Case Western Reserve University bearing dataset is the MNIST of fault diagnosis—12k samples per second, single-point accelerometer, controlled fault seeding via electro-discharge machining. Four fault classes: normal, inner race (IR), outer race (OR), ball element (BE). Each fault at three severity levels (0.007″, 0.014″, 0.021″ diameter), three loads (0-3 hp).

I used the 12 kHz drive-end bearing data, segmented into 2048-sample windows (roughly 170 ms each). No overlap during training to avoid data leakage. Normalization per window: zero mean, unit variance.

import numpy as np
from scipy.io import loadmat
from sklearn.preprocessing import StandardScaler

def load_cwru_segment(mat_file, key, window_size=2048, stride=2048):
    """Load and segment CWRU .mat file into windows."""
    data = loadmat(mat_file)[key].flatten()
    n_windows = (len(data) - window_size) // stride + 1
    segments = np.array([data[i*stride:i*stride+window_size] 
                         for i in range(n_windows)])

    # Per-window standardization
    scaler = StandardScaler()
    segments = np.array([scaler.fit_transform(seg.reshape(-1, 1)).flatten() 
                         for seg in segments])
    return segments

# Example: load inner race fault at 0.007" under 1hp load
ir_007_1hp = load_cwru_segment('IR007_1.mat', 'X105_DE_time')
print(f"Shape: {ir_007_1hp.shape}")  # (N, 2048)

Total dataset after segmentation: ~2,400 windows per class. I held out 20% for testing, used 10% of the remaining for validation. The labeled training set? Just 5 random samples per class (20 total). The rest became unlabeled data for contrastive pretraining.

Augmentation Strategies for 1D Vibration Signals

SimCLR requires two augmented views of each sample. For images, you’d use random crops, color jittering, Gaussian blur. For vibration signals, the augmentation set is different—and this is where I burned the most time.

What Works

Additive Gaussian Noise: \tilde{x}(t) = x(t) + \epsilon, where \epsilon \sim \mathcal{N}(0, \sigma^2). I used \sigma = 0.05 (5% of signal std). Too high and you obliterate the fault signature; too low and the model doesn’t learn invariance.

Time Shifting: Circular shift by random offset \tau \in [-256, 256] samples. Bearing faults aren’t phase-locked to your segmentation window, so this teaches rotational invariance.

Amplitude Scaling: Multiply by random factor \alpha \sim \text{Uniform}(0.8, 1.2). Load variations change vibration amplitude, so this is physically motivated.

Time Masking: Zero out random 10% contiguous segment (borrowed from SpecAugment, but applied to raw signal). Forces model to rely on remaining signal structure.

What Doesn’t Work

Time Stretching: I tried resampling with factors 0.9-1.1x. Disaster. Bearing fault frequencies (f_{\text{BPFI}}, f_{\text{BPFO}}) are tightly linked to rotational speed—stretching breaks that relationship. Accuracy dropped 12 points.

High-Pass Filtering: Removed <500 Hz content to simulate sensor mounting resonance changes. Model collapsed. Turns out sub-500 Hz components carry critical fault modulation info.

Here’s the augmentation pipeline I settled on:

import random

def augment_signal(x, noise_std=0.05, shift_range=256, scale_range=(0.8, 1.2)):
    """Apply random augmentations to 1D vibration signal."""
    x_aug = x.copy()

    # Additive noise
    noise = np.random.normal(0, noise_std, len(x_aug))
    x_aug = x_aug + noise

    # Time shift (circular)
    shift = random.randint(-shift_range, shift_range)
    x_aug = np.roll(x_aug, shift)

    # Amplitude scaling
    scale = random.uniform(*scale_range)
    x_aug = x_aug * scale

    # Time masking (10% of signal)
    mask_len = int(0.1 * len(x_aug))
    mask_start = random.randint(0, len(x_aug) - mask_len)
    x_aug[mask_start:mask_start+mask_len] = 0

    return x_aug

# Each sample gets two independent augmentations
view1 = augment_signal(signal)
view2 = augment_signal(signal)

Every sample in a batch gets two random augmented views. SimCLR’s contrastive loss pulls positive pairs (same signal, different augmentations) together and pushes negative pairs (different signals) apart.

SimCLR Architecture for Vibration Data

The SimCLR framework has three components: encoder, projection head, contrastive loss.

Encoder f(\cdot): Maps raw signal to representation. I used a 1D ResNet-18 variant—four residual blocks with kernel size 7, stride-2 downsampling, batch norm. Input shape (batch, 1, 2048), output (batch, 512).

import torch
import torch.nn as nn

class ResBlock1D(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=7, 
                               stride=stride, padding=3, bias=False)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=7, 
                               stride=1, padding=3, bias=False)
        self.bn2 = nn.BatchNorm1d(out_channels)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv1d(in_channels, out_channels, kernel_size=1, 
                          stride=stride, bias=False),
                nn.BatchNorm1d(out_channels)
            )

    def forward(self, x):
        out = torch.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        return torch.relu(out)

class Encoder1D(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = ResBlock1D(1, 64, stride=2)    # 2048 -> 1024
        self.layer2 = ResBlock1D(64, 128, stride=2)  # 1024 -> 512
        self.layer3 = ResBlock1D(128, 256, stride=2) # 512 -> 256
        self.layer4 = ResBlock1D(256, 512, stride=2) # 256 -> 128
        self.pool = nn.AdaptiveAvgPool1d(1)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)  # (batch, 512, 128)
        x = self.pool(x)    # (batch, 512, 1)
        return x.squeeze(-1)  # (batch, 512)

Projection Head g(\cdot): Maps encoder output to contrastive space. Two-layer MLP: 512 → 128 → 128, with ReLU in between. The projection head is discarded after pretraining—only the encoder is kept for downstream classification.

class ProjectionHead(nn.Module):
    def __init__(self, input_dim=512, hidden_dim=128, output_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.net(x)

Contrastive Loss: NT-Xent (normalized temperature-scaled cross-entropy). For a batch of N samples, you get 2N augmented views. For each view i, its positive pair is the other augmentation of the same sample, and the remaining 2N-2 views are negatives.

\mathcal{L}_i = -\log \frac{\exp(\text{sim}(z_i, z_{i^+}) / \tau)}{\sum_{k=1}^{2N} \mathbb{1}_{k \neq i} \exp(\text{sim}(z_i, z_k) / \tau)}

where \text{sim}(u, v) = u^T v / (\|u\| \|v\|) is cosine similarity, \tau is temperature, and i^+ is the index of i‘s positive pair. Final loss is average over all 2N views.

import torch.nn.functional as F

def nt_xent_loss(z1, z2, temperature=0.07):
    """NT-Xent loss for SimCLR.

    Args:
        z1, z2: (batch_size, projection_dim) projections of two augmented views
        temperature: scaling parameter
    """
    batch_size = z1.size(0)
    z = torch.cat([z1, z2], dim=0)  # (2*batch_size, projection_dim)

    # Cosine similarity matrix
    sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2)  # (2N, 2N)
    sim = sim / temperature

    # Mask to remove self-similarity (diagonal)
    mask = torch.eye(2 * batch_size, dtype=torch.bool, device=z.device)
    sim = sim.masked_fill(mask, -9e15)

    # Positive pairs: (i, i+N) and (i+N, i)
    pos_sim = torch.cat([sim[i, i+batch_size].unsqueeze(0) for i in range(batch_size)] +
                        [sim[i+batch_size, i].unsqueeze(0) for i in range(batch_size)])

    # Numerator: exp(positive similarity)
    numerator = torch.exp(pos_sim)

    # Denominator: sum of exp(all similarities except self)
    denominator = torch.exp(sim).sum(dim=1)

    loss = -torch.log(numerator / denominator).mean()
    return loss

Temperature \tau controls how hard the negatives are. Lower \tau = sharper distribution = model focuses more on hard negatives. I tested \tau \in \{0.05, 0.07, 0.1, 0.5\}. At 0.05, training was unstable (loss spiked randomly). At 0.5, loss converged too fast and representations were mushy. Sweet spot: 0.07.

Pretraining and Fine-Tuning

Pretraining: 200 epochs on unlabeled data (1920 samples), batch size 64, Adam optimizer with learning rate 3e-4, cosine annealing schedule. Takes about 2.5 hours on a single RTX 3090. I tracked alignment (average positive pair similarity) and uniformity (how spread out the representations are on the hypersphere)—both metrics from Wang & Isola (2020).

def train_simclr(encoder, proj_head, dataloader, epochs=200, lr=3e-4, device='cuda'):
    encoder.to(device)
    proj_head.to(device)
    optimizer = torch.optim.Adam(list(encoder.parameters()) + list(proj_head.parameters()), lr=lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

    for epoch in range(epochs):
        total_loss = 0
        for batch in dataloader:
            x = batch.to(device)  # (batch, 2048)

            # Generate two augmented views
            x1 = torch.stack([torch.from_numpy(augment_signal(s.cpu().numpy())) for s in x]).to(device)
            x2 = torch.stack([torch.from_numpy(augment_signal(s.cpu().numpy())) for s in x]).to(device)

            # Forward pass
            h1 = encoder(x1.unsqueeze(1))  # (batch, 512)
            h2 = encoder(x2.unsqueeze(1))
            z1 = proj_head(h1)  # (batch, 128)
            z2 = proj_head(h2)

            loss = nt_xent_loss(z1, z2, temperature=0.07)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        scheduler.step()
        if (epoch + 1) % 20 == 0:
            print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(dataloader):.4f}")

After pretraining, I froze the encoder and trained a linear classifier (512 → 4) on the 5-shot labeled set. 100 epochs, cross-entropy loss, learning rate 1e-3. This is where the magic happens—or doesn’t.

Results (5-shot, mean over 5 random seeds):
– SimCLR + linear probe: 89.3% ± 2.1%
– Supervised CNN (same architecture, trained from scratch): 55.2% ± 4.8%
– Fine-tuned encoder (unfroze last 2 ResBlocks): 91.7% ± 1.9%

Fine-tuning the encoder gains 2.4 points but risks overfitting. With 10 shots per class, the gap narrows: SimCLR hits 94.1%, supervised CNN reaches 78.3%. At 50 shots, supervised CNN catches up (96.2% vs 96.8%).

What Happens When Augmentations Go Wrong

I mentioned time stretching killed performance. Here’s why. Bearing fault characteristic frequencies depend on geometry and rotational speed:

f_{\text{BPFI}} = \frac{n}{2} f_r \left(1 + \frac{d}{D} \cos \phi \right)

where n is number of rolling elements, f_r is shaft frequency, d/D is ball-to-pitch diameter ratio, \phi is contact angle. Time stretching changes the apparent f_r, which breaks the physical relationship between signal and fault type. The model learns to classify based on augmentation artifacts instead of actual fault signatures.

I noticed this when inspecting t-SNE projections of the learned embeddings. With proper augmentations, fault classes formed tight clusters. With time stretching included, the clusters were smeared along a continuous manifold—the model had learned to encode stretch factor, not fault type.

Another failure: using too large noise \sigma = 0.2. At CWRU’s 12 kHz sampling rate and typical SNR ~15-20 dB, this overwhelmed the fault-induced modulation sidebands (usually 10-15 dB above noise floor). The model converged to a trivial solution: map everything to the same point.

Hyperparameter Sensitivity

Projection head dimensionality: I tried 64, 128, 256. At 64, loss plateaued early. At 256, training was slower with no accuracy gain. 128 was the Goldilocks zone.

Batch size: Contrastive learning benefits from large batches (more negatives). But CWRU samples are small. I tested 32, 64, 128. Performance peaked at 64—beyond that, diminishing returns due to limited dataset diversity.

Encoder depth: Tried ResNet-10 (2 blocks) and ResNet-34 (6 blocks). Shallower network underfit (86.1% accuracy). Deeper network overfit during fine-tuning (89.9% train, 88.2% test). ResNet-18 balanced capacity and generalization.

When Supervised Learning Still Wins

SimCLR’s advantage fades as labeled data grows. At 50 samples per class, the supervised baseline achieves 96.2%—only 0.6 points below SimCLR’s 96.8%. At 100 samples, they’re statistically tied.

And there’s the computational cost. Pretraining took 2.5 hours; supervised training from scratch took 8 minutes. If you’ve got abundant labeled data and tight deadlines, skip the contrastive pretraining.

SimCLR also assumes your unlabeled data distribution matches the labeled data. In the real world, you might collect vibration data from multiple machines, operating conditions, sensor placements. I tested this by pretraining on 0 hp load data, then fine-tuning on 1-3 hp labels. Accuracy dropped to 81.4%—domain shift bit hard.

Another edge case: sensor drift. CWRU data is pristine, collected over days. In a factory, accelerometer sensitivity drifts 5-10% per year due to mounting degradation, temperature cycling. If your pretraining data spans months, the learned representations might encode drift patterns instead of fault physics. I haven’t tested this empirically (don’t have long-term CWRU data), but it’s a known issue in PHM deployments.

Computational Constraints: Embedded vs Cloud

Can you run SimCLR inference on edge hardware? The encoder is ~2.1M parameters, 8.4 MB in FP32. On a Raspberry Pi 4, forward pass takes ~150 ms per sample (using PyTorch mobile). That’s fine for batch monitoring (analyze 1-second windows every 10 seconds), but not for real-time alarm systems.

Quantizing to INT8 via PyTorch’s torch.quantization.quantize_dynamic cuts size to 2.2 MB and latency to 45 ms, with only 0.8% accuracy drop. Still, if you need <10 ms latency, you’re looking at FPGA or custom ASIC.

Pretraining on-edge is impractical. Even with quantization-aware training, 200 epochs at 2.5 hours means 10+ days on a Pi. Better to pretrain in the cloud, deploy the frozen encoder.

Where SimCLR Fits in the PHM Toolbox

Use SimCLR when:
– You have lots of unlabeled vibration data but <10 labeled examples per fault class
– Your fault types are well-separated in the augmented feature space (check with t-SNE after pretraining)
– You can afford 2-3 hours of GPU time for pretraining
– Your labeled and unlabeled data come from similar operating conditions (same load, temperature range, sensor mounting)

Stick with supervised learning when:
– You have 50+ labeled samples per class
– Latency matters more than squeezing out 3% accuracy
– Your unlabeled data is noisy or from a different distribution
– You’re prototyping and need results in <1 hour

I’m curious whether momentum-based contrastive methods (MoCo, SimSiam) would outperform SimCLR on CWRU. MoCo maintains a queue of negatives, which might help with the small dataset size. But I haven’t tested it—my best guess is the gains would be marginal (<2%) given CWRU’s simplicity.

One thing I haven’t solved: how to choose augmentation hyperparameters (\sigma, shift range, scale range) without labeled validation data. Right now I’m tuning them via downstream accuracy, which defeats the point of unsupervised learning. Maybe there’s a way to use alignment/uniformity metrics as a proxy, but the correlation isn’t perfect. If anyone’s figured this out for vibration signals, I’d love to know.