Network Training is too slow for Custom Network even though network is not too large(Resnet + LSTM type network) - neural-network

Question

I have made a custom network for the image data. But the training on this network is very slow though the network is not too huge.
When I am training on resnet150 with lstm, the training is quite fast. Not sure which operation is lagging my training speed?
I have tried reducing Batch_size, seq_dim and network parameters. My guess is some operation is reducing the speed
The dataset I am using are images and they are also very small size(96*96) gray scale images.
***** CODE *****
class Residual(nn.Module):
def __init__(self, input_channels, num_channels,
use_1x1conv=True, strides=1, dilation=2, padding=1, kernel_size=5):
super(Residual, self).__init__()
self.conv1 = nn.Conv2d(input_channels, num_channels,
kernel_size=kernel_size, padding=padding, stride=strides, dilation=dilation)
self.conv2 = nn.Conv2d(num_channels, num_channels,
kernel_size=kernel_size, padding=2 * padding, dilation=(2 * dilation))
if use_1x1conv:
self.conv3 = nn.Conv2d(input_channels, num_channels,
kernel_size=1, stride=strides)
else:
self.conv3 = None
self.bn1 = nn.BatchNorm2d(num_channels)
self.bn2 = nn.BatchNorm2d(num_channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, X):
Y = F.relu(self.bn1(self.conv1(X)))
Y = self.bn2(self.conv2(Y))
if self.conv3:
X = self.conv3(X)
Y += X
return F.relu(Y)
def weight_init(m):
'''
Usage:
model = Model()
model.apply(weight_init)
'''
if isinstance(m, nn.Conv1d):
init.normal_(m.weight.data)
if m.bias is not None:
init.normal_(m.bias.data)
elif isinstance(m, nn.Conv2d):
init.xavier_normal_(m.weight.data)
if m.bias is not None:
init.normal_(m.bias.data)
elif isinstance(m, nn.BatchNorm1d):
init.normal_(m.weight.data, mean=1, std=0.02)
init.constant_(m.bias.data, 0)
elif isinstance(m, nn.BatchNorm2d):
init.normal_(m.weight.data, mean=1, std=0.02)
init.constant_(m.bias.data, 0)
elif isinstance(m, nn.Linear):
init.xavier_normal_(m.weight.data)
init.normal_(m.bias.data)
elif isinstance(m, nn.LSTM):
for param in m.parameters():
if len(param.shape) >= 2:
init.orthogonal_(param.data)
else:
init.normal_(param.data)
class Network(nn.Module):
def __init__(self):
super(Network, self).__init__()
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_channels = 64
strides = 1
dilation = 2
padding = 4
kernel_size = 5
input_channel = 3
self.fc = nn.Linear(config['hidden_dim'], config['output_dim'])
self.lstm = None
b1 = Residual(input_channels=input_channel, num_channels=num_channels,
use_1x1conv=True, strides=strides, dilation=dilation, padding=padding, kernel_size=kernel_size)
b2 = Residual(input_channels=num_channels, num_channels=2 * num_channels,
use_1x1conv=True, strides=strides, dilation=2 * dilation, padding=2 * padding,
kernel_size=kernel_size)
b3 = Residual(input_channels=2 * num_channels, num_channels=4 * num_channels,
use_1x1conv=True, strides=strides, dilation=4 * dilation, padding=4 * padding,
kernel_size=kernel_size)
self.net = nn.Sequential(b1, b2, b3, nn.AdaptiveMaxPool2d((2, 2)))
self.apply(weight_init)
def forward(self, x):
x = self.net(x)
x = x.view(config['batch_size'], config['seq_dim'], -1)
if self.lstm is None:
self.lstm = nn.LSTM(x.size(2), config['hidden_dim'], 1, batch_first=True).to(self.device)
for param in self.lstm.parameters():
if len(param.shape) >= 2:
init.orthogonal_(param.data)
else:
init.normal_(param.data)
h0 = torch.zeros(config['layer_dim'], x.size(0), config['hidden_dim']).to(self.device)
# Initialize cell state
c0 = torch.zeros(config['layer_dim'], x.size(0), config['hidden_dim']).to(self.device)
output, (hn, cn) = self.lstm(x, (h0,c0))
output = output[:, :, :]
output = self.fc(output)
return output, (hn, cn)
for epoch in range(config['num_epochs']):
print('epoch', epoch)
running_loss = 0
nb_classes = config['output_dim']
confusion_matrix = torch.zeros(nb_classes, nb_classes)
for i, (image, label) in enumerate(trainLoader):
print('batch: ',i)
image = image.float().to(device)
label = label.to(device)
optimizer.zero_grad()
batch_size, timesteps, H, W, C = image.size()
# Change Image shape
image = image.view(batch_size * timesteps, H, W, C)
image = image.permute(0, 3, 1, 2) # from NHWC to NCHW
output, (hn,cn) = model(image)
label = label.view(-1)
output = output.view(-1, output.size(2))
loss = criterion(output, label)
loss *= config['seq_dim']
loss.backward() # Backward pass
optimizer.step() # Now we can do an optimizer step
running_loss += loss.item()