Thank you for your reply.
I have tried the initialization of parameters but still something wrong with my results. The output has a variance of close to 0.
My result:
Note: Left-hand side is the distribution of training data, while the right-hand side is the distribution of outputs from my model. Every row corresponds to one fold in cross validation.
download (1)3471×4425 531 KB
My train function:
from sklearn.model_selection import KFold

EPOCH = 20
batch_size = 128
n_splits = 10

def train(n_splits, X_train, Y_train):
    fold = KFold(n_splits=n_splits, shuffle=False)

    # Score of folds
    valid_loss_per_fold = []
    training_loss_per_fold = []
    train_target_fold, train_output_fold = [], []
    val_target_fold, val_output_fold = [], []
    
    for fold_idx, (train_idx, val_idx) in enumerate(fold.split(X_train)):
        
        X_train_ = [X_train[i] for i in train_idx] 
        Y_train_ = [Y_train[i] for i in train_idx] 
        
        X_val_ = [X_train[i] for i in val_idx]
        Y_val_ = [Y_train[i] for i in val_idx]
        
        train_set = data_loader(X_train_, Y_train_, transform)
        train_loader = DataLoader(train_set, shuffle=True, batch_size=batch_size)
        
        val_set = data_loader(X_val_, Y_val_, transform)
        val_loader = DataLoader(val_set, shuffle=True, batch_size=batch_size)

        train_loss = []
        valid_loss = []
        
        net = RegressionNet(input_size=[batch_size, 3, n_grid, n_grid] , output_size=1)
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        net = net.to(device)

        optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
        criterion = torch.nn.MSELoss()
        
        print('Training for fold {}...'.format(fold_idx + 1))
        
        for epoch in range(EPOCH):
            net.train()
            running_train_loss = 0.0
            with torch.set_grad_enabled(True):
                for inputs, labels in train_loader:
                    inputs = inputs.view(-1, 3, 100, 100)  
                    inputs = inputs.to(device).float()
                    labels = labels.to(device).float()
                    optimizer.zero_grad()
                    outputs = net(inputs)
                    loss = criterion(outputs, labels)
                    running_train_loss += loss.item()
                    loss.backward()
                    optimizer.step()
            train_loss.append(running_train_loss / len(train_set))

            # Validation
            net.eval()
            running_valid_loss = 0.0
            with torch.set_grad_enabled(False):
                for inputs, labels in val_loader:
                    inputs = inputs.view(-1, 3, 100, 100)  
                    inputs = inputs.to(device).float()
                    labels = labels.to(device).float()
                    outputs = net(inputs)
                    running_valid_loss += loss.item()

            valid_loss.append(running_valid_loss / len(val_set))

            print('#epoch:{}\ttrain loss: {}\tvalid loss: {}'.format(epoch, running_train_loss / len(train_loss), running_valid_loss / len(valid_loss)))
        
        net.eval()
        train_target, train_output = [], []
        val_target, val_output = [], []
        with torch.set_grad_enabled(False):
            for inputs, labels in train_loader:
                inputs = inputs.view(-1, 3, 100, 100)  
                inputs = inputs.to(device).float()
                labels = labels.to(device).float()
                outputs = net(inputs)
                train_target.append(labels)
                train_output.append(outputs)
                
            for inputs, labels in train_loader:
                inputs = inputs.view(-1, 3, 100, 100)  
                inputs = inputs.to(device).float()
                labels = labels.to(device).float()
                outputs = net(inputs)
                val_target.append(labels)
                val_output.append(outputs)
        
        flat_train_target = [item.cpu().numpy() for sublist in train_target for item in sublist]
        flat_train_target = np.vstack(flat_train_target)
        flat_train_output = [item.cpu().numpy() for sublist in train_output for item in sublist]
        flat_train_output = np.vstack(flat_train_output)
        flat_val_target = [item.cpu().numpy() for sublist in val_target for item in sublist]
        flat_val_target = np.vstack(flat_val_target)
        flat_val_output = [item.cpu().numpy() for sublist in val_output for item in sublist]
        flat_val_output = np.vstack(flat_val_output)
        train_target_fold.append(flat_train_target)
        train_output_fold.append(flat_train_output)
        val_target_fold.append(flat_val_target)
        val_output_fold.append(flat_val_output)
        valid_loss_per_fold.append(np.mean(valid_loss) / fold.n_splits)
        training_loss_per_fold.append(np.mean(train_loss) / fold.n_splits)
        fold_idx += 1
    

    return training_loss_per_fold, valid_loss_per_fold, train_target_fold, train_output_fold, val_target_fold, val_output_fold

Is there something wrong with my model architecture or train function?