Transformer for Remaining Useful Lifetime (RUL) Prediction

Predictive maintenance is a generalized methodology that focuses on real-time monitoring and diagnosis of systems and processes. The main task is to take action as soon as a component exhibits a certain behavior that results in partial or total machine failure, performance degradation or quality reduction. A core component of predictive maintenance deals with predicting of the Remaining Useful Lifetime (RUL) of systems. This article presents the basics for training data-driven models, more precisely Transformers, and using them to predict the RUL. The following points are addressed:

1. The preparation of the dataset for extracting RUL information from multivariate time series data
2. The formation of the Transformer model with integration in Lightning
3. The training of the model with Lightning and Weights&Biases
4. The inference of the model and generation of Run-To-Failure (RTF) plots

You can find the repository for this article on GitHub.

Information extraction pipeline

The system parameters at runtime are used to predict the RUL. To apply the methodology to arbitrary systems, the data is not recorded between fixed time points, but at inspection intervals. The design matrix has the form Xₜ ∈ ℝᵏ with t = (1, … ,T), where T describes the number of inspection intervals and k the number of parameters of the system. The prediction of the RUL can be understood as mapping yₜ = f(Xₜ), where yₜ is the RUL and f is an arbitrary model. In this case the Vanilla Transformer.

The dataset used is the Commercial Modular Aero-Propulsion System Simulation (C-MAPSS) dataset, which originates from the C-MAPSS simulator and can generate commercial turbofan engine data. A total of four different multivariate time series subdatasets are included, which contain different operating states and fault modes. The dataset contains a total of 26 features: the unit, the current cycle, three operating conditions and the remaining sensor data.

[embed]CMAPSS Jet Engine Simulated Data | NASA Open Data Portal Edit descriptiondata.nasa.gov

There are three important points to note regarding data preprocessing:

1. For the first inspection intervals after Begin-of-Life (BOL), the RUL is set to a maximum value of 125
2. After the raw dataset has been loaded, it must first be normalized using a min-max scaler
3. To take advantage of the multivariate time series, a Sliding Time Window (STW) procedure is applied

STW procedure

To generate the training data, it is iterated over the normalized raw training dataset. The index i stands for the individual unit. An auxiliary dataset is now generated for each unit. By applying the STW procedure, it is not iterated over all cycles, but over (cycles-window_size+1) and the design matrix in the form Xⱼ with shape (window_size x features) is generated for a specific value of the RUL.

for i in range(1, int(np.max(train[:, 0])) + 1):
    ind = np.where(train[:, 0] == i)
    ind = ind[0]
    data_temp = train[ind, :]
    for j in range(len(data_temp) - window_size + 1):
        train_X.append(data_temp[j:j + window_size, 2:].tolist())
        train_RUL = len(data_temp) - window_size - j
        if train_RUL > RUL_max:
            train_RUL = RUL_max
        train_y.append(train_RUL)

To generate the test data, the normalized raw test dataset is iterated accordingly. The index i also stands for the individual unit and an auxiliary dataset is also generated for each unit. It is possible that fewer cycles are available for the test data than are required for the STW procedure. If this is the case, a spline interpolation is performed.

for i in range(1, int(np.max(test[:, 0])) + 1):
    ind = np.where(test[:, 0] == i)
    ind = ind[0]
    data_temp = test[ind, :]
    if len(data_temp) < window_size:
        data_temp_a = []
        for myi in range(data_temp.shape[1]):
            x1 = np.linspace(0, window_size - 1, len(data_temp))
            x_new = np.linspace(0, window_size - 1, window_size)
            tck = interpolate.splrep(x1, data_temp[:, myi])
            a = interpolate.splev(x_new, tck)
            data_temp_a.append(a.tolist())
        data_temp_a = np.array(data_temp_a)
        data_temp = data_temp_a.T
        data_temp = data_temp[:, 2:]
    else:
        data_temp = data_temp[-window_size:, 2:]

    data_temp = np.reshape(data_temp, (1, data_temp.shape[0], data_temp.shape[1])) 

    if i == 1:
        test_X = data_temp
    else:
        test_X = np.concatenate((test_X, data_temp), axis=0)

    if RUL[i - 1] > RUL_max:
        test_y.append(RUL_max)
    else:
        test_y.append(RUL[i - 1])

Once the training and test datasets have been successfully generated, they are converted from lists into arrays, reshaped and saved in an “h5” file.

train_X = (np.array(train_X)).reshape(len(train_X), window_size, features)
train_y = (np.array(train_y)/RUL_max).transpose()
test_X = (np.array(test_X)).reshape(len(test_X), window_size, features)
test_y = (np.array(test_y)/RUL_max).transpose()

save_dir = f"{cwd}/{subdataset}"
with h5py.File(f"{save_dir}/{subdataset}.h5", 'w') as f:
    f.create_dataset('X_train', data=train_X)
    f.create_dataset('Y_train', data=train_y)
    f.create_dataset('X_test', data=test_X)
    f.create_dataset('Y_test', data=test_y)

A Run-to-Failure (RTF) dataset is also generated for the evaluation. A unit is first selected and then the dataset is created in the same way as the previous process.

ind = np.where(test[:, 0] == rtf["unit"])
ind = ind[0]
data_temp = test[ind, :]
data_RUL = RUL[rtf["unit"] - 1]
for j in range(len(data_temp) - window_size + 1):
    rtf_X.append(data_temp[j:j + window_size, 2:].tolist())
    test_RUL = len(data_temp) + data_RUL - window_size - j
    if test_RUL > RUL_max:
        test_RUL = RUL_max
    rtf_y.append(test_RUL)

rtf_X = (np.array(rtf_X)).reshape(len(rtf_X), window_size, features)
rtf_y = (np.array(rtf_y)/RUL_max).transpose()

print(rtf_X.shape)
print(rtf_y.shape)

with h5py.File(f"{save_dir}/RTF.h5", 'w') as f:
    f.create_dataset('RTF_X', data=rtf_X)
    f.create_dataset('RTF_Y', data=rtf_y)

Building the Vanilla Transformer model as Lightning module

In the following, we will take a look at the modeling of the Vanilla Transformer. It is created as a Lightning module, which results in simplified training. The following model components are created as individual modules:

• Encoder
• Decoder
• Parameter extraction for the decoder

The encoder and decoder are designed in the pre-LN version.

Vanilla Transformer in pre-LN variant

A PyTorch module consists of an init function, which initializes the declared layers when the model is initialized, and a forward pass, which defines the forward calculation of the module. In addition to the LayerNorm and the Linear layers, the MultiheadAttention variant already implemented in PyTorch is used. It is important to set the “batch_first” variable here, as otherwise the batch dimension is not in first place.

class Encoder(torch.nn.Module):
    """
    Encoder
    """
    def __init__(self, dimv, dimatt, n_heads, drop):
        super().__init__()
        self.ln1 = torch.nn.LayerNorm(dimv, eps=1e-5)
        self.attn = torch.nn.MultiheadAttention(
            dimatt,
            n_heads,
            drop,
            batch_first=True
        )
        self.ln2 = torch.nn.LayerNorm(dimv, eps=1e-5)
        self.ffn1 = torch.nn.Linear(dimv, dimv)
        self.ffn2 = torch.nn.Linear(dimv, dimv)

    def forward(self, x, mask=None):
        """
        Forward pass
        """
        a = self.ln1(x)
        a, _ = self.attn(a, a, a, attn_mask=mask)
        x = self.ln2(a + x)
        a = self.ffn2(torch.nn.ReLU()(self.ffn1(x)))
        return x + a

The decoder module is set up in the same way as the encoder module. However, two MultiheadAttention layers must be initialized, as the first performs a so-called self-attention and the second a cross-attention to the generated feature map of the encoder. The Linear layers and the LayerNorm are implemented according to the encoder.

class Decoder(torch.nn.Module):
    """
    Decoder
    """
    def __init__(self, dimv, dimatt, n_heads, drop):
        super().__init__()
        self.ln1 = torch.nn.LayerNorm(dimv, eps=1e-5)
        self.attn1 = torch.nn.MultiheadAttention(
            dimatt,
            n_heads,
            drop,
            batch_first=True
        )
        self.ln2 = torch.nn.LayerNorm(dimv, eps=1e-5)
        self.attn2 = torch.nn.MultiheadAttention(
            dimatt,
            n_heads,
            drop,
            batch_first=True
        )
        self.ln3 = torch.nn.LayerNorm(dimv, eps=1e-5)
        self.ffn1 = torch.nn.Linear(dimv, dimv)
        self.ffn2 = torch.nn.Linear(dimv, dimv)

    def forward(self, x, enc):
        """
        Forward pass
        """
        a = self.ln1(x)
        a, _ = self.attn1(a, a, a, key_padding_mask=None)
        x = self.ln2(a + x)
        a, _ = self.attn2(x, enc, enc, key_padding_mask=None)
        x = self.ln3(a + x)
        a = self.ffn2(torch.nn.ReLU()(self.ffn1(x)))
        return x + a

The following is an important point in the application of Transformers for multivariate time series and RUL prediction. The encoder of the Transformer gets to see the time series data from the current STW, but not the decoder. The decoder is presented with the last two points in time from the current STW, including extracted statistical parameters for the entire STW. This is done by the parameter extraction module for the decoder. In this case, the minimum, maximum, mean and standard deviation are used as statistical parameters.

class ParamExtraction(torch.nn.Module):
    """
    Parameter Extraction
    """
    def __init__(self) -> None:
        super().__init__()

    def forward(self, x):
        """
        Forward pass
        """
        t_min = torch.unsqueeze(torch.min(x, dim=1).values, dim=1)
        t_max = torch.unsqueeze(torch.max(x, dim=1).values, dim=1)
        t_mean = torch.unsqueeze(torch.mean(x, dim=1), dim=1)
        t_std = torch.unsqueeze(torch.std(x, dim=1), dim=1)
        ret = torch.cat([x[:, -2:, :], t_min, t_max, t_mean, t_std], dim=1)
        return ret

The Vanilla Transformer can now be assembled from the individual modules. First, the model is derived as a LightningModule. Second, the hyperparameters of the model are now set in the initialization. The dimension of the input data, the hidden dimension of the model, as well as the number of encoder and decoder blocks and their heads play a role here. Third, the previously defined modules are then stacked as a module list. Fourth, simple linear layers are used for the embedding layer of the encoder and decoder. And finally, as this task is a regression, a linear layer is used as the regression head of the model.

class VanTransLitModule(pl.LightningModule):
    """
    Vanilla Transformer Module
    """
    def __init__(self):
        super().__init__()

        self.input_size = (40, 17)
        self.d_model = 64
        self.heads = 4
        self.nencoder = 2
        self.ndecoder = 1
        self.dim_val = self.d_model
        self.dec_l = 6
        self.output_size = 1

        self.encoder = torch.nn.ModuleList(
            [Encoder(self.dim_val, self.d_model, self.heads, 0)
             for _ in range(self.nencoder)])

        self.decoder = torch.nn.ModuleList(
            [Decoder(self.dim_val, self.d_model, self.heads, 0)
             for _ in range(self.ndecoder)])

        self.pos = torch.nn.ModuleList(
            [PositionalEncoding(self.input_size[0], self.dim_val)])

        self.decinp = torch.nn.ModuleList(
            [ParamExtraction()]
        )

        self.encemb = torch.nn.Linear(self.input_size[1], self.dim_val)
        self.decemb = torch.nn.Linear(self.input_size[1], self.dim_val)

        self.ln1 = torch.nn.LayerNorm(self.dim_val, eps=1e-5)
        self.out = torch.nn.Linear(self.dec_l*self.dim_val,
                                   self.output_size)

    def forward(self, x):
        """
        Forward pass
        """
        e = self.encoder[0](self.pos[0](self.encemb(x)))

        for enc in self.encoder[1:]:
            e = enc(e)

        p = self.ln1(e)

        d = self.decoder[0](self.decemb(self.decinp[0](x)), p)
        for dec in self.decoder[1:]:
            d = dec(d, p)

        x = self.out(torch.nn.ReLU()(d.flatten(start_dim=1)))

        return x

PS: training_step, validation_step, test_step and other Lightning-specific functions can also be defined, but these are not listed here.

Transformer training using Lightning and Weights&Biases

This section deals with the training of the Vanilla Transformer on the C-MAPSS dataset. The Lightning Framework for training and Weights&Biases as an MLOps platform for model monitoring are discussed.

With the publication of the Transformer architecture, an important aspect in the context of training was also mentioned. Transformers require Learning Rate Scheduling (LRS) with a warm-up phase for stable learning behavior. This is not implemented natively in PyTorch. There are a number of other schedulers, but we want to use the approach from the publication. The following class implements the LRS for Transformers. The hidden dimension of the model and the number of warm-up steps serve as hyperparameters. The function get_lr is called with every update of the optimizer and calculates the learning rate of the optimizer.

class TransformerLRS(_LRScheduler):
    """
    Custom Transformer Learning Rate Scheduler
    """
    def __init__(self, 
                 optimizer: Optimizer,
                 dim_embed: int,
                 warmup_steps: int,
                 last_epoch: int=-1,
                 verbose: bool=False) -> None:

        self.dim_embed = dim_embed
        self.warmup_steps = warmup_steps
        self.num_param_groups = len(optimizer.param_groups)

        super().__init__(optimizer, last_epoch, verbose)

    def get_lr(self) -> float:
        lr = self._calc_lr(self._step_count)
        return [lr] * self.num_param_groups

    def _calc_lr(self, step):
        return self.dim_embed**(-0.5) * min(step**(-0.5), step * self.warmup_steps**(-1.5))

Since Weights&Biases is used as a monitoring and logging tool, it is necessary to log in first. An API key and the URL to the server may be required for private instances. Below a dictionary as config for the training run is defined.

wandb.login()

config = {
    "trainer": {
        "epochs": 300,
        "batch_size": 256
    },
    "architecture": "Vanilla Transformer",
    "dataset": "FD001"
}

The actual run is then initialized. Optional metadata can be attached to this run. Within the run environment, the database contained in the “h5” file is loaded and converted into a TensorDataset. It is more efficient to first convert the database[“<dataset>”] datasets into a Numpy array and then into a PyTorch tensor.

with wandb.init(
    project="RUL Prediction",
    job_type="training",
    notes="Training Vanilla Transformer for RUL prediction",
    tags=["baseline", "Vanilla", "RUL"],
    config=config
) as run:
    save_dir = f"{cwd}/{wandb.config['dataset']}/{wandb.config['dataset']}.h5"
    database = h5py.File(save_dir, "r")

    training_set = TensorDataset(
        torch.tensor(np.array(database["X_train"]), dtype=torch.float),
        torch.tensor(np.array(database["Y_train"]), dtype=torch.float)
    )
    validation_set = TensorDataset(
        torch.tensor(np.array(database["X_test"]), dtype=torch.float),
        torch.tensor(np.array(database["Y_test"]), dtype=torch.float)
    )

The DataLoaders are then generated, which wrap an iterable around the dataset. The batch_size is passed as a hyperparameter from the config.

training_loader = DataLoader(
    training_set,
    batch_size=wandb.config["trainer"]["batch_size"],
    shuffle=True,
    num_workers=4
)
validation_loader = DataLoader(
    validation_set,
    batch_size=wandb.config["trainer"]["batch_size"],
    num_workers=4
)

The model is now loaded and the callbacks are defined. When training the model, the callbacks are used to save model checkpoints if a parameter state of the model leads to a minimum value of the validation metric RMSE and to monitor the learning rate per epoch. Weights&Biases is integrated as a logger. It is important to note that the entire model is saved with the log_model flag being true.

model = vanilla_transformer.VanTransLitModule()

callbacks = [
    ModelCheckpoint(
        monitor="val_RMSE",
        mode="min"
    ),
    LearningRateMonitor(logging_interval="epoch")
]

logger = WandbLogger(
    name="Vanilla Transformer",
    checkpoint_name="best_model",
    project="RUL Prediction",
    log_model=True
)

Finally, the actual trainer is created. The callbacks and the logger are passed to this. If a CUDA-capable GPU is installed, then the GPU is used as the accelerator, otherwise the CPU. The keyword “auto” takes over this selection automatically. The training is finally started with trainer.fit().

trainer = Trainer(
    logger=logger,
    callbacks=callbacks,
    accelerator="auto",
    max_epochs=wandb.config["trainer"]["epochs"]
)

trainer.fit(model, training_loader, validation_loader)

wandb.finish()

Infer the trained model to create Run-To-Failure (RTF) plots

The last section deals with the inference of the trained model in order to create an RTF plot. First the RTF dataset is loaded, then the model is downloaded from Weights&Biases, followed by a prediction and finally the result is plotted.

Similar to the training process, you must first log in to Weights&Biases. The database is then loaded. A TensorDataset is also created, which is then passed to a DataLoader.

wandb.login()

save_dir = f"{cwd}/{subdataset}/RTF.h5"
database = h5py.File(save_dir, "r")

rtf_set = TensorDataset(
    torch.tensor(np.array(database["RTF_X"]), dtype=torch.float),
    torch.tensor(np.array(database["RTF_Y"]), dtype=torch.float)
)

rtf_loader = DataLoader(
    rtf_set,
    batch_size=256,
    num_workers=4

A prediction run is now initialized, whereby the best state of the model is downloaded as an artifact from the Weights&Biases server.

config = {
    "model": "best_model:latest"
}

with wandb.init(
    project="RUL Prediction",
    job_type="inference",
    notes="Testing Vanilla Transformer for RUL prediction",
    tags=["baseline", "Vanilla", "RUL"],
    config=config
) as run:
    path = run.use_artifact(wandb.config["model"]).download()

A logger and a trainer are created in the same way as for training. The model is now instantiated from the downloaded checkpoint and the prediction of the trainer is called up.

wandb_logger = WandbLogger(
    project="RUL Prediction",
    job_type="inference",
    notes="Testing Vanilla Transformer for RUL prediction",
    tags=["baseline", "Vanilla", "RUL"]
)

trainer = Trainer(
    logger=wandb_logger,
    accelerator="gpu"
)

model = vanilla_transformer.VanTransLitModule.load_from_checkpoint(os.path.join(path, "model.ckpt"))

predictions = trainer.predict(model=model, dataloaders=rtf_loader)

Since the data for the model processing was normalized to the value range [0, 1], the true and predicted RUL must now be normalized back.

y_true = rtf_loader.dataset.tensors[1].numpy() * RUL_max
y_pred = predictions[0].numpy() * RUL_max

The following code snippet creates a plot of the true and predicted RUL of the Vanilla Transformer. In addition, a confidence interval of ±10% of the maximum RUL is specified and some formatting settings are made.

fig, ax = plt.subplots(figsize=(8, 4))
plt.rcParams["font.family"] = "Times New Roman"

plt.plot(
    y_true,
    color="tab:blue",
    label=f"True RUL for unit {unit}"
)

ci_lower = np.squeeze(y_true - ci*RUL_max)
ci_upper = np.squeeze(y_true + ci*RUL_max)
t = np.arange(len(y_true))
ax.fill_between(t, ci_lower, ci_upper, color='grey', alpha=.3)

plt.plot(
    y_pred,
    color="tab:orange",
    label="Predicted RUL"
)

plt.legend(fontsize=12)
plt.grid()
plt.xlabel("Inspection intervals", fontsize=12)
plt.ylabel("RUL", fontsize=12)
title_str = "Vanilla Transformer for RUL Prediction"
plt.title(title_str, fontsize=12)
fig.tight_layout()

RTF plot for unit 24 of subdataset FD001

Finally, the RTF plot can be displayed. As an example, the RTF plot for unit 24 was generated from the sub-data set FD001. It can be seen that the model approximates the true RUL relatively well in the initial inspection intervals. In the middle range of approximately 40 to 100 inspection intervals, an undershoot and overshoot is recognizable, but this is reduced again as the End-of-Life (EOL) is approached.

Thanks for reading! My name is Dominic, I am working as a Data Scientist with a passion for Generative AI and Predictive Maintenance.