From raw vibration data to a working fault detector, and what to check along the way

This post is about deep learning on industrial sensor data, in this case vibration signals from machines on a factory floor, used to detect bearing faults. Most posts on this topic are about architectures. In my experience the model is the easy part. The time goes into labels that don't exist, sensors that lie, splits that leak, and results that look great offline and fall apart on the floor.

I go through how I deal with those problems, and I use one complete training example to make each step concrete. It's a 1D CNN on the CWRU bearing dataset, built as a small training script (YAML config, Pydantic validation, optional mixed precision, optional gradient accumulation, Weights & Biases logging). Each section starts with how the step usually goes in industrial work, then shows it in the example, with code and tensor shapes.

In this post

• What to find out about industrial sensor data before you pick a model
• How to split time series data without leaking test information into training
• A small 1D CNN on raw vibration signals, with the shape math worked out
• Mixed precision, gradient accumulation and a config driven training script
• Sanity checks to run before you trust a high accuracy number
• What changes once you leave the offline numbers behind

Prerequisites

• Comfortable with Python and NumPy.
• Pytorch basics: nn.Module, Dataset and DataLoader, and how a training loop works.
• Basic calculus, enough for the intuition behind gradients, backprop and a loss function. You don't need to derive anything.

You don't need an industrial or signal processing background. I explain the machine, the sensor and the signal terms as they come up.

A note on scope. This is a small experiment that shows how I approach training, including the order of the steps, the checks and the habits that make reruns cheap. It is not a production system. A real deployment has far more moving parts than anything here, such as data pipelines that run continuously, messy and drifting sensors, many machines and operating conditions, monitoring, retraining, alerting, and people who have to trust the output. Everything here ran on a laptop with an RTX 3060 (6 GB). The numbers in this post come from one dataset, one model and one seed, and they show that the pipeline works, not that the model is ready for a plant.

Contents

1. Start with the data
2. Decide what the model predicts
3. Set up the project so reruns are cheap
4. Prepare the data
5. Pick a model you can debug
6. Mixed precision
7. Training settings
8. Run it
9. Check the result before trusting it
10. After the offline numbers

---

1. Start with the data

Before choosing any model, I write down what I know about the data. The questions that matter most:

• What is the sampling rate, and is it constant?
• Which sensors are there, and where are they mounted?
• What operating conditions exist (load, speed, product, ambient temperature), and how many of them does the data cover?
• How are timestamps produced, and are they aligned across sensors?
• What labels exist, who made them, and how many events of interest are in there?

Anything between the sensor and the file you receive (resampling, forward filling, tag renames, timestamp alignment) changes the signal before a model sees it. Find out what was done, because it decides which features you can trust.

The example uses the CWRU bearing dataset from Case Western Reserve University. A motor drives a shaft with bearings at the drive end and the fan end. The researchers put small defects into the bearings (0.007, 0.014 and 0.021 inches wide) using electric discharge machining, then recorded vibration with accelerometers while the motor ran under different loads.

I use the drive end accelerometer at 12 kHz with four motor loads (0, 1, 2 and 3 HP, shaft speed roughly 1730 to 1797 rpm). That gives 10 classes, which are normal plus inner race, ball and outer race faults at three sizes each. Each recording is a .mat file holding a few long arrays, and a fault recording runs about 120,000 samples, around 10 seconds.

The file numbers for the drive end data are below. The four numbers in each row are loads 0, 1, 2 and 3 HP, in order:

normal      97  98  99  100
inner_007   105 106 107 108
ball_007    118 119 120 121
outer_007   130 131 132 133
inner_014   169 170 171 172
ball_014    185 186 187 188
outer_014   197 198 199 200
inner_021   209 210 211 212
ball_021    222 223 224 225
outer_021   234 235 236 237

Outer race faults were recorded at several sensor positions (6, 3 and 12 o'clock). I only use the 6 o'clock ones, since mixing positions would turn this into a different experiment.

A sampling rate trap inside the dataset. The normal baseline files (97 to 100) look like they were recorded at 48 kHz. The fault files above are the 12 kHz set. The CWRU pages don't state a rate for the baseline files, but files 98 to 100 hold about 484,000 samples against about 121,000 in each fault file. That is four times as many samples for the same roughly 10 second recording, which fits 48 kHz against 12 kHz. If you feed both to a model as they are, the "normal" class differs from every other class in sampling rate, and a network can learn that instead of anything about bearings. It would score well and mean nothing. The fix is to bring the normal recordings down to 12 kHz with an anti-aliased decimation by 4 before anything else happens (code in section 4). This is exactly the "is the sampling rate constant?" question from the list above, and it showed up in a benchmark that people treat as clean.

The window length should come from the physics of the machine. I cut windows of 1200 samples, which is 0.1 seconds at 12 kHz. With the shaft turning at about 30 Hz, an outer race fault produces an impact roughly every 9 ms (around 107 Hz) and an inner race fault roughly every 6 ms (around 162 Hz). One window holds somewhere between 10 and 16 impacts, enough for a repeating pattern to show up. For a different machine you'd redo this arithmetic with its own fault frequencies.

2. Decide what the model predicts

This choice shapes everything after it, including the metrics. With plant data you typically have few faults, labels like "something went wrong sometime in March", and hardly any recordings of a clearly named fault type. That usually points to one of three setups:

• Supervised classification, when you have enough clean, named fault examples.
• Anomaly detection, where the model learns what normal looks like from normal data and flags whatever doesn't fit.
• Semi-supervised methods, a mix of lots of unlabeled data and a few labels.

The example uses classification because CWRU has a clean label for every recording, since the researchers made each fault themselves. The data pipeline (splitting, windowing, normalization, quality checks) carries over to the other setups unchanged. What changes is the head of the model and the loss.

3. Set up the project so reruns are cheap

Industrial work means rerunning the same experiment on new data over and over, whether it's a new month of recordings, a different machine or a revised label set. That goes better when nothing is hard coded and every setting lives in one file.

I kept the project to a handful of flat files, because one dataset and one model don't need more. I’m keeping the example deliberately small so each step is easy to follow, rather than building a reusable training framework.

industrial-data-training/
  config.yaml      every setting lives here
  config.py        Pydantic models for config.yaml
  download_data.py fetch the 40 CWRU .mat files into data/raw/cwru
  cwru.py          load the .mat files, split, window, normalize, build the DataLoaders
  cnn.py           the model
  train.py         the training loop
  evaluate.py      per-class report and confusion matrix for a checkpoint
  requirements.txt

train.py reads the YAML through config.py, asks cwru.py for the loaders, builds the model from cnn.py and runs the loop. I considered a registry (a name to class lookup, so a new dataset or model is one new file plus a line of YAML) and a trainer class, and left both out. With one dataset and one model they would be code that exists only to be read around. If a second dataset shows up, adding that structure then is easy.

Why Pydantic

YAML gives you nested dictionaries, so you end up writing cfg["train"]["lr"], and a typo in a key either fails late or quietly falls back to a default. Pydantic turns the dict into objects (so cfg.train.lr works, and your editor autocompletes it), checks types and ranges, and with extra="forbid" it rejects any key it doesn't know.

# config.py
class Section(BaseModel):
    model_config = ConfigDict(extra="forbid")  # a typo in the YAML is an error, not silently ignored


class DataConfig(Section):
    root: str
    window: int = Field(1200, gt=0)
    train_stride: int = Field(600, gt=0)
    eval_stride: int = Field(1200, gt=0)
    gap: int = Field(1200, ge=0)
    train_frac: float = Field(0.70, gt=0)
    val_frac: float = Field(0.15, gt=0)

    @model_validator(mode="after")
    def leave_room_for_test(self):
        if self.train_frac + self.val_frac >= 1:
            raise ValueError("train_frac + val_frac must be < 1 to leave a test split")
        return self


class TrainConfig(Section):
    seed: int = 42
    output_dir: str = "runs/cwru-cnn1d"
    epochs: int = Field(30, gt=0)
    batch_size: int = Field(64, gt=0)
    num_workers: int = Field(0, ge=0)
    lr: float = Field(1e-3, gt=0)
    weight_decay: float = Field(0.01, ge=0)
    warmup_steps: int = Field(100, gt=0)
    label_smoothing: float = Field(0.0, ge=0, lt=1)
    grad_clip: float = Field(1.0, gt=0)
    precision: Literal["fp32", "fp16", "bf16"] = "fp32"
    grad_accum_steps: int = Field(1, gt=0)

The ModelConfig and WandbConfig sections follow the same pattern, and a top level Config holds all four. The validators are small, but they mean a bad edit to the fractions at midnight fails in the first second and not an hour into a run.

There are no command line overrides. To try a different learning rate I edit config.yaml, or copy it and pass --config my_experiment.yaml. A copied file also doubles as a record of what a run used.

The config for the example run (this is the whole file):

data:
  root: data/raw/cwru      # folder with the numbered CWRU .mat files
  window: 1200
  train_stride: 600
  eval_stride: 1200
  gap: 1200                # samples dropped between time splits to avoid leakage
  train_frac: 0.70
  val_frac: 0.15           # the rest is test

model:
  widths: [32, 64, 128]
  first_kernel: 64
  first_stride: 8
  dropout: 0.3

train:
  seed: 42
  output_dir: runs/cwru-cnn1d
  epochs: 30
  batch_size: 64
  num_workers: 0
  lr: 0.001
  weight_decay: 0.01
  warmup_steps: 100
  label_smoothing: 0.0
  grad_clip: 1.0
  precision: fp32          # fp32 = mixed precision OFF; fp16 / bf16 = ON
  grad_accum_steps: 1      # 1 = accumulation OFF; e.g. 4 = ON

wandb:
  enabled: false
  project: cwru-cnn1d
  run_name: null

4. Prepare the data

Three things happen here, and the order matters.

Check quality first

With plant data I run checks before cutting a single window. A sensor stuck at one value for hours looks like perfectly calm operation to a model. Dropouts filled with zeros, clipped peaks and duplicated or out of order timestamps do similar damage, and they're much cheaper to catch before training than to explain after it. The checks are simple. Look for missing values, runs of identical samples, a std near zero and timestamp gaps. Thresholds depend on the sensor, so treat any numbers as a starting point.

CWRU is a clean lab dataset, so the example code only does the cheapest check, rejecting an empty recording or one with NaN or inf values. On real plant data I'd add the rest of that list at the same spot, in the function that loads a file.

Split by time, then window

The shortcut is to cut windows first, shuffle them all together and split afterwards. It gives a great number, but the number is inflated. Training windows overlap, so window 41 and window 42 share half their samples, and after shuffling one lands in train while the other lands in test. Without overlap it's still bad, because neighbouring windows from one recording share the same machine state, noise floor, sensor and fault, so the model has effectively seen them.

The better order is:

1. Load a recording and run the quality checks.
2. Split the raw signal by time, with the first 70% for train, the next 15% for val and the last 15% for test.
3. Leave a gap between the pieces.
4. Cut windows separately inside each piece.

No window crosses a boundary, no train window overlaps a val or test window, and the test windows come from a later stretch of the recording than anything the model trained on. With plant data the same idea applies at a bigger scale. Split across calendar time (the first few months for training, the next stretch for validation, the one after for testing), because machines wear, sensors drift and raw material changes. CWRU recordings are each one steady condition, so a split inside each recording is the closest equivalent.

Normalize with training statistics only

Z-scoring each window on its own is popular, but it throws away amplitude, and amplitude carries information (a 0.021 inch fault shakes harder than a 0.007 inch one). I use one global mean and std computed from the training windows and apply them everywhere. Val and test never feed into those numbers, otherwise test information leaks into training through the scaling. The stats are computed in float32 numpy, which matters again in the mixed precision section, and saved into meta.

meta goes into the checkpoint together with the weights, so the mean and std that went with a model travel with it. If you've ever had a model that worked in the notebook and fell apart in the script because someone forgot the scaler, you know why that matters.

The code

load_signal looks for the key that ends in _DE_time instead of building it, because the prefix (like X105) doesn't always match the file name. It's also where the quality check lives.

# cwru.py
def load_signal(path):
    path = Path(path)
    if not path.is_file():
        raise FileNotFoundError(f"Missing CWRU file: {path}. Put the numbered .mat files in data.root.")
    mat = loadmat(path)
    key = next((k for k in mat if k.endswith("_DE_time")), None)
    if key is None:
        raise ValueError(f"{path} has no MATLAB variable ending in '_DE_time'")
    signal = np.asarray(mat[key]).ravel().astype(np.float32)
    if signal.size == 0 or not np.isfinite(signal).all():
        raise ValueError(f"{path}: empty signal or NaN/inf values")
    return signal


def make_windows(signal, size, stride):
    if len(signal) < size:
        return np.empty((0, size), dtype=np.float32)
    n = (len(signal) - size) // stride + 1
    indices = np.arange(size)[None, :] + stride * np.arange(n)[:, None]
    return signal[indices].astype(np.float32)

make_windows builds an index matrix from np.arange(size) + stride * np.arange(n) and grabs signal[indices] in one go. There's no loop, and row i starts at sample i * stride, which makes it easy to check by hand. It copies the data, which is fine since everything fits in memory many times over. A segment shorter than one window gives a correctly shaped empty array, and prepare_data below raises if that ever leaves a whole split empty.

Training windows use a stride of 600, so they overlap by half. That roughly doubles the number of training examples from the same raw signal. The overlap is safe only because the split already happened on the raw signal. Val and test use a stride of 1200, so each sample is scored once and no two scored windows share data. The gap is one window (1200 samples) between pieces, so no window of any kind can touch both sides of a boundary.

Then the function that does the split, the windowing and the normalization:

def prepare_data(cfg):
    """Split each recording in time BEFORE windowing so windows never overlap across splits."""
    window, gap = cfg.window, cfg.gap
    classes = list(CWRU_FILES)
    parts = {s: ([], []) for s in ("train", "val", "test")}
    for label, name in enumerate(classes):
        for number in CWRU_FILES[name]:
            signal = load_signal(Path(cfg.root) / f"{number}.mat")
            if name == "normal":  # recorded at 48 kHz, the faults at 12 kHz
                signal = decimate(signal, 4, zero_phase=True).astype(np.float32)
            n = len(signal)
            a, b = int(n * cfg.train_frac), int(n * (cfg.train_frac + cfg.val_frac))
            pieces = {
                "train": (signal[:a], cfg.train_stride),
                "val": (signal[a + gap:b], cfg.eval_stride),
                "test": (signal[b + gap:], cfg.eval_stride),
            }
            for split, (piece, stride) in pieces.items():
                windows = make_windows(piece, window, stride)
                parts[split][0].append(windows)
                parts[split][1].append(np.full(len(windows), label, dtype=np.int64))
    arrays = {s: (np.concatenate(xs), np.concatenate(ys)) for s, (xs, ys) in parts.items()}
    if any(len(y) == 0 for _, y in arrays.values()):
        raise ValueError("An empty split: check recording lengths, window, gap and fractions")
    # Normalization statistics come from training windows only.
    mean, std = float(arrays["train"][0].mean()), float(arrays["train"][0].std())
    arrays = {s: ((x - mean) / std, y) for s, (x, y) in arrays.items()}
    return arrays, {"classes": classes, "mean": mean, "std": std}

The decimate call is the 48 kHz fix. Taking every fourth sample would alias anything above 6 kHz down into the band I care about, so I use scipy.signal.decimate, which low-pass filters first. The filter has a little passband ripple, about 1.1% in amplitude when I decimated a 100 Hz tone from 48 kHz, which I'm fine with here.

The slicing is the whole leakage protection. It's signal[:a] for train, signal[a + gap:b] for val and signal[b + gap:] for test. If you want to convince yourself, feed it a fake signal where each sample's value is its own index, and check that the largest train value is below the smallest val value and so on. Just make the fake signal long enough that val and test each hold at least one full window after the gaps. With 10,000 samples, val is left with only 300 and the arrays come back empty.

Finally the loaders. A TensorDataset is enough here, since a window is just a row of a tensor:

def build_loaders(cfg, batch_size, num_workers=0):
    arrays, meta = prepare_data(cfg)
    loaders = {}
    for split, (x, y) in arrays.items():
        dataset = TensorDataset(torch.from_numpy(x).unsqueeze(1), torch.from_numpy(y))  # (N, 1, window)
        loaders[split] = DataLoader(dataset, batch_size=batch_size, shuffle=split == "train",
                                    num_workers=num_workers)
    return loaders, meta

The unsqueeze(1) is there because Conv1d expects (batch, channels, samples), and a vibration window has one channel. Only the training loader shuffles.

One effect of the decimation to keep in mind is that the normal files come out slightly shorter than the fault files once they're at 12 kHz, so the normal class ends up with 490 training windows against about 563 for each fault class. That's a mild imbalance and I left it alone, but it's the first thing I'd look at if the normal row of the confusion matrix looks off.

5. Pick a model you can debug

For a first model on industrial data, small is better. When something goes wrong (and with messy data something usually does), a model with 35 thousand parameters has fewer places for the problem to hide, and it trains in minutes on whatever hardware you have. A bigger model can come later, once the pipeline is trusted.

The example is three conv blocks and a linear head.

Input                             (B, 1, 1200)
Conv1d(1, 32, k=64, s=8, p=28)    (B, 32, 150)
BatchNorm1d + ReLU
MaxPool1d(2)                      (B, 32, 75)
Conv1d(32, 64, k=3, p=1)          (B, 64, 75)
BatchNorm1d + ReLU
MaxPool1d(2)                      (B, 64, 37)
Conv1d(64, 128, k=3, p=1)         (B, 128, 37)
BatchNorm1d + ReLU
MaxPool1d(2)                      (B, 128, 18)
AdaptiveAvgPool1d(1) + Flatten    (B, 128)
Dropout(0.3)
Linear(128, 10)                   (B, 10)

B is the batch size, 64 in the config.

A conv's output length is floor((L + 2p - k) / s) + 1. For the first layer that's (1200 + 56 - 64) / 8 + 1 = 150. The padding of 28 was picked so the division comes out exact. Max pooling with size 2 halves the length and rounds down, so 150 becomes 75, 75 becomes 37 (the odd leftover sample is dropped) and 37 becomes 18.

The first conv is the unusual one. A kernel of 64 samples covers about 5 ms of signal, the same scale as the gap between fault impacts, so it acts like a set of learned filters sitting directly on the raw waveform. A stride of 8 shrinks the sequence from 1200 to 150 right away, which makes every later layer cheaper.

Each conv is followed by BatchNorm, ReLU and a max pool. The convs have bias=False because BatchNorm adds its own shift. Pooling keeps the stronger response out of every pair of neighbours, so a fault impact landing a few samples earlier or later changes the features very little.

The second and third convs use kernel 3. They sit on sequences that have already been shrunk, so each output position covers a long stretch of the original signal, and the receptive field grows with every pooling step. Channels go 32, 64, 128, the usual trade of more channels for shorter sequences.

After the last block, adaptive average pooling averages over time. Flattening the (128, 18) map would give 2304 numbers and a big linear layer. Averaging keeps the head tiny and makes the model indifferent to where in the window things happen. Dropout of 0.3 comes next, then one linear layer to the 10 class scores (logits).

The parameter count:

Conv1   1 * 32 * 64      = 2,048
Conv2   32 * 64 * 3      = 6,144
Conv3   64 * 128 * 3     = 24,576
BatchNorm (2 per channel, 32+64+128 channels) = 448
Linear  128 * 10 + 10    = 1,290
Total                      34,506

I checked this against sum(p.numel() for p in model.parameters()) and it matches.

# cnn.py
def conv_block(c_in, c_out, kernel, stride=1, padding=0):
    return nn.Sequential(
        nn.Conv1d(c_in, c_out, kernel, stride=stride, padding=padding, bias=False),
        nn.BatchNorm1d(c_out),
        nn.ReLU(inplace=True),
        nn.MaxPool1d(2),
    )


class CNN1D(nn.Module):
    def __init__(self, num_classes=10, widths=(32, 64, 128), first_kernel=64, first_stride=8, dropout=0.3):
        super().__init__()
        layers = [conv_block(1, widths[0], first_kernel, first_stride, padding=(first_kernel - first_stride) // 2)]
        layers += [conv_block(a, b, 3, padding=1) for a, b in zip(widths[:-1], widths[1:])]
        self.features = nn.Sequential(*layers)
        self.pool = nn.AdaptiveAvgPool1d(1)
        self.head = nn.Sequential(nn.Dropout(dropout), nn.Linear(widths[-1], num_classes))

    def forward(self, x):
        return self.head(self.pool(self.features(x)).flatten(1))

The model is a plain nn.Module that takes a batch of windows in and gives class scores out. The loss lives in the training loop, as a single F.cross_entropy call. That also keeps the anomaly detection swap from section 2 small, since it only needs a different head in this file and a different loss line in train.py. The loop around them doesn't care.

I don't use gradient checkpointing (trading extra compute for lower activation memory) here. 34k parameters on a batch of 64 use so little memory that there's nothing to save, and checkpointing doesn't play well with BatchNorm anyway.

Before any training, the loss on the first batch should sit near ln(10), about 2.30, since there are 10 balanced classes and the model starts out guessing. Anything from 2.2 to 2.6 is normal for a random init. A first loss of 0.4 or 14 points at wrong labels, the wrong loss, or unnormalized input.

6. Mixed precision

Mixed precision does the heavy math in 16 bit numbers and keeps the touchy parts in 32 bit. It saves memory and can run faster on the right GPU, which starts to matter as models and datasets grow (long recordings at high sampling rates get big fast). I keep it as a switch in the config, off by default (precision: fp32), and turning it on is a one word change.

There are two 16 bit formats. fp16 has a small range, so tiny gradients can underflow to zero, and fp16 training therefore needs a gradient scaler that multiplies the loss by a big number before backward and divides the gradients afterwards. bf16 has the same range as fp32 with fewer digits of precision. Gradients don't underflow, so there's no scaler. bf16 is easier to live with, and it's what I'd try first.

In train.py, the one config value turns into an autocast context and a scaler:

# train.py
dtype = {"fp32": None, "fp16": torch.float16, "bf16": torch.bfloat16}[t.precision]
if t.precision == "fp16" and device.type != "cuda":
    raise ValueError("fp16 requires CUDA; use fp32 or bf16 on CPU")
autocast = (lambda: torch.autocast(device.type, dtype=dtype)) if dtype else nullcontext
scaler = torch.amp.GradScaler(device.type, enabled=t.precision == "fp16")

With fp32 the autocast is a no-op, and the scaler is only enabled for fp16, so one code path in the loop serves all three precisions. Asking for fp16 without a GPU is an error and not a silent fallback, since I'd rather know that the run isn't doing what I asked for.

Inside the autocast block, pytorch decides op by op. Convolutions and matrix multiplies run in the low precision type. Softmax, cross entropy and big sums stay in float32, and BatchNorm keeps its running statistics in float32. The weights are always stored in float32 and the optimizer updates those copies, so nothing here trains a pure 16 bit model. Parts of the forward and backward pass run in 16 bit.

Two details are easy to miss. First, bf16 keeps only about 3 significant digits, so the dataset mean and std are computed in float32 numpy, and the tensors only get cast to bf16 inside the first conv. A mean computed on bf16 data would be slightly off and nothing would warn you. Second, the logits are cast with .float() before the loss, which spells out what autocast does anyway for cross entropy.

The training step, with gradient accumulation (next section) in the same loop:

for i, (x, y) in enumerate(loader, 1):
    x, y = x.to(device), y.to(device)
    with autocast():
        logits = model(x)
    loss = F.cross_entropy(logits.float(), y, label_smoothing=t.label_smoothing)
    scaler.scale(loss / accum).backward()

    if i % accum == 0 or i == len(loader):  # also flush a partial last group
        scaler.unscale_(optimizer)
        nn.utils.clip_grad_norm_(model.parameters(), t.grad_clip)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)
        scheduler.step()

The unscale has to come before the clip. Clipping the scaled gradients compares them against a threshold meant for the real ones, and a clip at 1.0 then does something quite different from what was intended. In fp16 the scaler also notices an inf or NaN gradient and skips that optimizer step by itself. In bf16 and fp32 nothing checks for it, so a bad batch would go straight into the weights. If you see the loss turn to NaN, that's where to look first.

Don't expect a speedup on this model. With 34k parameters the GPU mostly waits on Python and data loading, and on CPU bf16 can even be slower. It's worth flipping on because it's the same code path bigger models will use, and because it makes a cheap experiment. Train once with precision: bf16 and once with fp32, same seed. The two runs should land close to each other. If they don't, there's a bug somewhere.

7. Training settings

These are starting points, and I'll say why I picked each one.

Batch size is 64. With 5,559 training windows that's 87 optimizer steps per epoch and 2,610 steps over 30 epochs. Each step is a little noisy, and the batch is big enough for BatchNorm statistics to settle.

Gradient accumulation is off (grad_accum_steps: 1). With grad_accum_steps: 4, four micro batches get added up before each optimizer step, so the effective batch is 64 x 4 = 256 (each micro batch loss is divided by 4, which makes the gradients an average and not a sum). It's meant for models that don't fit in memory at a large batch size. This one fits, so each optimizer step uses one micro batch. When the number of batches in an epoch doesn't divide evenly, the leftover micro batches at the end still get their own optimizer step. The learning rate schedule is counted in optimizer steps, not micro batches, so turning accumulation on changes how many steps the schedule has.

The optimizer is AdamW with weight decay 0.01, AdamW's usual default. I made one simplification here. The decay applies to every parameter, including biases and BatchNorm scales. The more careful setup puts those in a separate group with no decay, since decaying a BatchNorm scale just pulls it toward zero for no reason. For a model this small I haven't seen it matter, so I dropped the parameter groups. It's a few lines if you want it back.

The peak learning rate is 1e-3, which usually works for small CNNs with BatchNorm. Without normalization layers, start lower. The schedule climbs linearly for the first 100 steps (a bit over one epoch), then follows a cosine curve down to zero. Warmup protects the weights from the early, chaotic steps, and the slow decay at the end lets the model settle. The scheduler steps once per optimizer step and not once per epoch, so the curve is smooth.

def lr_scale(step):
    if step < t.warmup_steps:
        return (step + 1) / t.warmup_steps
    progress = (step - t.warmup_steps) / max(1, total_steps - t.warmup_steps)
    return 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_scale)

Gradient clipping is at 1.0, so any gradient with a norm above that gets scaled down. For this model it should rarely trigger.

Dropout is 0.3 and label smoothing is 0. If the model turns out overconfident on val windows, smoothing is the first thing to try.

Validation runs once per epoch and the loop keeps the checkpoint with the best validation accuracy. There's no early stopping and no moving average of the weights. Both are common, but for 30 epochs of a small model I'd rather see the whole curve than have a rule cut it short, and the best checkpoint is kept anyway. At the end, the script loads that best checkpoint and scores it once on the test split.

Checkpoints. Whenever validation accuracy improves, the loop writes best.pt to the run folder. It holds the weights, the model settings and meta (class names, mean and std), which is everything evaluate.py needs to rebuild the model. There's no resume. A 34k parameter model trains in minutes, so if a run dies I start it again. The moment a run takes hours, that's the first thing I'd add back.

if val_acc > best_acc:
    best_acc = val_acc
    torch.save({"model": model.state_dict(), "meta": meta, "model_cfg": cfg.model.model_dump()},
               out_dir / "best.pt")

One seed in the config seeds Python, numpy and torch, and the device is picked automatically (GPU if there is one). The script runs on a single device and has no DDP, which is plenty for 34k parameters.

8. Run it

First you need the data. The 40 .mat files come from the Case Western Reserve University Bearing Data Center, from its "Normal Baseline Data" page and its "12k Drive End Bearing Fault Data" page. You can download them by hand, or run python download_data.py in the repo, which fetches every file listed in CWRU_FILES into data/raw/cwru and skips the ones you already have.

python download_data.py

Then train:

python train.py

Run the default config first, then change one value at a time, ideally in a copy of the config (python train.py --config exp1.yaml), so every experiment is "the same config plus one change". Each run writes best.pt to its output_dir under runs/.

Every epoch prints one line with the train loss and accuracy, the validation loss and accuracy, and the current learning rate. With wandb.enabled: true (after a one time wandb login) the same numbers go to Weights & Biases as curves, together with the full config of the run, and the final test loss, test accuracy and best validation accuracy land in the run summary. Keeping every run in W&B is what makes "the same config plus one change" comparisons easy a month later.

Three things in the log are worth watching, in this order. The first epoch's train loss should be somewhere near 2.30 (it's an average over the epoch, so a bit lower than the very first batch). Train loss should then drop smoothly through warmup, and a spike there usually means a learning rate that's too high or inputs that aren't normalized. Val loss should follow train loss down and then flatten, since a widening gap between them is the classic overfit sign.

Then evaluation:

python evaluate.py --checkpoint runs/cwru-cnn1d/best.pt

evaluate.py loads best.pt, rebuilds the model from the settings stored inside it, runs the test split (--split val for validation) and prints the accuracy, per-class precision and recall and a confusion matrix. It rebuilds the data from the config, so the CWRU files need to be reachable, and the mean and std get recomputed from the same training windows. There is no script for predicting on new signals. When you need one, the saved meta has the mean and std to apply, and the signal has to be at 12 kHz and cut into 1200 sample windows the same way as in make_windows.

Results of the example run (the default config, with precision: bf16 and wandb.enabled: true, on an RTX 3060 Laptop GPU):

Train / val / test windows:  5559 / 552 / 552
Best val accuracy:           1.0000 (first reached at epoch 4)
Test accuracy:               1.0000 (test loss 0.0158)
bf16 vs fp32 test accuracy:  1.0000 vs 1.0000 (test loss 0.0158 vs 0.0078)

The first epoch's train loss was 1.60, which fits the "near 2.30, a bit lower as an epoch average" expectation from above, and train accuracy passed 99% by epoch 4. Validation accuracy hit 100% at epoch 4 and stayed there, and the validation loss kept shrinking to 0.0003.

Val loss sits below train loss for the whole run. That's expected here, because the train number is averaged over the epoch while the weights are still changing, and it's measured with dropout on, while validation uses the final weights of the epoch with dropout off. It isn't a sign of a bug.

About that 100%. My first reaction to this log was that the model had overfit. It hasn't, in the usual sense, since val and test are perfect too, and nothing here shows a train/val gap. It does mean this evaluation is easy, for two reasons. CWRU is a saturated benchmark, and even a plain random split of it is routinely scored above 99%. And my time split, while leak-free at the window level, still draws train, val and test from the same 40 recordings, with the same motor, sensor, four loads and faults. The model never has to handle a situation it hasn't seen. The test loss (0.016) being 50 times the val loss (0.0003) hints at a small shift between the early and late parts of each recording, but not enough to cost a single window. So the number says the pipeline works end to end. It says very little about how the model would do on a new machine or a new load, which is what the held-out operating condition test in section 10 is for.

9. Check the result before trusting it

A high test accuracy means little until the simple explanations (a bug, a leak, bad labels) are ruled out. The point of these checks isn’t to squeeze a better score out of this run. It’s to give you a few ways to catch mistakes before you trust a result. These checks go from cheapest to most telling, and they apply to any dataset, not just this one. The first and third work with the code as it stands. The second and fourth need a small temporary edit, which I show as a diff, and which you revert after the run.

1. Overfit a tiny batch. Train on 64 windows with dropout off. The loss should fall toward zero within a few hundred steps. A model that can't memorize 64 examples has a bug in the model or the loop, and tuning on the full data won't fix it.

import torch
from torch.nn import functional as F

from cnn import CNN1D
from config import load_config
from cwru import build_loaders

cfg = load_config()
loaders, meta = build_loaders(cfg.data, batch_size=64)
x, y = next(iter(loaders["train"]))        # the train loader shuffles, so all classes show up

net = CNN1D(dropout=0.0)
opt = torch.optim.AdamW(net.parameters(), lr=1e-3)
net.train()
for step in range(300):
    logits = net(x)
    loss = F.cross_entropy(logits, y)
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 100 == 0:
        print(step, loss.item(), (logits.argmax(-1) == y).float().mean().item())

I ran exactly this on a batch from the CWRU training windows (all 10 classes present). The loss went from 2.65 to 0.0035 in 300 steps, and accuracy hit 100% well before the end.

2. Shuffle the labels. Permute the training labels and train again. Accuracy should stay near 10%, which is chance for 10 classes, and val loss should hover around ln(10). Anything noticeably higher means labels are leaking in from somewhere. In train.py, right after build_loaders, add:

ds = loaders["train"].dataset
ds.tensors = (ds.tensors[0], ds.tensors[1][torch.randperm(len(ds))])

Shuffled labels, val accuracy:  0.0851 at epoch 30 (best 0.1522), val loss 2.34, test accuracy 0.1413

That's what I got. Val accuracy sat around chance (8.5% at the end, and it bounced around between epochs), and val loss stayed near ln(10) = 2.30. Train accuracy did creep up to about 22%, since a network can slowly memorize shuffled labels, but nothing carried over to val. Chance is the right result here.

3. Read the confusion matrix. One average number hides where the errors are, and a messy row can sit under a 98% overall figure. evaluate.py prints it with rows as the true class and columns as the predicted class, plus per-class precision and recall. Ball faults are often the shakiest class on CWRU, and the smallest faults get confused with normal more than the big ones do.

python evaluate.py --checkpoint runs/cwru-cnn1d/best.pt

4. Measure what leaking costs. This one is the most telling, because section 4 argued that the shortcut (cut windows first, shuffle, split randomly) inflates the result, and here you can try to measure by how much. Train the same model with the same config, but with the leaky split. In prepare_data in cwru.py, the time split goes away. Each whole recording is windowed with the train stride, everything is shuffled together, and only then cut 70/15/15.

             n = len(signal)
-            a, b = int(n * cfg.train_frac), int(n * (cfg.train_frac + cfg.val_frac))
-            pieces = {
-                "train": (signal[:a], cfg.train_stride),
-                "val": (signal[a + gap:b], cfg.eval_stride),
-                "test": (signal[b + gap:], cfg.eval_stride),
-            }
-            for split, (piece, stride) in pieces.items():
-                windows = make_windows(piece, window, stride)
-                parts[split][0].append(windows)
-                parts[split][1].append(np.full(len(windows), label, dtype=np.int64))
-    arrays = {s: (np.concatenate(xs), np.concatenate(ys)) for s, (xs, ys) in parts.items()}
+            windows = make_windows(signal, window, cfg.train_stride)  # whole recording, no time split
+            parts["train"][0].append(windows)
+            parts["train"][1].append(np.full(len(windows), label, dtype=np.int64))
+    x, y = (np.concatenate(p) for p in parts["train"])
+    order = np.random.default_rng(0).permutation(len(y))  # shuffle all windows together, then cut
+    a, b = int(len(y) * cfg.train_frac), int(len(y) * (cfg.train_frac + cfg.val_frac))
+    arrays = {"train": (x[order[:a]], y[order[:a]]), "val": (x[order[a:b]], y[order[a:b]]),
+              "test": (x[order[b:]], y[order[b:]])}

I ran this edit on the real files. With it in place, train windows overlap by half and neighbours from the same recording end up on both sides of the split (5,579 train, 1,195 val and 1,196 test windows). Put the two test numbers next to each other:

Time ordered split:  1.0000 test accuracy (test loss 0.0158)
Random window split: 1.0000 test accuracy (test loss 0.0204)

On CWRU the two numbers came out the same, because even the time ordered split already scores 100% and accuracy can't go higher. The leak is real here, since neighbouring windows from the same recording sit on both sides of the random split, but this dataset is too easy to show it. On a harder dataset the random split would come out higher, and that gap is the extra optimism you'd walk into a meeting with. Don't read the identical numbers as proof that leakage doesn't matter. Revert the edit afterwards, so the repo keeps the honest split.

10. After the offline numbers

Offline results are where the checking starts. With plant data I'd add a few more steps before anyone relies on the model:

• Accuracy is a poor headline when faults are rare. False alarms per week, and how many real problems got caught early enough to matter, are closer to what the people reading the report care about.
• Run the model in the background next to the existing process, watching with no control over anything, and read its alerts for a few weeks. That's a much better test than any offline split.
• Operators won't act on "the model says 0.93". They'll ask which signal moved and by how much, so plan for some explanation (which channel, which frequency band, how far from normal) next to the score.
• Normalization stats and model weights from an early period go stale as the machine wears and sensors drift. Decide up front how drift gets noticed and when the model gets retrained.

One more test I'd run, though I haven't here, is to hold out a whole operating condition. Train on loads 0, 1 and 2 and test on recordings from load 3, which the model has never seen in any form. In a plant, a change of operating condition (a new product, a different speed, a worn part) is the normal case, and this is the closest thing to it that CWRU offers. A number far below the normal test accuracy would show how badly the model handles new conditions.

CWRU is also a friendly benchmark. Faults were seeded on purpose, conditions are steady, and the labels are clean. A high score here means the pipeline works. It tells you nothing about whether a model would survive in a real plant. This whole example is a simple experiment, and a production setup has many more variables than the ones covered here.

Next things to try on the example itself are an anomaly detection head trained only on normal data and scored on how well it separates the fault classes, augmentations that imitate sensor noise and amplitude changes, and a plot of the 32 first layer filters to see what they picked up.

The full code is in the repo. If you spot a mistake in the file number map or the shape math, tell me and I'll fix it.