Save labels name when save model

aazz: <p>Hello,<br>
I’m new in image classification.<br>
I start with eleven image type and I would like know how save the labels when the training is finish.<br>
When I predict a picture with this model, I need to have the label name and I don’t want add this on the prediction code:</p>
<pre><code class="lang-auto">    classes = ('Cat', 'Dog', 'Bird', 'Airplane', 'Boat', 'Car', 'moto', 'bicycle', 'wheel',  'pen', 'apple')
    
</code></pre>
<p>Thank you for your help</p>
<pre><code class="lang-auto">import torch
import torchvision
import  torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np

import os 
os.environ['KMP_DUPLICATE_LIB_OK']='True'

if __name__ == "__main__":
    path_training_folder="Image" 
    
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    transform = transforms.Compose(
        [transforms.Resize(255),
         transforms.CenterCrop(255),         
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    
    dataset = torchvision.datasets.ImageFolder(root=path_training_folder,transform=transform)
    
    # Creating data indices for training and validation splits:
    validation_split = 0.3
    shuffle_dataset = True
    random_seed = 217
    dataset_size = len(dataset)
    indices = list(range(dataset_size))
    split = int(np.floor(validation_split * dataset_size))
    if shuffle_dataset :
        np.random.seed(random_seed)
        np.random.shuffle(indices)
    train_indices, val_indices = indices[split:], indices[:split]
    for i in range(7):      # * 2^4
        train_indices = train_indices + train_indices
    
    print("train indices number="+str(len(train_indices)))
    print("valid indices number="+str(len(val_indices)))
    
    # Creating PT data samplers and loaders:
    train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_indices)
    valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_indices)
    
    trainloader = torch.utils.data.DataLoader(dataset, batch_size=4,
                                              shuffle=False, num_workers=2, sampler=train_sampler)
     
    testloader = torch.utils.data.DataLoader(dataset, batch_size=4,
                                             shuffle=False, num_workers=2, sampler=valid_sampler)
    
    classes = ('Cat', 'Dog', 'Bird', 'Airplane', 'Boat', 'Car', 'moto', 'bicycle', 'wheel',  'pen', 'apple')
    
    def imshow(img):
        img = img / 2 + 0.5     # unnormalize
        npimg = img.numpy()
        plt.imshow(np.transpose(npimg, (1, 2, 0)))
        plt.show()
        
    # get some random training images
    dataiter = iter(trainloader)
    images, labels = dataiter.next()

    # print labels
    print('Check: ', ' - '.join('%5s' % classes[labels[j]] for j in range(4)))
       

    # show images
    imshow(torchvision.utils.make_grid(images))

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(3, 4, 5)  ####
            self.pool = nn.MaxPool2d(2, 2)
            self.conv2 = nn.Conv2d(4, 16, 5) ####
            self.fc1 = nn.Linear(57600 , 1024) ####
            self.fc2 = nn.Linear(1024, 700)  ####
            self.fc3 = nn.Linear(700, 11)  #### 11 nb of classes

        def forward(self, x):
            x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
            x = F.max_pool2d(F.relu(self.conv2(x)), 2)
            x = x.view(-1, self.num_flat_features(x))
            #x = x.view(-1, 16 * 61 * 61)
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
            x = self.fc3(x)
            return x

        def num_flat_features(self, x):
            size = x.size()[1:]  # all dimensions except the batch dimension
            num_features = 1
            for s in size:
                num_features *= s
            return num_features

    net = Net()
    net.to(device)

    # Training
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    epochs = 2
    steps = 0
    running_loss = 0
    print_every = 10
    train_losses, test_losses = [], []
    enum_max = 0

    for epoch in range(epochs):  # loop over the dataset multiple times
        running_loss = 0.0
        enum = 0
        enum_max = 0
        for i, data in enumerate(trainloader, 0):
            enum = enum+1
            steps += 1
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data[0].to(device), data[1].to(device)

            # zero the parameter gradients
            optimizer.zero_grad()           
            
            logps = net(inputs)
            loss = criterion(logps, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            if steps % print_every == 0:
              test_loss = 0
              accuracy = 0
              net.eval()
              with torch.no_grad():                
                for inputs, labels in testloader:
                    inputs, labels = inputs.to(device), labels.to(device)
                    logps = net.forward(inputs)
                    batch_loss = criterion(logps, labels)
                    test_loss += batch_loss.item()
                    ps = torch.exp(logps)
                    top_p, top_class = ps.topk(1, dim=1)
                    equals = top_class == labels.view(*top_class.shape)
                    accuracy += torch.mean(equals.type(torch.FloatTensor)).item() 
              train_losses.append(running_loss/len(trainloader))
              test_losses.append(test_loss/len(testloader))                    
              print(f"Epoch {epoch+1}/{epochs}.. "
                    f"Train loss: {running_loss/print_every:.3f}.. "
                    f"Test loss: {test_loss/len(testloader):.3f}.. "
                    f"Test accuracy: {accuracy/len(testloader):.3f}")
              
              if steps % 100 == 0:
                plt.plot(train_losses, label='Training loss')
                plt.plot(test_losses, label='Validation loss')
                plt.legend(frameon=False)
                plt.show()
              if (round((running_loss/print_every),3) == 0.000 and enum_max == 0):
                print(enum)
              running_loss = 0
              net.train()
    PATH = 'trainv2.pth'
    torch.save(net.state_dict(), PATH)
    print('Finished Training')

    # get some random test images
    dataiter = iter(testloader)
    images, labels = dataiter.next()

    # print images and label
    imshow(torchvision.utils.make_grid(images))
    print('GroundTruth: ', ' - '.join('%5s' % classes[labels[j]] for j in range(4)))




      
     
     

</code></pre>

ptrblck: <p>You can use <code>torch.save</code> to save e.g. a <code>dict</code> object, which could contain the <code>model.state_dict()</code>, your class name mapping and any additional objects.</p>