Forward pass with all samples

Question

import torch
import torch.nn as nn

class PINN(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_layers, neurons_per_layer):
        super(PINN, self).__init__()
        layers = []
        layers.append(nn.Linear(input_dim, neurons_per_layer))
        for _ in range(hidden_layers):
            layers.append(nn.Linear(neurons_per_layer, neurons_per_layer))
        layers.append(nn.Linear(neurons_per_layer, output_dim))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

# Example: generating random input data
inputs = torch.rand((1000, 3))  # 3D input coordinates


model = PINN(input_dim=3, output_dim=3, hidden_layers=4, neurons_per_layer=64)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

epochs = 10000
for epoch in range(epochs):
    optimizer.zero_grad()  
    nn_output = model(inputs) # Compute the NN prediction
    # Compute e.g gradient of nn_output
    loss.backward()  
    optimizer.step() 

I want to implement a physics-informed NN where the inputs are N 3d points (x,y,z) and the NN output is a vector-valued quantitiy at this point, that is, both input dimension and output dimension are the same.

To calculate the loss at every epoch, I need to have the value of the quantity at all points. Example: For N=1000points, I need all 1000 NN-predictions before I can proceed with the loss calculation. In my code, I am basically giving a 1000x3 object to the input layer assuming that pytorch passes each row (1x3) separately to the network and at the end organizes it again as an 1000x3object.

Does pytorch work like that or do I have to rethink this approach?

Answer 1

Yes, PyTorch is set up to work with batches of samples exactly as you suggest. You can pass all 1000 samples in and you'll get 1000 samples out of your network, where, internally, each one of those 1000 samples will get passed through the network. Try it!

Take this simple network as an example:

# linear network with 3 input parameters and 3 output parameters
network = torch.nn.Linear(3, 3, bias=False)

# input with 2 samples
inputs = torch.randn(2, 3)

outputs = network(inputs)
print(outputs)
tensor([[0.1277, 0.5881, 0.1048],
        [0.3140, 0.2438, 0.1175]], grad_fn=<MmBackward0>)

# which is equivalent to matrix multiplying the network weights
# by each input sample 
for sample in inputs:
    print(torch.matmul(sample, network.weight.T))
tensor([0.1277, 0.5881, 0.1048], grad_fn=<SqueezeBackward4>)
tensor([0.3140, 0.2438, 0.1175], grad_fn=<SqueezeBackward4>)

In general, in many cases of training a neural network training, you will have a very large number of training samples (hundreds of thousands, millions, etc!), and during each epoch you will want to pass all those samples through the network before calculating the overall loss. Often, that's not practical, and during each epoch you will have to pass through the samples in smaller batches ("minibatches"), calculate the loss on the output of each minibatch as an approximation of the overall loss, and optimize on that. If you only have 1000 samples, which are each only 3 parameters in size, then there's no issue in just passing them all through the network in one go during a training epoch.