File size: 2,282 Bytes
c8ddb9b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
"""Frequently used convolution modules"""
from torch import nn
from typing import Tuple
def conv2d(
in_channels: int,
out_channels: int,
kernel_size: int = 3,
stride: int = 1,
padding: int = 1,
) -> nn.Conv2d:
"""
Template convolution which is typically used throughout the project
:param int in_channels: Number of input channels
:param int out_channels: Number of output channels
:param int kernel_size: Size of sliding kernel
:param int stride: How many steps kernel does when sliding
:param int padding: How many dimensions to pad
:return: Convolution layer with parameters
:rtype: nn.Conv2d
"""
return nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
)
def conv1d(
in_channels: int,
out_channels: int,
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
) -> nn.Conv1d:
"""
Template 1d convolution which is typically used throughout the project
:param int in_channels: Number of input channels
:param int out_channels: Number of output channels
:param int kernel_size: Size of sliding kernel
:param int stride: How many steps kernel does when sliding
:param int padding: How many dimensions to pad
:return: Convolution layer with parameters
:rtype: nn.Conv2d
"""
return nn.Conv1d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
)
def calc_out_conv(
h_in: int, w_in: int, kernel_size: int = 3, stride: int = 1, padding: int = 0
) -> Tuple[int, int]:
"""
Calculate the dimensionalities of images propagated through conv layers
:param h_in: Height of the image
:param w_in: Width of the image
:param kernel_size: Size of sliding kernel
:param stride: How many steps kernel does when sliding
:param padding: How many dimensions to pad
:return: Height and width of image through convolution
:rtype: tuple[int, int]
"""
h_out = int((h_in + 2 * padding - kernel_size) / stride + 1)
w_out = int((w_in + 2 * padding - kernel_size) / stride + 1)
return h_out, w_out
|