Pytorch ：OneHot encoding 标签转换 代码

import torch
import numpy as np
import pandas as pd

def convertLabel(datasetLabel):
    """
    Labels (product ratings) from the dataset are provided to you as
    floats, taking the values 1.0, 2.0, 3.0, 4.0, or 5.0.
    You may wish to train with these as they are, or you you may wish
    to convert them to another representation in this function.
    Consider regression vs classification.
    """
    # Convert to longTensor (i.e. int)
    convertedLabel = datasetLabel.type(torch.LongTensor)
    # Subtract 1 for classifier target values 0-4
    convertedLabel = torch.add(convertedLabel, -1)
    # Create Tensor as one-hot encoding (makes it easier to do MSE)
    oneHot = torch.zeros(convertedLabel.shape[0], 5)
    for index in range(oneHot.shape[0]):
        oneHot[index, convertedLabel[index]] = 1

    return oneHot  # one hot encoding

if __name__ == '__main__':
    data = np.array(pd.read_csv(r"D:\OCdata\HDI2.csv"))
    x = data[:, :-1]
    y = data[:, -1]
    x = torch.from_numpy(x).type(torch.FloatTensor)
    y = torch.from_numpy(y).type(torch.FloatTensor)
    y_OneHot = convertLabel(y)
    print(y_OneHot)