Pytorch_CNN-RNN/utils/CNN_methods.py

from torch import add
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

"""
Returns a function that convolutes or separable convolutes, normalizes, activates (ELU), pools and dropouts input.
 
Kernel_size = (height, width, depth)

CAN DO SEPARABLE CONVOLUTION IF GROUP = 2!!!! :))))
"""

def conv_elu_maxpool_drop(in_channel, filters, kernel_size, stride=(1,1,1), padding=0, dilation=1,
                      groups=1, bias=True, padding_mode='zeros', pool=False, drop_rate=0, sep_conv = False):
    def f(input):

        # SEPARABLE CONVOLUTION
        if(sep_conv):

            # SepConv depthwise, Normalizes, and ELU activates
            sepConvDepthwise = nn.Conv3d(in_channel, filters, kernel_size, stride=stride, padding=padding,
                                         groups=in_channel, bias=bias, padding_mode=padding_mode)(input)

            # SepConv pointwise
            # Todo, will stride & padding be correct for this?
            conv = nn.Conv3d(in_channel, filters, kernel_size=1, stride=stride, padding=padding,
                                         groups=1, bias=bias, padding_mode=padding_mode)(sepConvDepthwise)

        # CONVOLUTES
        else:
            # Convolutes, Normalizes, and ELU activates
            conv = nn.Conv3d(in_channel, filters, kernel_size, stride=stride, padding=padding, dilation=dilation,
                             groups=groups, bias=bias, padding_mode=padding_mode)(input)

        normalization = nn.BatchNorm3d(filters)(conv)
        elu = nn.ELU()(normalization)

        # Pools
        if (pool):
            elu = nn.MaxPool2d(kernel_size=3, stride=2, padding=0)(elu)

        return nn.Dropout(p=drop_rate)(elu)
    return f


'''
Mid_flow in CNN. sep_convolutes 3 times, adds residual (initial input) to 3 times convoluted, and activates through ELU()
'''

def mid_flow(I, drop_rate, filters):
    in_channel = None   # TODO, IN_CHANNEL

    residual = I        # TODO, DOES THIS ACTUALLY COPY?

    x = conv_elu_maxpool_drop(in_channel, filters, (3,3,3), drop_rate=drop_rate)(I)
    x = conv_elu_maxpool_drop(in_channel, filters, (3,3,3), drop_rate=drop_rate)(x)
    x = conv_elu_maxpool_drop(in_channel, filters, (3, 3, 3), drop_rate=drop_rate)(x)

    x = add(x, residual)
    x = nn.ELU()(x)
    return x


"""
Returns a function that Fully Connects (FC), normalizes, activates (ELU), and dropouts input.
"""

def fc_elu_drop(in_features, units, drop_rate=0):
    def f(input):

        fc = nn.Linear(in_features, out_features=units)(input)
        fc = nn.BatchNorm3d(units)(fc)          # TODO 3d or 2d???
        fc = nn.ELU()(fc)
        fc = nn.Dropout(p=drop_rate)
        return fc

    return f