Model returning target average

1

I am having trouble generating adequate results in my regression model. For context, a given data sample’s inputs consists of a dictionary of terms {‘Label1’: [data],‘Label2’: [data]} in which the corresponding labels are fed to specific NN models in a ModuleList with the output of both of those models being the ‘predicted’ value. However, upon training my model merely results in predictions equivalent to the mean of the target values - thus a RMSE equivalent to the standard deviation of the targets. Clearly there is some issue here? A copy of my model and trainer are as follows:

Class Model (nn.Module):
    def __init__(self,labels):
        super(Model,self).__init__()
        self.labels=labels
        self.n_labels=len(labels)
        label_models=nn.ModuleList()
        for n_models in range(n_labels):
            label_models.append(MLP())
        self.label_models=label_models

    def forward(self,data):
        pred=OrderedDict()
        for index,label in enumerate(self.labels):
            model_inputs=data[label][0]
            contribution_index=data[label][1]
            label_outputs=self.label_models[index].forward(model_inputs)
            for index,label_output in enumerate(label_outputs):
                if contribution_index[index] not in pred.keys():
                    pred[contribution_index[index]]=label_output
                else:
                    pred[contribution_index[index]]+=label_output
        output=pred.values()
        output=torch.stack(output)
        return output

            scheduler.step()
            model.train()

            MSE=0.0

            for data_sample in dataloader:
                input_data=data_sample[0]
                target=data_sample[1]
                batch_size=len(target)

                for element in unique_atoms:
                    input_data[element][0]=input_data[element][0].to(device)
                target=target.to(device)
                optimizer.zero_grad()

                with torch.set_grad_enabled(True):
                    output=model(input_data)
                    loss=criterion(output,target)
                    loss.backward()
                    optimizer.step()

                MSE += loss.item()*batch_size

            MSE=MSE/dataset_size
            RMSE=np.sqrt(MSE)
            epoch_loss = RMSE
1 Like