How to apply a CNN from PyTorch to your images. | by Alexey Khrustalev | Towards Data Science

Hello everyone! Today I’d like to talk about uploading your images into your PyTorch CNN.
Today we will cover the following:
How to store images properly, so that you can easily get your data labeled.
How to access the data using PyTorch and make some preprocessing on the way.
In the end, we will use a very simple CNN to classify our images.
I will use this data set, which contains the images of cats and dogs.
The easiest way to store your images is to create a folder for each class, naming the folder with the name of the class.
Let me show an example:
By simply naming your folders properly, you’ll let PyTorch know to which class to assert an image.
import matplotlib.pyplot as pltfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoaderimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optimimport torchdef get_data():    data_dir = '/your_path/Data/'    train_set = datasets.ImageFolder(data_dir + '/training_set')    test_set = datasets.ImageFolder(data_dir + '/test_set')    return train_set, test_set
The function above gets the data from the directory. You have to simply specify the path to your Train Set and Test Set folders. PyTorch will then automatically assign the labels to images, using the names of the folders in the specified directory.
However, you might want to make some preprocessing before using the images, so let’s do it and, furthermore, let’s create a DataLoader right away. To do so, let’s add some new lines to the code above.
def get_data():    data_dir = '/your_path/Data/'       transform = transforms.Compose([        #transforms.RandomRotation(20),        transforms.RandomResizedCrop(128),        #transforms.RandomHorizontalFlip(),        transforms.ToTensor()])    train_set = datasets.ImageFolder(data_dir + '/training_set', transform=transform)    test_set = datasets.ImageFolder(data_dir + '/test_set', transform=transform)    train = DataLoader(train_set, batch_size=32, shuffle=True)    test = DataLoader(test_set, batch_size=32, shuffle=True)    return train, test
As the transformations, you may want to crop, flip, resize, rotate, etc the images. To do so you can specify all the things you need in one place, using transforms. Then you can create the DataLoader with your images uploaded.
Now, when you have the data ready, you might want to take a quick look at it. To do so, you can use this simple function, which will show the first 5 images.
def train_imshow():    classes = ('cat', 'dog') # Defining the classes we have    dataiter = iter(train)    images, labels = dataiter.next()    fig, axes = plt.subplots(figsize=(10, 4), ncols=5)    for i in range(5):        ax = axes[i]        ax.imshow(images[i].permute(1, 2, 0))         ax.title.set_text(' '.join('%5s' % classes[labels[i]]))    plt.show()
Let’s define our Convolutional Neural Network:
class Net(nn.Module):    def __init__(self):        super(Net, self).__init__()        self.conv1 = nn.Conv2d(3, 6, 5)        self.pool = nn.MaxPool2d(8, 8)        self.conv2 = nn.Conv2d(6, 16, 5)        self.fc1 = nn.Linear(16, 100)        self.fc2 = nn.Linear(100, 50)        self.fc3 = nn.Linear(50, 2)    def forward(self, x):        x = self.pool(F.relu(self.conv1(x)))        x = self.pool(F.relu(self.conv2(x)))        x = x.view(-1, 16)        x = F.relu(self.fc1(x))        x = F.relu(self.fc2(x))        x = self.fc3(x)        return x
I used the network defined here and modified it a bit so that now it can work with my images.
Now we can specify the criterion, optimizer, learning rate, and train the network.
criterion = nn.CrossEntropyLoss()optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)if torch.cuda.is_available(): # Checking if we can use GPU    model = net.cuda()    criterion = criterion.cuda()def train_net(n_epoch): # Training our network    losses = []    for epoch in range(n_epoch):  # loop over the dataset multiple times        running_loss = 0.0        for i, data in enumerate(train, 0):            # get the inputs; data is a list of [inputs, labels]            inputs, labels = data            # zero the parameter gradients            optimizer.zero_grad()            # forward + backward + optimize            outputs = net(inputs)            loss = criterion(outputs, labels)            loss.backward()            optimizer.step()            # print statistics            losses.append(loss)            running_loss += loss.item()            if i % 100 == 99:  # print every 2000 mini-batches                print('[%d, %5d] loss: %.10f' %                      (epoch + 1, i + 1, running_loss / 2000))                running_loss = 0.0    plt.plot(losses, label='Training loss')    plt.show()    print('Finished Training')
You can save and load your trained network.
PATH = './cat_dog_net.pth'torch.save(net.state_dict(), PATH)# Loading the trained networknet.load_state_dict(torch.load(PATH))
Let’s feed our test images into the network.
correct = 0total = 0with torch.no_grad():    for data in test:        images, labels = data        outputs = net(images)        _, predicted = torch.max(outputs.data, 1)        total += labels.size(0)        correct += (predicted == labels).sum().item()print('Accuracy of the network on the %d test images: %d %%' % (len(test),    100 * correct / total))
The test set error is 66%, which is not bad, though getting the best possible accuracy is not the goal of this article.
Here you can find the complete code.
Well done! Now you are ready to practice in deep learning using your images!
Thank you for reading!