pytorch, Using nn.DataParallel in LSTM

Question

/pytorch/aten/src/ATen/native/cudnn/RNN.cpp:1266: UserWarning: RNN module weights are not part of single contiguous chunk of memory. 
This means they need to be compacted at every call, possibly greatly increasing memory usage. To compact weights again call flatten_parameters().

Hello. I am using pytorch. I am trying to use DataParallel function in pytorch, but the model is LSTM. I'm warned to flatten the model again, but I don't know when and where to flatten. Can you let me know?

This is my model

    import torch.nn as nn
from torchvision import models

class ConvLstm(nn.Module):
    def __init__(self, latent_dim, model, hidden_size, lstm_layers, bidirectional, n_class):
        super(ConvLstm, self).__init__()
        self.conv_model = Pretrained_conv(latent_dim, model)
        self.Lstm = Lstm(latent_dim, hidden_size, lstm_layers, bidirectional)
        self.output_layer = nn.Sequential(
            nn.Linear(2 * hidden_size if bidirectional ==
                    True else hidden_size, n_class),
            nn.Softmax(dim=-1)
        )

    def forward(self, x):
        batch_size, timesteps, channel_x, h_x, w_x = x.shape
        conv_input = x.view(batch_size * timesteps, channel_x, h_x, w_x)
        conv_output = self.conv_model(conv_input)
        lstm_input = conv_output.view(batch_size, timesteps, -1)
        lstm_output = self.Lstm(lstm_input)
        lstm_output = lstm_output[:, -1, :]
        output = self.output_layer(lstm_output)
        return output

class Pretrained_conv(nn.Module):
    def __init__(self, latent_dim, model):
        if model == 'resnet152':
            super(Pretrained_conv, self).__init__()
            self.conv_model = models.resnet152(pretrained=True)
            # ====== freezing all of the layers ======
            for param in self.conv_model.parameters():
                param.requires_grad = False
            # ====== changing the last FC layer to an output with the size we need. this layer is un freezed ======
            self.conv_model.fc = nn.Linear(
                self.conv_model.fc.in_features, latent_dim)

    def forward(self, x):
        return self.conv_model(x)


class Lstm(nn.Module):
    def __init__(self, latent_dim, hidden_size, lstm_layers, bidirectional):
        super(Lstm, self).__init__()
        self.Lstm = nn.LSTM(latent_dim, hidden_size=hidden_size,
                            num_layers=lstm_layers, batch_first=True, bidirectional=bidirectional)
        self.hidden_state = None

    def reset_hidden_state(self):
        self.hidden_state = None

    def forward(self, x):
        output, self.hidden_state = self.Lstm(x, self.hidden_state)
        return output

Enter LSTM and execute the following code.

def foward_step(model, images, labels, criterion, mode=''):

model.module.Lstm.reset_hidden_state()

if mode == 'test':

    with torch.no_grad():

        output = model(images)

else:

    output = model(images)

loss = criterion(output, labels)

# Accuracy calculation

predicted_labels = output.detach().argmax(dim=1)

acc = (predicted_labels == labels).cpu().numpy().sum()

return loss, acc, predicted_labels.cpu()

This is main

model = nn.DataParallel(model, device_ids=[0,1,2,3]).cuda()

pytorch, Using nn.DataParallel in LSTM

0 Answers