demo_lstm_test.py

# -*- coding: utf-8 -*-
"""Demo_LSTM_Test.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1ErrkLpB8iMH7p80WwDe9RJsyQOoDCBim

# Part 2: Network Testing

This tutorial demonstrates how to test the LSTM-based model for the hysteresis loop prediction. The network model will be loaded through the saved state dictionary (.sd) file and the prediction results will be saved as (.csv) files.

# Step 0: Import Packages
In this demo, the neural network is synthesized using the PyTorch framework. Please install PyTorch according to the [official guidance](https://pytorch.org/get-started/locally/) , then import PyTorch and other dependent modules.
"""

# Import necessary packages

import torch
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import random
import numpy as np
import json
import math
import os
import re

if "COLAB_TPU_ADDR" in os.environ or "COLAB_GPU" in os.environ:
    print("Running on Google Colab")
    from google.colab import drive

    drive.mount("/content/drive")
    dir = "drive/MyDrive/ColabDrive/magnetchallenge-repo/src/tutorials_replicas/"

else:
    print("Running locally")
    dir = "/Users/tom/Library/CloudStorage/GoogleDrive-twaltermck@gmail.com/My Drive/ColabDrive/magnetchallenge-repo/src/tutorials_replicas/"

print(os.listdir(dir))

"""# Step 1: Define Network Structure
In this part, we define the structure of the LSTM-based encoder-projector-decoder neural network. Refer to the [PyTorch document](https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html) for more details.
"""

# Define model structures and functions

class Encoder(nn.Module):
   def __init__(self, input_dim, hidden_dim):
       super(Encoder, self).__init__()

       self.hidden_dim = hidden_dim
       self.input_dim = input_dim

       self.lstm = nn.LSTM(1, self.hidden_dim, num_layers=1, batch_first=True)

   def forward(self, x):
       outputs, (hidden, cell) = self.lstm(x)
       return hidden, cell

class Decoder(nn.Module):
   def __init__(self, output_dim, hidden_dim):
       super(Decoder, self).__init__()

       self.hidden_dim = hidden_dim
       self.output_dim = output_dim

       self.lstm = nn.LSTM(1, self.hidden_dim, num_layers=1, batch_first=True)
       self.out = nn.Sequential(
            nn.Linear(self.hidden_dim, self.hidden_dim*2),
            nn.Tanh(),
            nn.Linear(self.hidden_dim*2, self.output_dim))

   def forward(self, x, hidden, cell):

       batch = x.shape[0]
       x = x.reshape(batch,1,1)
       output, (hidden, cell) = self.lstm(x, (hidden, cell))

       prediction = self.out(output)
       prediction = prediction.squeeze(0)

       return prediction, hidden, cell

class Projector(nn.Module):
   def __init__(self, num_var, hidden_dim, mod_dim):
       super(Projector, self).__init__()

       self.hidden_dim = hidden_dim
       self.num_var = num_var
       self.mod_dim = mod_dim

       self.out = nn.Sequential(
            nn.Linear(self.hidden_dim + self.num_var, self.mod_dim),
            nn.Tanh(),
            nn.Linear(self.mod_dim, self.mod_dim),
            nn.Tanh(),
            nn.Linear(self.mod_dim, self.hidden_dim))

   def forward(self, x, var1):

       x = x.squeeze(0)
       y = self.out(torch.cat([x,var1],dim=1))
       y = y.unsqueeze(0)

       return y

class Seq2Seq(nn.Module):
   def __init__(self, encoder, projector_hidden, projector_cell, decoder, device):
       super().__init__()

       self.encoder = encoder
       self.projector_hidden = projector_hidden
       self.projector_cell = projector_cell
       self.decoder = decoder
       self.device = device

   def forward(self, source, target, var1, teacher_forcing_ratio=0.5):

       target_len = target.shape[1]
       batch_size = source.shape[0]

       trg_vocab_size = self.decoder.output_dim

       outputs = torch.zeros(target_len+1, batch_size, trg_vocab_size).to(self.device)

       hidden, cell = self.encoder(source)

       hidden = self.projector_hidden(hidden, var1)
       cell = self.projector_cell(hidden, var1)

       trg = torch.add(torch.zeros(batch_size, trg_vocab_size),10).to(self.device)

       for t in range(1, target_len+1):
           prediction, hidden, cell = self.decoder(trg, hidden, cell)

           outputs[t] = prediction.squeeze(2)

           if random.random() < teacher_forcing_ratio:
             trg = target[:,t-1]
           else:
             trg = prediction

       return outputs[1:]

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

"""# Step 2: Load the Dataset
In this part, we load and pre-process the dataset for the network training and testing. In this demo, a small dataset containing sinusoidal waveforms measured with N87 ferrite material under different frequency, temperature, and dc bias conditions is used, which can be downloaded from the [MagNet GitHub](https://github.com/PrincetonUniversity/Magnet) repository under "tutorial".
"""

# Load the dataset

def load_dataset(data_length=128):
    # Load .json Files
    with open(dir+'Dataset/Dataset_sine.json','r') as load_f:
        DATA = json.load(load_f)
    B = DATA['B_Field']
    B = np.array(B)
    Freq = DATA['Frequency']
    Freq = np.log10(Freq)  # logarithm, optional
    Temp = DATA['Temperature']
    Temp = np.array(Temp)
    Hdc = DATA['Hdc']
    Hdc = np.array(Hdc)
    H = DATA['H_Field']
    H = np.array(H)

    # Format data into tensors
    in_B = torch.from_numpy(B).float().view(-1, data_length, 1)
    in_F = torch.from_numpy(Freq).float().view(-1, 1)
    in_T = torch.from_numpy(Temp).float().view(-1, 1)
    in_D = torch.from_numpy(Hdc).float().view(-1, 1)
    out_H = torch.from_numpy(H).float().view(-1, data_length, 1)

    # Normalize
    in_B = (in_B-torch.mean(in_B))/torch.std(in_B)
    in_F = (in_F-torch.mean(in_F))/torch.std(in_F)
    in_T = (in_T-torch.mean(in_T))/torch.std(in_T)
    in_D = (in_D-torch.mean(in_D))/torch.std(in_D)
    out_H = (out_H-torch.mean(out_H))/torch.std(out_H)

    # Save the normalization coefficients for reproducing the output sequences
    # For model deployment, all the coefficients need to be saved.
    normH = [torch.mean(out_H),torch.std(out_H)]

    print(in_B.size())
    print(in_F.size())
    print(in_T.size())
    print(in_D.size())
    print(out_H.size())

    return torch.utils.data.TensorDataset(in_B, in_F, in_T, in_D, out_H), normH

"""# Step 3: Testing the Model
In this part, we program the testing procedure of the network model. The loaded dataset is entirely used as the test set. The pre-trained model is loaded through the state dictionary file (.sd), and the output of the testing part is the (.csv) file containing the predicted sequences.
"""

# Config the model testing

def main():

    # Reproducibility
    random.seed(1)
    np.random.seed(1)
    torch.manual_seed(1)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # Hyperparameters
    NUM_EPOCH = 2000
    BATCH_SIZE = 128
    DECAY_EPOCH = 150
    DECAY_RATIO = 0.9
    LR_INI = 0.004

    # Select default device
    if torch.cuda.is_available():
        device = torch.device("cuda")
        kwargs = {"num_workers": 0, "pin_memory": True, "pin_memory_device": "cuda"}
        print("using GPU")
    else:
        device = torch.device("cpu")
        kwargs = {
            "num_workers": 0,
            "pin_memory": False,
        }
        print("using CPU")


    # Load dataset
    dataset, normH = load_dataset()
    test_loader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, **kwargs)

    # Setup network
    encoder = Encoder(input_dim=100, hidden_dim=32).to(device)
    decoder = Decoder(output_dim=1, hidden_dim=32).to(device)
    projector_hidden = Projector(num_var=3, hidden_dim=32, mod_dim=64).to(device)
    projector_cell = Projector(num_var=3, hidden_dim=32, mod_dim=64).to(device)
    net = Seq2Seq(encoder, projector_hidden, projector_cell, decoder, device).to(device)

    # Load trained parameters
    state_dict = torch.load(dir+'LSTM/Output/Model_LSTM.sd', map_location=device)
    net.load_state_dict(state_dict, strict=True)
    net.eval()
    print("Model is loaded!")

    # Log the number of parameters
    print("Number of parameters: ", count_parameters(net))

    # Test the network
    with torch.no_grad():
        for in_B, in_F, in_T, in_D, out_H in test_loader:

            outputs = net(in_B.to(device),out_H.to(device),torch.cat((in_F.to(device), in_T.to(device), in_D.to(device)), dim=1),teacher_forcing_ratio=0.0)
            outputs = outputs.transpose(0,1)

            # Save results
            with open(dir+"LSTM/Output/pred.csv", "a") as f:
                np.savetxt(f, (outputs*normH[1]+normH[0]).squeeze(2).cpu().numpy())
                f.close()
            with open(dir+"LSTM/Output/meas.csv", "a") as f:
                np.savetxt(f, (out_H*normH[1]+normH[0]).squeeze(2).cpu().numpy())
                f.close()

        print("Testing finished! Results are saved!")


if __name__ == "__main__":
    main()