Ajout d'une couche linéaire à un modèle existant sur Pytorch

Question

J'essaie d'ajouter une nouvelle couche à un réseau existant (en tant que première couche) et de l'entraîner sur l'entrée d'origine. Quand j'ajoute un calque convolutif, tout fonctionne parfaitement mais quand je le change en linéaire, il ne semble pas s'entraîner. Des idées pourquoi? Voici l'ensemble du réseau:

class ActorCritic(torch.nn.Module): #original model 
    def __init__(self, num_inputs, action_space): 
     super(ActorCritic, self).__init__() 
     self.conv1 = nn.Conv2d(num_inputs, 32, 3, stride=2, padding=1) 
     self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1) 
     self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1) 
     self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1) 

     self.lstm = nn.LSTMCell(32 * 3 * 3, 256) 

     num_outputs = action_space.n 
     self.critic_linear = nn.Linear(256, 1) 
     self.actor_linear = nn.Linear(256, num_outputs) 

    def forward(self, inputs): 
     inputs, (hx, cx) = inputs 
     x = F.elu(self.conv1(inputs)) 
     x = F.elu(self.conv2(x)) 
     x = F.elu(self.conv3(x)) 
     x = F.elu(self.conv4(x)) 
     x = x.view(-1, 32 * 3 * 3) 
     hx, cx = self.lstm(x, (hx, cx)) 
     x = hx 
     return self.critic_linear(x), self.actor_linear(x), (hx, cx) 

class TLModel(torch.nn.Module): #new model 
    def __init__(self, pretrained_model, num_inputs): 
     super(TLModel, self).__init__() 
     self.new_layer = nn.Linear(1*1*42*42, 1*1*42*42) 
     self.pretrained_model = pretrained_model 

    def forward(self, inputs): 
     inputs, (hx, cx) = inputs 
     x = F.elu(self.new_layer(inputs.view(-1, 1*1*42*42))) 
     return self.pretrained_model((x.view(1,1,42,42), (hx, cx))) 

J'ai essayé différentes fonctions d'activation (non seulement ELU). il fonctionne avec conv:

class TLModel(torch.nn.Module): 
    def __init__(self, pretrained_model, num_inputs): 
     super(TLModel, self).__init__() 
     self.new_layer = nn.Conv2d(num_inputs, num_inputs, 1) 
     self.pretrained_model = pretrained_model 

    def forward(self, inputs): 
     inputs, (hx, cx) = inputs 
     x = F.elu(self.new_layer(inputs)) 
     return self.pretrained_model((x, (hx, cx))) 

Le nombre d'entrées est 1 et la taille d'une entrée est 1x1x42x42

Answer 1

Il serait utile si vous aviez fourni le message d'erreur. De ce que vous avez écrit, je peux seulement deviner que vous avez oublié de presser votre entrée. Vous écrivez que votre entrée est de taille 1x1x42x42, c'est-à-dire qu'elle est en 4 dimensions. nn.Conv2D attend une entrée à 4 dimensions. nn.Linear s'attend plutôt à une entrée bidimensionnelle.

Par conséquent, essayez d'appeler input = input.squeeze() avant de l'envoyer à votre modèle. Cela supprime les dimensions singleton, et rendra votre entrée bidimensionnelle comme il y a deux dimensions singleton.

En note, nn.Linear attend une entrée de dimension batch_size x feat_dim. Une couche linéaire a-t-elle vraiment un sens pour vos données? Comme une autre note de côté, quand on ajoute généralement des couches à des réseaux, ils les mettent à la fin pas le début, mais j'espère que vous avez de bonnes raisons de le faire et de savoir ce que vous faites :)

Bonne chance!

Answer 2

Les dimensions ne devraient pas poser de problème puisque vous changez la vue dans la fonction de transfert. Je crois que l'ajout du FCN dans le démarrage du réseau ne fait pas vraiment beaucoup de différence, vous pouvez vérifier en supprimant la couche et en recyclant le réseau, cela produirait des résultats similaires. Essayez d'imaginer que si vous essayez de détecter les caractéristiques d'un chat (moustaches, oreilles, etc.), vous n'avez pas vraiment besoin d'informations sur tous les pixels de l'image mais seulement sur les pixels voisins, donc une couche entièrement connectée vient compliquer le réseau.