Que fait la fonction de collecte dans pytorch en termes simples?
J'ai parcouru le document officiel et this mais il est difficile de comprendre ce qui se passe.
J'essaie de comprendre un code source DQN et il utilise la fonction de collecte de la ligne 197.
Quelqu'un pourrait-il expliquer en termes simples ce que la fonction de collecte fait? Quel est le but de cette fonction?
Le torch.gather fonction (ou torch.Tensor.gather) est une méthode de sélection multi-index. Regardez l'exemple suivant de la documentation officielle:
t = torch.tensor([[1,2],[3,4]])
r = torch.gather(t, 1, torch.tensor([[0,0],[1,0]]))
# r now holds:
# tensor([[ 1, 1],
# [ 4, 3]])
Commençons par la sémantique des différents arguments: Le premier argument, input, est le tenseur source à partir duquel nous souhaitons sélectionner des éléments. Le second, dim, est la dimension (ou l’axe dans tensorflow/numpy) que nous souhaitons collecter. Et enfin, index sont les index pour indexer input. En ce qui concerne la sémantique de l'opération, voici comment la documentation officielle l'explique:
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
Alors passons en exemple.
le tenseur d'entrée est [[1, 2], [3, 4]], et l’argument dim est 1, c’est-à-dire que nous souhaitons collecter à partir de la deuxième dimension. Les indices pour la deuxième dimension sont donnés sous la forme [0, 0] et [1, 0].
En "sautant" la première dimension (la dimension que nous voulons collecter est 1), la première dimension du résultat est implicitement donnée en tant que première dimension de index. Cela signifie que les index contiennent la deuxième dimension, ou les index de colonne, mais pas les index de ligne. Celles-ci sont données par les indices du tenseur index lui-même. Pour l'exemple, cela signifie que la sortie aura dans sa première ligne une sélection des éléments de la première ligne du tensorisateur input, comme indiqué par la première rangée du premier tensionneur index rangée. Comme les index-colonnes sont donnés par [0, 0], nous sélectionnons donc deux fois le premier élément de la première ligne de l’entrée, ce qui a pour résultat [1, 1]. De même, les éléments de la deuxième ligne du résultat résultent de l'indexation de la deuxième ligne du tenseur input par les éléments de la deuxième ligne du tenseur index, ce qui donne [4, 3].
Pour illustrer cela davantage, échangeons la dimension dans l'exemple:
t = torch.tensor([[1,2],[3,4]])
r = torch.gather(t, 0, torch.tensor([[0,0],[1,0]]))
# r now holds:
# tensor([[ 1, 2],
# [ 3, 2]])
Comme vous pouvez le constater, les indices sont maintenant collectés le long de la première dimension.
Pour l'exemple que vous avez cité,
current_Q_values = Q(obs_batch).gather(1, act_batch.unsqueeze(1))
gather indexera les lignes des valeurs q (c'est-à-dire les valeurs q par échantillon d'un lot de valeurs q) par la liste d'actions. Le résultat sera le même que si vous aviez fait ce qui suit (bien que ce soit beaucoup plus rapide qu'une boucle):
q_vals = []
for qv, ac in Zip(Q(obs_batch), act_batch):
q_vals.append(qv[ac])
q_vals = torch.cat(q_vals, dim=0)
torch.gather crée un nouveau tenseur à partir du tenseur en prenant les valeurs de chaque ligne le long de la dimension en entrée dim. Les valeurs dans torch.LongTensor, passé sous la forme index, spécifie la valeur à prendre dans chaque "ligne". La dimension du tenseur de sortie est identique à celle du tenseur d'indice. L'illustration suivante des documents officiels l'explique plus clairement: 
(Remarque: dans l'illustration, l'indexation commence à 1 et non à 0).
Dans le premier exemple, la dimension donnée est le long des lignes (de haut en bas). Ainsi, pour (1,1) la position de result, elle prend la valeur de la ligne du index pour le src C'est 1. À (1,1) dans la valeur source est 1 donc, sorties 1 à (1,1) dans result. De même pour (2,2) la valeur de la ligne de l'index pour src est 3. En (3,2) la valeur dans src est 8 et donc les sorties 8 etc.
De même pour le deuxième exemple, l'indexation se fait le long des colonnes, et donc à la position (2,2) de result, la valeur de la colonne de l'index pour src est 3, donc à (2,3) de src, 6 est pris et sorti en result à (2,2)