diff --git a/steganogan/encoders.py b/steganogan/encoders.py index d1a93c8..8742234 100644 --- a/steganogan/encoders.py +++ b/steganogan/encoders.py @@ -2,6 +2,20 @@ import torch from torch import nn +import torch +from torch import nn +import torch.onnx +from torchvision.ops.deform_conv import DeformConv2d + +input = torch.rand(4, 3, 10, 10) +kh, kw = 3, 3 +weight = torch.rand(5, 3, kh, kw) +# offset and mask should have the same spatial size as the output +# of the convolution. In this case, for an input of 10, stride of 1 +# and kernel size of 3, without padding, the output size is 8 +offset = torch.rand(4, 2 * kh * kw, 8, 8) +mask = torch.rand(4, kh * kw, 8, 8) +out = deform_conv2d(input, offset, weight, mask=mask) class BasicEncoder(nn.Module): @@ -16,7 +30,7 @@ class BasicEncoder(nn.Module): add_image = False def _conv2d(self, in_channels, out_channels): - return nn.Conv2d( + return DeformConv2d( in_channels=in_channels, out_channels=out_channels, kernel_size=3,