Not quite. This profile definitely suggests it should, but in the actual network I am running (a fairly simple 16-layer encoder/decoder), a batch with the cudnn benchmark on still takes ~3.6s on average after 100’s of batches have been run through. That is, after the algorithm selection should have been optimized, it still takes 3.6s / batch (forward+backward). My implementation drops that down to ~1.6s / batch (forward+backward). I still can’t account for the difference yet. This is with an input size of 3 x 20480 x 3 x 4 x 4 (B x P x C x H x W).
I have yet to try the unfold method in the network, so it’s also possible that will end up winning in total. I’m going to try that now.
All in all, it seems silly that these methods should run so much less efficiently when compared over sequential calls given that they all profile so similarly on an individual basis.
If you’re feeling adventurous (or particularly bored), feel free to see what kind of speed you can get. I’d be curious to know if you discover a better method of doing it. I’ve pasted my network below (mind you it currently uses the PatchConvolution with my custom CUDA code, but that can be swapped out easily for any of these methods we’ve evaluated). I am training with ADAM(lr=1e-4) and using the MSELoss if that makes any difference (it shouldn’t).
import torch
import torch.nn as nn
import torch.nn.functional as F
from patch_convolution import *
##-----------------------------------------------------------------------------
class Network(nn.Module):
def __init__(self):
# Call the super constructor
super(Network, self).__init__()
# 1x
self.encoder0_0 = ConvELUBlock(in_channels=3,
out_channels=32,
kernel_size=3,
padding=1)
self.encoder0_1 = ConvELUBlock(in_channels=32,
out_channels=64,
kernel_size=3,
padding=1)
# 2x
self.encoder1_0 = ConvELUBlock(in_channels=64,
out_channels=128,
kernel_size=3,
padding=1,
stride=2)
self.encoder1_1 = ConvELUBlock(in_channels=128,
out_channels=128,
kernel_size=3,
padding=1)
self.encoder1_2 = ConvELUBlock(in_channels=128,
out_channels=128,
kernel_size=3,
padding=1)
# 4x
self.encoder2_0 = ConvELUBlock(in_channels=128,
out_channels=256,
kernel_size=3,
padding=1,
stride=2)
self.encoder2_1 = ConvELUBlock(in_channels=256,
out_channels=256,
kernel_size=3,
padding=1)
self.encoder2_2 = ConvELUBlock(in_channels=256,
out_channels=256,
kernel_size=3,
padding=1)
# 4x
self.encoder3_0 = ConvELUBlock(in_channels=256,
out_channels=512,
kernel_size=3,
padding=1)
self.encoder3_1 = ConvELUBlock(in_channels=512,
out_channels=512,
kernel_size=3,
padding=1)
# 2x
self.decoder0_0 = ConvTransposeELUBlock(in_channels=512,
out_channels=256,
kernel_size=4,
padding=1,
stride=2)
self.decoder0_1 = ConvELUBlock(in_channels=256,
out_channels=256,
kernel_size=3,
padding=1)
# 1x
self.decoder1_0 = ConvTransposeELUBlock(in_channels=256,
out_channels=128,
kernel_size=4,
padding=1,
stride=2)
self.decoder1_1 = ConvELUBlock(in_channels=128,
out_channels=128,
kernel_size=3,
padding=1)
self.decoder1_2 = ConvELUBlock(in_channels=128,
out_channels=64,
kernel_size=1)
# Prediction
self.prediction = ConvELUBlock(in_channels=64,
out_channels=1,
kernel_size=3,
padding=1)
self.apply(xavier_init)
def forward(self, x):
encoder0_0_out = self.encoder0_0(x)
encoder0_1_out = self.encoder0_1(encoder0_0_out)
encoder1_0_out = self.encoder1_0(encoder0_1_out)
encoder1_1_out = self.encoder1_1(encoder1_0_out)
encoder1_2_out = self.encoder1_2(encoder1_1_out)
encoder2_0_out = self.encoder2_0(encoder1_2_out)
encoder2_1_out = self.encoder2_1(encoder2_0_out)
encoder2_2_out = self.encoder2_2(encoder2_1_out)
encoder3_0_out = self.encoder3_0(encoder2_2_out)
encoder3_1_out = self.encoder3_1(encoder3_0_out)
decoder0_0_out = self.decoder0_0(encoder3_1_out)
decoder0_1_out = self.decoder0_1(decoder0_0_out)
decoder1_0_out = self.decoder1_0(decoder0_1_out)
decoder1_1_out = self.decoder1_1(decoder1_0_out)
decoder1_2_out = self.decoder1_2(decoder1_1_out)
pred = self.prediction(decoder1_2_out)
return pred
##-----------------------------------------------------------------------------
class ConvELUBlock(nn.Module):
'''Convenience function combining a convolution with an ELU activation'''
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1):
super(ConvELUBlock, self).__init__()
self.conv = PatchConvolution(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation)
def forward(self, x):
return F.elu(self.conv(x), inplace=True)
##-----------------------------------------------------------------------------
class ConvTransposeELUBlock(nn.Module):
'''Convenience function combining a convolution with an ELU activation'''
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1):
super(ConvTransposeELUBlock, self).__init__()
self.conv = TransposedPatchConvolution(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation)
def forward(self, x):
return F.elu(self.conv(x), inplace=True)
.