Pytorch 重命名經過訓練的 model 的標簽

Question

我有一個訓練有素的 model。 在訓練期間，字符串字符出現了一些問題。 所以我把我的標簽轉換成數字，比如：

red: 0

blue: 1

green: 2

現在可以將我的 label 重命名為實際的 label 名稱。 小時的訓練。 如果有人有想法會很有幫助。

訓練和驗證 model

for epoch in range(1, epoch_num + 1):
    loss_train, acc_train = train(train_loader, model, criterion, optimizer, epoch)
    loss_val, acc_val = validate(val_loader, model, criterion, epoch)
    total_loss_val.append(loss_val)
    total_acc_val.append(acc_val)

測試單個圖像：

def eval_image(file_path):
    model = torch.load(file_path)
    model.eval()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    X = Image.open('red.jpeg')
    test_transforms = transforms.Compose([transforms.ToTensor()])
    image_tensor = test_transforms(X).float()
    image_tensor = image_tensor.unsqueeze_(0)
    input = Variable(image_tensor)
    input = input.to(device)
    output = model(input)
    index = output.data.cpu().numpy().argmax()
    print(index)

訓練完成后，評估腳本不能通用。 我將不得不通過以下代碼

idx_to_class = {
        0: "red",
        1: "blue",
        2: "green",
    }
    class_name = idx_to_class[index]

但我不想通過上面的代碼。 因為我的評估腳本需要通用。

Answer 1

使用字典！

labels = [0, 1, 2, 1, 0, 2, 1]

dct = {0: "red", 1: "blue", 2: "green"}

renamed_labels = [dct[x] for x in labels]

renamed_labels ## ["red", "blue", "green", "blue", "red", "green", "blue"]

Answer 2

labels = ["red", "blue", "green"]
dataset = ["red", "blue", "green", "green", "green", "red", "blue"] 

創建 label 字典：

l_dict = {v:k for k, v in enumerate(labels)}
print(l_dict)

Output： {'red': 0, 'blue': 1, 'green': 2}

l_dict_reverse = {k:v for k, v in enumerate(labels)}
print(l_dict_reverse)

Output： {0: 'red', 1: 'blue', 2: 'green'}

char_to_num = [l_dict[x] for x in dataset]
print(char_to_num) 

Output: [0, 1, 2, 2, 2, 0, 1]

num_to_char = [l_dict_reverse[x] for x in char_to_num]
print(num_to_char)                        

Output： ['red', 'blue', 'green', 'green', 'green', 'red', 'blue']

完整的代碼片段：

labels = ["red", "blue", "green"]
dataset = ["red", "blue", "green", "green", "green", "red", "blue"] 

l_dict = {v:k for k, v in enumerate(labels)}
# Output: {'red': 0, 'blue': 1, 'green': 2}

l_dict_reverse = {k:v for k, v in enumerate(labels)}
# Output: {0: 'red', 1: 'blue', 2: 'green'}

char_to_num = [l_dict[x] for x in dataset]
# Output: `[0, 1, 2, 2, 2, 0, 1]

num_to_char = [l_dict_reverse[x] for x in char_to_num]
# Output: `['red', 'blue', 'green', 'green', 'green', 'red', 'blue']