To implement a simple feedforward neural network in PyTorch, we need to define the architecture of the model, the optimizer, and the loss function. Here’s a simple example of how to do this:
import torch
import torch.nn as nn
import torch.optim as optim
# Define the neural network architecture
class NeuralNet(nn.Module):
def __init__(self):
super(NeuralNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# Create an instance of the NeuralNet class
model = NeuralNet()
# Define the optimizer
optimizer = optim.SGD(model.parameters(), lr=0.01)
# Define the loss function
criterion = nn.CrossEntropyLoss()
# Generate some sample data
inputs = torch.randn(3, 10)
labels = torch.tensor([0, 1, 0])
# Train the model
for i in range(100):
# Zero the gradients
optimizer.zero_grad()
# Forward pass
outputs = model(inputs)
# Compute the loss
loss = criterion(outputs, labels)
# Backward pass
loss.backward()
# Update the weights
optimizer.step()
# Print the loss every 10 iterations
if (i+1) % 10 == 0:
print(f"Epoch {i+1}, Loss = {loss.item():.2f}")
In this example, ‘NeuralNet‘ is a subclass of ‘nn.Module‘ that defines the architecture of the neural network. The ‘__init__‘ method defines the layers of the network, and the ‘forward‘ method specifies how the data flows through the network.
We then create an instance of the ‘NeuralNet‘ class and define the optimizer and loss function. In this example, we use stochastic gradient descent (‘optim.SGD‘) as the optimizer and cross-entropy loss (‘nn.CrossEntropyLoss‘) as the loss function.
We generate some sample data, which consists of three input vectors of length 10 and their corresponding labels (0 or 1). We then train the model for 100 epochs using a for loop, where we perform forward and backward passes and update the weights using the optimizer. We print the loss every 10 iterations.
Note that this is just a simple example of how to implement a feedforward neural network in PyTorch. In practice, neural networks can be much more complex and may require different architectures, optimizers, and loss functions.