Pytorch重命名已训练模型的标签

问题描述

我有训练有素的模型。在培训期间，字符串字符存在一些问题。所以我将标签转换成数字：

red : 0

blue: 1

green: 2

现在可以将我的标签重命名为实际的标签名称了。 培训时间。如果有人有主意，将会很有帮助。

训练并验证模型

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,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]

但是我不想传递上面的代码。由于我的评估脚本需要通用。

解决方法

为此使用字典！

labels = [0,1,2,1]

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

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

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

,

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

创建标签字典：

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

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

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

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

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

输出：[0,1]

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

输出：['red','blue','green','red','blue']

完整的代码段：

labels = ["red","blue"] 

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

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

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

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