Как спроектировать последний слой моей нейронной сети? - PullRequest
Новая
       9

Как спроектировать последний слой моей нейронной сети?

0 голосов
/ 10 июля 2020

У меня есть нейронная сеть с таким дизайном:

class CONTROLLER (nn.Module): 

def __init__(self, numb_features ,k):
    super(CONTROLLER, self).__init__()
    self.fc1 = nn.Linear(numb_features, 100)
    self.fc2 = nn.Linear(100, 100)
    self.fc3 = nn.Linear(100,k)
    self.relu = nn.ReLU()
    self.softmax = nn.Softmax(k)

    for m in self.modules():
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight)
           
def forward(self, x):
    #x = x.view(-1, 1 * 28 * 28)
    x = self.fc1(x)
    x = self.relu(x)
    x = self.fc2(x)
    x = self.relu(x)
    x = self.fc3(x)
    k_probabilities = self.softmax(x)
    
    return k_probabilities

Я хочу накормить эту сеть mini_batch_size = 100. Выход k = 20 - это 20 вероятностей, и мне нужно значение индекса + 1 наивысшей вероятности. Я получаю это значение для каждого входного изображения, и для 100 изображений у меня должно быть 100 таких значений. Тогда я получаю среднее. Я хотел бы знать, могу ли я встроить этот шаг в сеть в качестве слоя? Я получаю матрицу 20 100, могу ли я получить матрицу 20 1, которая дает мне самую высокую вероятность для всех 100 изображений в мини-пакете?

...