Building Neural Network Prediction Models with PyTorch

Click on the above “Mechanical and Electronic Engineering Technology” to follow us

To build a neural network model based on PyTorch for predicting outputs, you can follow these steps:

Define the Problem: First, clarify whether the problem is a regression problem or a classification problem. If the output is a continuous value, it may be a regression problem; if the output is discrete, it may be a classification problem.
Prepare the Dataset: Ensure that the dataset has been cleaned and preprocessed, including standardization or normalization, and split into training, validation, and testing sets.
Define the Model: Use PyTorch to define a neural network model. For a model with 5 inputs and 1 output, you can define a simple fully connected layer (also known as a linear layer) network.
Define the Loss Function and Optimizer: For regression problems, mean squared error (MSE) is commonly used as the loss function. The optimizer can be Adam or SGD.
Train the Model: Use the training data to train the model while monitoring performance on the validation set to avoid overfitting.
Evaluate the Model: Evaluate the model’s performance on the test set using appropriate evaluation metrics such as MSE or RMSE.
Tuning the Model: Adjust the model structure or hyperparameters as needed to improve model performance.

In this example, the input data is stored in a txt document with 500 rows, each containing 6 data points separated by spaces, with the first 5 as inputs and the last 1 as output. The test data is stored in another txt document with 100 rows, each containing 6 data points separated by spaces, with the first 5 as inputs and the last 1 as output.

The following code implements the creation of a neural network model, reads data for training, and saves the entire model object after training:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Define the neural network model
class Net(nn.Module):
    def __init__(self, input_size, output_size):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, 128)  # First hidden layer
        self.fc2 = nn.Linear(128, 64)         # Second hidden layer
        self.fc3 = nn.Linear(64, output_size)  # Output layer

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Read data and store it in a list
file_path = r'xxx.txt'
data_list = []
with open(file_path, 'r') as file:
    for line in file:
        values = line.strip().split()
        data_list.append(list(map(float, values)))

# Split data into inputs and outputs
input_size = 5  # 5 inputs
output_size = 1  # 1 output
inputs = torch.tensor([data[:-1] for data in data_list], dtype=torch.float32)
outputs = torch.tensor([data[-1] for data in data_list], dtype=torch.float32)

# Create data loader
batch_size = 10
dataset = TensorDataset(inputs, outputs)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Instantiate the model
model = Net(input_size, output_size)

# Define loss function and optimizer
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Train the model
num_epochs = 10000  # Train for 100 epochs
for epoch in range(num_epochs):
    model.train()  # Set model to training mode
    running_loss = 0.0
    for inputs_batch, outputs_batch in dataloader:
        optimizer.zero_grad()  # Clear previous gradients
        outputs_pred = model(inputs_batch)  # Forward pass
        loss = criterion(outputs_pred, outputs_batch.view(-1, 1))  # Calculate loss
        loss.backward()  # Backward pass
        optimizer.step()  # Update weights
        running_loss += loss.item()
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(dataloader)}')

# Evaluate the model (optional)
model.eval()  # Set model to evaluation mode
with torch.no_grad():
    for inputs_batch, outputs_batch in dataloader:
        outputs_pred = model(inputs_batch)
        loss = criterion(outputs_pred, outputs_batch.view(-1, 1))
        print(f'Test Loss: {loss.item()}')
model_path = r'xxx.pth'
# Save the entire model object
torch.save(model, model_path)

The following code implements saving the entire model object and using test data for predictions, saving the prediction results to a txt document.

import torch
import torch.nn as nn
import numpy as np

# Define the neural network model
class Net(nn.Module):
    def __init__(self, input_size, output_size):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, 128)  # First hidden layer
        self.fc2 = nn.Linear(128, 64)         # Second hidden layer
        self.fc3 = nn.Linear(64, output_size)  # Output layer

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Load the saved model weights
model_path = r'xxx.pth'
model = torch.load(model_path)
# Set model to evaluation mode
model.eval()

# Read test data
file_path = r'xxx.txt'
test_data_list = []
with open(file_path, 'r') as file:
    for line in file:
        values = line.strip().split()
        test_data_list.append(list(map(float, values)))

# Split test data into inputs and outputs
test_inputs = torch.tensor([data[:-1] for data in test_data_list], dtype=torch.float32)
test_outputs_actual = torch.tensor([data[-1] for data in test_data_list], dtype=torch.float32)

# Optional: Create a TensorDataset and DataLoader for batch processing
# test_dataset = TensorDataset(test_inputs, test_outputs_actual)
# test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# Make predictions
predictions = []
with torch.no_grad():
    # If using DataLoader
    # for inputs_batch in test_dataloader:
    #     outputs_pred = model(inputs_batch)
    #     predictions.extend(outputs_pred.cpu().numpy())
        # If not using DataLoader, predict all data at once
    outputs_pred = model(test_inputs)
    predictions.extend(outputs_pred.cpu().numpy())

# Print prediction results
print(predictions)
# Use np.savetxt to save the array to a text file
np.savetxt('output.txt', predictions, fmt='%f', delimiter=' ')
# If needed, calculate model performance metrics, such as MSE
mse_loss = torch.nn.MSELoss()
test_loss = mse_loss(outputs_pred, test_outputs_actual.view(-1, 1))
print(f'Test MSE Loss: {test_loss.item()}')

Building Neural Network Prediction Models with PyTorch

Want to learn more

Quickly scan the code to follow

Leave a Comment Cancel reply