Primitives

neps.search_spaces.architecture.primitives

AvgPool

AvgPool(kernel_size, stride, **kwargs)

Bases: AbstractPrimitive

Implementation of Avergae Pooling.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, kernel_size, stride, **kwargs):
    stride = int(stride)
    super().__init__(locals())
    self.avgpool = nn.AvgPool2d(3, stride=stride, padding=1, count_include_pad=False)

AvgPool1x1

AvgPool1x1(
    kernel_size, stride, C_in, C_out, affine=True, **kwargs
)

Bases: AbstractPrimitive

Implementation of Avergae Pooling with an optional 1x1 convolution afterwards. The convolution is required to increase the number of channels if stride > 1.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(
    self,
    kernel_size,
    stride,
    C_in,
    C_out,
    affine=True,
    **kwargs,
):
    super().__init__(locals())
    stride = int(stride)
    self.stride = int(stride)
    self.avgpool = nn.AvgPool2d(3, stride=stride, padding=1, count_include_pad=False)
    if stride > 1:
        assert C_in is not None and C_out is not None
        self.conv = nn.Conv2d(C_in, C_out, 1, stride=1, padding=0, bias=False)
        self.bn = nn.BatchNorm2d(C_out, affine=affine)

Concat1x1

Concat1x1(num_in_edges, C_out, affine=True, **kwargs)

Bases: AbstractPrimitive

Implementation of the channel-wise concatination followed by a 1x1 convolution to retain the channel dimension.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(
    self, num_in_edges, C_out, affine=True, **kwargs
):
    super().__init__(locals())
    self.conv = nn.Conv2d(
        num_in_edges * C_out, C_out, kernel_size=1, stride=1, padding=0, bias=False
    )
    self.bn = nn.BatchNorm2d(C_out, affine=affine)

forward

forward(x)

Expecting a list of input tensors. Stacking them channel-wise and applying 1x1 conv

Source code in neps/search_spaces/architecture/primitives.py

def forward(self, x):
    """
    Expecting a list of input tensors. Stacking them channel-wise
    and applying 1x1 conv
    """
    x = torch.cat(x, dim=1)
    x = self.conv(x)
    x = self.bn(x)
    return x

ConvBn

ConvBn(
    C_in,
    C_out,
    kernel_size,
    stride=1,
    affine=True,
    **kwargs
)

Bases: AbstractPrimitive

Implementation of 2d convolution, followed by 2d batch normalization and ReLU activation.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, C_in, C_out, kernel_size, stride=1, affine=True, **kwargs):
    super().__init__(locals())
    self.kernel_size = kernel_size
    pad = 0 if stride == 1 and kernel_size == 1 else 1
    self.op = nn.Sequential(
        nn.Conv2d(C_in, C_out, kernel_size, stride=stride, padding=pad, bias=False),
        nn.BatchNorm2d(C_out, affine=affine),
    )

ConvBnReLU

ConvBnReLU(
    C_in,
    C_out,
    kernel_size,
    stride=1,
    affine=True,
    **kwargs
)

Bases: AbstractPrimitive

Implementation of 2d convolution, followed by 2d batch normalization and ReLU activation.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, C_in, C_out, kernel_size, stride=1, affine=True, **kwargs):
    super().__init__(locals())
    self.kernel_size = kernel_size
    pad = 0 if stride == 1 and kernel_size == 1 else 1
    self.op = nn.Sequential(
        nn.Conv2d(C_in, C_out, kernel_size, stride=stride, padding=pad, bias=False),
        nn.BatchNorm2d(C_out, affine=affine),
        nn.ReLU(inplace=False),
    )

DilConv

DilConv(
    C_in,
    C_out,
    kernel_size,
    stride,
    padding,
    dilation,
    affine=True,
    **kwargs
)

Bases: AbstractPrimitive

Implementation of a dilated separable convolution as used in the DARTS paper.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(
    self, C_in, C_out, kernel_size, stride, padding, dilation, affine=True, **kwargs
):
    super().__init__(locals())

    C_in = int(C_in)
    C_out = int(C_out)
    kernel_size = int(kernel_size)
    stride = int(stride)
    padding = int(padding)
    dilation = int(dilation)
    affine = bool(affine)

    self.kernel_size = kernel_size
    self.op = nn.Sequential(
        nn.ReLU(inplace=False),
        nn.Conv2d(
            C_in,
            C_in,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=C_in,
            bias=False,
        ),
        nn.Conv2d(C_in, C_out, kernel_size=1, padding=0, bias=False),
        nn.BatchNorm2d(C_out, affine=affine),
    )

Identity

Identity(**kwargs)

Bases: AbstractPrimitive

An implementation of the Identity operation.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, **kwargs):
    super().__init__(locals())

MaxPool1x1

MaxPool1x1(
    kernel_size, stride, C_in, C_out, affine=True, **kwargs
)

Bases: AbstractPrimitive

Implementation of MaxPool with an optional 1x1 convolution in case stride > 1. The 1x1 convolution is required to increase the number of channels.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, kernel_size, stride, C_in, C_out, affine=True, **kwargs):
    super().__init__(locals())

    kernel_size = int(kernel_size)
    stride = int(stride)
    C_in = int(C_in)
    C_out = int(C_out)
    affine = bool(affine)

    self.stride = stride
    self.maxpool = nn.MaxPool2d(kernel_size, stride=stride, padding=1)
    if stride > 1:
        assert C_in is not None and C_out is not None
        self.conv = nn.Conv2d(C_in, C_out, 1, stride=1, padding=0, bias=False)
        self.bn = nn.BatchNorm2d(C_out, affine=affine)

SepConv

SepConv(
    C_in,
    C_out,
    kernel_size,
    stride,
    padding,
    affine=True,
    **kwargs
)

Bases: AbstractPrimitive

Implementation of Separable convolution operation as in the DARTS paper, i.e. 2 sepconv directly after another.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, C_in, C_out, kernel_size, stride, padding, affine=True, **kwargs):
    super().__init__(locals())

    C_in = int(C_in)
    C_out = int(C_out)
    kernel_size = int(kernel_size)
    stride = int(stride)
    padding = int(padding)
    affine = bool(affine)

    self.kernel_size = kernel_size
    self.op = nn.Sequential(
        nn.ReLU(inplace=False),
        nn.Conv2d(
            C_in,
            C_in,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            groups=C_in,
            bias=False,
        ),
        nn.Conv2d(C_in, C_in, kernel_size=1, padding=0, bias=False),
        nn.BatchNorm2d(C_in, affine=affine),
        nn.ReLU(inplace=False),
        nn.Conv2d(
            C_in,
            C_in,
            kernel_size=kernel_size,
            stride=1,
            padding=padding,
            groups=C_in,
            bias=False,
        ),
        nn.Conv2d(C_in, C_out, kernel_size=1, padding=0, bias=False),
        nn.BatchNorm2d(C_out, affine=affine),
    )

Sequential

Sequential(*args, **kwargs)

Bases: AbstractPrimitive

Implementation of torch.nn.Sequential to be used as op on edges.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, *args, **kwargs):
    super().__init__(locals())
    self.primitives = args
    self.op = nn.Sequential(*args)

Stem

Stem(C_out, C_in=3, **kwargs)

Bases: AbstractPrimitive

This is used as an initial layer directly after the image input.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, C_out, C_in=3, **kwargs):
    super().__init__(locals())

    C_out = int(C_out)

    self.seq = nn.Sequential(
        nn.Conv2d(C_in, C_out, 3, padding=1, bias=False), nn.BatchNorm2d(C_out)
    )

Zero

Zero(stride, **kwargs)

Bases: AbstractPrimitive

Implementation of the zero operation. It removes the connection by multiplying its input with zero.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, stride, **kwargs):
    """
    When setting stride > 1 then it is assumed that the
    channels must be doubled.
    """
    super().__init__(locals())
    self.stride = int(stride)

Zero1x1

Zero1x1(stride, **kwargs)

Bases: AbstractPrimitive

Implementation of the zero operation. It removes the connection by multiplying its input with zero.

Source code in neps/search_spaces/architecture/primitives.py

def __init__(self, stride, **kwargs):
    """
    When setting stride > 1 then it is assumed that the
    channels must be doubled.
    """
    super().__init__(locals())
    self.stride = int(stride)