Pytorch Neural Network Trapped in Local Minimum

Question

I apologise in advance for the maths formatting - I do not have the reputation to post equations.

I have been using pytorch to build a neural network to learn the function,

f(x,y,t)=-x.10^y.cos(t)

but so far within a short number (~10) epochs the weights and biases all drop to 0 and never change from there. I believe this is because the network is stuck in a local minimum. The network is structured as:

    def __init__(self,learning_rate,beta_params):
    super(PINN, self).__init__()
    self.E_c=1.0

    self.lin1 = nn.Linear(3,20)
    self.lin2 = nn.Linear(20,20)
    self.lin3 = nn.Linear(20,20)
    self.lin4 = nn.Linear(20,20)
    self.lin5 = nn.Linear(20,1)

    self.optimizer = tr.optim.Adam(self.parameters(),lr=0.01)
    self.LossCrit = tr.nn.MSELoss()
    self.scaler=tr.cuda.amp.GradScaler()

with a forward function defined as:

    def forward(self, x):
    x = tr.sin(self.lin1(x))
    x = tr.tanh(self.lin2(x))
    x = tr.tanh(self.lin3(x))
    x = tr.tanh(self.lin4(x))
    x = self.lin5(x)

    return x

And the training cycle is defined as:

    def InitCycle(self, lb, ub, nTrain, batch_size):
    itNum = []
    itLoss = []
    for it in range(nTrain):
        #Choose a random selection of batch_size points from between lb and ub
        if it % 10 == 0:
            X_f = tr.rand((batch_size, 3)).cuda() * (ub - lb) + lb
            x = X_f[:,0:1]
            B_0 = tr.pow(10.0,X_f[:,2:3])
            t = X_f[:,1:2]
            
            #define a target
            f_Target = (- x * tr.cos(t))

        #zero the gradients and evaluate the prediction of the neural network
        self.optimizer.zero_grad()
        f_pred = self.forward(X_f)

        #define an 'output' which divides the neural network output by B_0 to keep all values close to unity and calculate loss and scaled gradients.
        output = f_pred / B_0
        loss = self.LossCrit(output,f_Target)
        self.scaler.scale(loss).backward()
       
        #record loss every 10 iterations
        if it % 10 == 0:
            if it > 0:
                itNum.append(np.log10(float(it)))
                itLoss.append(np.log10(loss.item()))
            print('loss at it. %d:  %.8e, lr at it. %d: %.2e' % (
            it, loss.item(), it, self.optimizer.param_groups[0]['lr']))
        
        #update the weights and biases
        self.scaler.step(self.optimizer)
        self.scaler.update()
     
    #save the loss file at the end of the training run
    LFile = pd.DataFrame({'it#': itNum, 'Loss': itLoss})
    LFile.to_csv("PATH.csv")

I have tried every pytorch optimizer with various learning rates and hyperparameters but they all get stuck in the local minimum at {W,b}=0. I have also tried a different loss function which compares the reciprocals (with an added eps for numerical stability) but this also gets stuck at {W,b}=0.

The reason I am using gradient scaling is because with the default pytorch initialization my initial gradients are of the order 1E8 without scaling. I have tried using gradient clipping instead.

I've also tried to train a slightly different network to learn the function:

f(x,y,t)=-x.y.cos(t)

but this didn't work either so I stuck with the original form since I want the network to cover 6 orders of magnitude in 10^y and wanted to keep the values of the input y close to unity (y=-3->3).

I chose the sinusoidal activation function for the first layer in an effort to make it easier for the network to learn the cos(t) term but using a tanh activation function does not work either.

I am now out of ideas for how to get this to train and was hoping someone might be able to shed some light on the situation.

Pytorch Neural Network Trapped in Local Minimum

0 Answers