需要长时间训练的模特

Question

I have added an LSTM layer after a convolution in the VGG-16 model using PyTorch. Overtime, the model learns just fine. However, after adding just one LSTM layer, which consists of 32 LSTM cells, the process of training and evaluating takes about 10x longer.

I added the LSTM layer to a VGG framework as follows

def make_layers(cfg, batch_norm=False):
   # print("Making layers!")
    layers = []
    in_channels = 3
    count=0
    for v in cfg:
        count+=1
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels=v
        if count==5:
            rlstm =RLSTM(v)
            rlstm=rlstm.cuda()
            layers+=[rlstm]

RLSTM is my custom class, which implements RowLSTM, from Google's Pixel RNN paper.

Is this a common issue? Do LSTM layers just take long to train in general?

Answer 1

Yes, since LSTM (and many other RNNs) rely on sequential feeding of information you lose a big portion of parallelization speed ups you generally have with CNNs. There are other types of RNNs you can explore that leverage more parallelizable algorithms but the verdict on their predictive performance compared to LSTM/GRU is still not out