pytorch простой пользовательский рекуррентный слой очень медленный - PullRequest
Новая
       26

pytorch простой пользовательский рекуррентный слой очень медленный

1 голос
/ 09 ноября 2019

Я реализовал очень простой пользовательский рекуррентный слой в pytorch, используя PackedSequence. Слой замедляет мою сеть в порядке х20. Я читал о замедлении на пользовательских слоях без использования JIT, но в порядке x1.7, с чем я мог бы жить.

Я просто индексирую упакованные последовательности для каждой последовательности и выполняю рекурсию. У меня есть подозрение, что какой-то код не выполняется на GPU?

Я также благодарен за любые другие советы, как реализовать этот тип RNN (по сути, не имеющий плотного слоя, без какого-либо смешения между функциями). 

import torch
import torch.nn as nn
from torch.nn.utils.rnn import PackedSequence

def getPackedSequenceIndices(batch_sizes):
    """input: batch_sizes from PackedSequence object
    requires length-sorted sequences!
    """
    nBatches = batch_sizes[0]
    seqIdx = []
    for ii in range(nBatches):
        seqLen = torch.sum((batch_sizes - ii) > 0).item()
        idx = torch.LongTensor(seqLen)
        idx[0] = ii
        idx[1:] = batch_sizes[0:seqLen-1]
        seqIdx.append( torch.cumsum(idx, dim=0) )
    return seqIdx


class LinearRecursionLayer(nn.Module):
    """Linear recursive smoothing layer with trainable smoothing constants."""

    def __init__(self, feat_dim, alpha_smooth=0.5):
        super(LinearRecursionLayer, self).__init__()
        self.feat_dim = feat_dim
        # trainable parameters
        self.alpha_smooth = nn.Parameter(alpha_smooth*torch.ones(self.feat_dim))
        self.wx = nn.Parameter(torch.ones(self.feat_dim))
        self.activ = nn.Tanh

    def forward(self, x):
        if isinstance(x, PackedSequence):
            seqIdx = getPackedSequenceIndices(x.batch_sizes)
            ydata = torch.zeros_like(x.data)
            for idx in seqIdx:
                y_frame = x.data[idx[0]] # init with first frame
                # iterate over sequence
                for nn in idx:
                    x_frame = x.data[nn]
                    y_frame = self.alpha_smooth*y_frame + (1-self.alpha_smooth)*x_frame # smoothing recurrence
                    ydata[nn,:] = self.activ(self.wx*(y_frame))
            y = PackedSequence(ydata, x.batch_sizes) # pack

        else:  # tensor
            raise ValueError('not implemented')
        return y
...