Определите нейронную сеть с помощью PyTorch для обучения изображений с метками

Question

У меня есть 3 папки (поезд, проверка и тест), где папка содержит 1000, 100 и 100 изображений соответственно с размерами 128x128 и файл label.txt, который содержит имя изображения, центральные точки (как x, y) и радиус кругов для всех изображений в каталоге "images". Вот несколько первых примеров деталей в файле label.txt:

0000,   0.68 ,   0.61 ,   0.24 
0001,   0.7  ,   0.33 ,   0.26 
0002,   0.52 ,   0.61 ,   0.22 
0003,   0.5  ,   0.49 ,   0.46 
0004,   0.43 ,   0.43 ,   0.36 
0005,   0.36 ,   0.38 ,   0.19 
0006,   0.4  ,   0.61 ,   0.3  
0007,   0.81 ,   0.5  ,   0.15 
0008,   0.7  ,   0.36 ,   0.26 
0009,   0.29 ,   0.56 ,   0.2  

и this - пример того, как выглядит изображение (Первое тренировочное изображение в наборе обучающих данных ).

Изображение подается в нейронную сеть, выходные данные которой являются центром и радиусом круга внутри изображения (изображение содержит прямоугольник, треугольник и круг). Рассматриваемая область - «[0,1] x [0,1]« прямоугольник ». Цель состоит в том, чтобы нарисовать круг. Изображения являются RGB, поэтому я знаю, что форма прочитанного изображения (3 128 128). I Я создал класс набора данных и загрузчик данных, выполнив некоторые исследования, но поскольку я новичок в pytorch, я не уверен, как определить нейронную сеть. Код выглядит следующим образом:

Класс набора данных

class ShapesDataset(Dataset):

    def __init__(self, dataset_dir):
        """
        Initializing dataset by generating a dicitonary of labels, where an image file name is the key 
        and its labels are the contents of that entry in the dictionary. Images are not loaded. This way it
        is possible to iterate over arbitrarily large datasets (limited by labels dicitonary fitting 
        in memory, which is not a problem in practice)

        Args:
            dataset_dir : path to directory with images and labels. In this directory we expect to find
                          a directory called "images" containing the input images, and a file called 
                          "labels.txt" containing desired labels (coefficients)
        """

        self.dataset_dir = dataset_dir
        self.labels_dict = self.gen_labels_dict()
        self.images_keys = list(self.labels_dict)  # getting the keys of the dictionary as list
        self.images_keys.sort()                    # sorting so as to have in alphabetical order 

    def __len__(self):
        return len(self.labels_dict)

    def __getitem__(self, index):    
        """
        This funtion makes it possible to iterate over the ShapesDataset
        Args:
            index: running index of images

        Returns:
            sample: a dicitionary with three entries:
                    1. 'image'  contains the image
                    2. 'labels' contains labels (coeffs) corresponding to image
                    3. 'fname'  contains name of file (image_key) - may be useful for debugging
        """
        image_key = self.images_keys[index]     # recall - key is the file name of the corresponding image
        image = np.array(Image.open(image_key)) # image has shape: (128, 128, 3)
        image = image/255.0                     # simple normalization - just to maintain small numbers
        image = np.transpose(image, (2, 0, 1))  # network needs RGB channels to be first index
        labels = self.labels_dict[image_key]
        sample = {'image': image, 'labels': labels, 'fname':image_key}

        return sample


    def gen_labels_dict(self):
        """
        This fucntion generates a dictionary of labels

        Returns:
            labels_dict: the key is image file name and the value is the corresponding 
            array of labels  
        """

        labels_fname = self.dataset_dir + "/labels.txt"
        labels_dict = {}
        with open(labels_fname, "r") as inp:
            for line in inp:
                line = line.split('\n')[0]                                      # remove '\n' from end of line 
                line = line.split(',')
                key  = self.dataset_dir + '/images/' + line[0].strip() + ".png" # image file name is the key
                del line[0]

                list_from_line = [float(item) for item in line]
                labels_dict[key] = np.asarray(list_from_line, dtype=np.float32)

        return labels_dict             

Загрузчик данных

train_dir      = "./train/"
validation_dir = "./validation/"
test_dir       = "./test/"


train_dataset = ShapesDataset(train_dir)

train_loader = DataLoader(train_dataset, 
                          batch_size=32,
                          shuffle=True)



validation_dataset = ShapesDataset(validation_dir)

validation_loader = DataLoader(validation_dataset, 
                               batch_size=1,
                               shuffle=False)



test_dataset = ShapesDataset(test_dir)

test_loader = DataLoader(test_dataset, 
                          batch_size=1,
                          shuffle=False)


print("train loader examples     :", len(train_dataset)) 
print("validation loader examples:", len(validation_dataset))
print("test loader examples      :", len(test_dataset))

#Output:
#train loader examples     : 1000
#validation loader examples: 100
#test loader examples      : 100

И, наконец, определение нейронной сети:

* * РЕДАКТИРОВАТЬ: определение NN **

class CircleNet(nn.Module):    # nn.Module is parent class  
    #Not sure if I am on the right track.
    def __init__(self):
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)

        # We'll apply max pooling with a kernel size of 2
        self.pool = nn.MaxPool2d(kernel_size=2)

        # A second convolutional layer takes 12 input channels, and generates 12 outputs
        self.conv2 = nn.Conv2d(in_channels=12, out_channels=12, kernel_size=3, stride=1, padding=1)

        # A third convolutional layer takes 12 inputs and generates 24 outputs
       self.conv3 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=3, stride=1, padding=1)

        # A drop layer deletes 20% of the features to help prevent overfitting
        self.drop = nn.Dropout2d(p=0.2)

        # Our 128x128 image tensors will be pooled twice with a kernel size of 2. 128/2/2 is 32.
        # So our feature tensors are now 32 x 32, and we've generated 24 of them, so the array is 32x32x24
        # We need to flatten these and feed them to a fully-connected layer
        # to map them to  the probability for each class
        self.fc = nn.Linear(in_features=32 * 32 * 24, out_features=num_classes) 

    #I am struggling to move forward from here            
    def forward(self, x):
        """
        Feed forward through network
        Args:
            x - input to the network

        Returns "out", which is the network's output
        """
        # Use a relu activation function after convolution 1 and pool
        x = F.relu(self.pool(self.conv1(x)))

        # Use a relu activation function after convolution 2 and pool
        x = F.relu(self.pool(self.conv2(x)))

        # Select some features to drop after the 3rd convolution to prevent overfitting
        x = F.relu(self.drop(self.conv3(x)))

        # Only drop the features if this is a training pass
        x = F.dropout(x, training=self.training)

        # Flatten
        x = x.view(-1, 32 * 32 * 24)
        # Feed to fully-connected layer to predict class
        x = self.fc(x)
        # Return class probabilities via a softmax function 

        return F.log_softmax(x, dim=1)