Search code examples
pythonpytorchlossinferencecustom-training

Train model in Pytorch with custom loss how to set up optimizer and run training?


I am new to pytorch and I am trying to run a github model I found and test it. So the author's provided the model and the loss function.

like this:

#1. Inference the model
model = PhysNet_padding_Encoder_Decoder_MAX(frames=128)
rPPG, x_visual, x_visual3232, x_visual1616 = model(inputs)

#2. Normalized the Predicted rPPG signal and GroundTruth BVP signal
rPPG = (rPPG-torch.mean(rPPG)) /torch.std(rPPG)     # normalize
BVP_label = (BVP_label-torch.mean(BVP_label)) /torch.std(BVP_label)     # normalize

#3. Calculate the loss
loss_ecg = Neg_Pearson(rPPG, BVP_label)

Dataloading

    train_loader = torch.utils.data.DataLoader(train_set, batch_size = 20, shuffle = True)

    batch = next(iter(train_loader))

    data, label1, label2 = batch

    inputs= data

Let's say I want to train this model for 15 epochs. So this is what I have so far: I am trying to set the optimizer and training, but I am not sure how to tie the custom loss and data loading to the model and set the 15 epoch training correctly.

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

for epoch in range(15):
  ....

Any suggestions?


Solution

  • I assumed BVP_label is label1 of train_loader

    train_loader = torch.utils.data.DataLoader(train_set, batch_size = 20, shuffle = True)
    
    # Using GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    model = PhysNet_padding_Encoder_Decoder_MAX(frames=128)
    model.to(device)
    
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    
    for epoch in range(15):
        model.train()
        for inputs, label1, label2 in train_loader:
            rPPG, x_visual, x_visual3232, x_visual1616 = model(inputs)
            BVP_label = label1 # assumed BVP_label is label1
    
            rPPG = (rPPG-torch.mean(rPPG)) /torch.std(rPPG)
            BVP_label = (BVP_label-torch.mean(BVP_label)) /torch.std(BVP_label)
            
            loss_ecg = Neg_Pearson(rPPG, BVP_label)
            
            optimizer.zero_grad()
            loss_ecg.backward()
            optimizer.step()
    

    PyTorch training steps are as belows.

    • Create DataLoader
    • Initialize model and optimizer
    • Create a device object and move model to the device

    in the train loop

    • select a mini-batch of data
    • use the model to make predictions
    • calculate the loss
    • loss.backward() updates the gradients of the model
    • update the parameters using optimizer

    As you may know you can also check PyTorch Tutorials.

    Learning PyTorch with Examples

    What is torch.nn really?