diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100755 index 0000000..937733c --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,23 @@ +name: torch_sig_container_${PROJECT_NAME} +services: + torchsig_service: + build: . + image: torchsig:v0.5.0 + container_name: torchsig_${PROJECT_NAME} + stdin_open: true + tty: true + volumes: + - ./:/workspace/code + ports: + - '${JUP_PORT}:${JUP_PORT}' + environment: + - NVIDIA_VISIBLE_DEVICES=all + - NVIDIA_DRIVER_CAPABILITIES=all + command: jupyter lab --allow-root --ip=0.0.0.0 --no-browser --port ${JUP_PORT} --NotebookApp.token='' + shm_size: 512GB + deploy: + resources: + reservations: + devices: + - capabilities: [gpu] + driver: nvidia diff --git a/torchsig/models/iq_models/efficientnet/efficientnet1d.py b/torchsig/models/iq_models/efficientnet/efficientnet1d.py new file mode 100755 index 0000000..ba5f148 --- /dev/null +++ b/torchsig/models/iq_models/efficientnet/efficientnet1d.py @@ -0,0 +1,45 @@ +import timm +from torch.nn import Linear + +from torchsig.models.model_utils.model_utils_1d.conversions_to_1d import convert_2d_model_to_1d + +__all__ = ["EfficientNet1d"] + +def EfficientNet1d( + input_channels: int, + n_features: int, + efficientnet_version: str = "b0", + drop_path_rate: float = 0.2, + drop_rate: float = 0.3, +): + """Constructs and returns a 1d version of the EfficientNet model described in + `"EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks" `_. + + Args: + + input_channels (int): + Number of 1d input channels; e.g., common practice is to split complex number time-series data into 2 channels, representing the real and imaginary parts respectively + + n_features (int): + Number of output features; should be the number of classes when used directly for classification + + efficientnet_version (str): + Specifies the version of efficientnet to use. See the timm efficientnet documentation for details. Examples are 'b0', 'b1', and 'b4' + + drop_path_rate (float): + Drop path rate for training + + drop_rate (float): + Dropout rate for training + + """ + mdl = convert_2d_model_to_1d( + timm.create_model( + "efficientnet_" + efficientnet_version, + in_chans=input_channels, + drop_path_rate=drop_path_rate, + drop_rate=drop_rate, + ) + ) + mdl.classifier = Linear(mdl.classifier.in_features, n_features) + return mdl \ No newline at end of file diff --git a/torchsig/models/iq_models/xcit/xcit1d.py b/torchsig/models/iq_models/xcit/xcit1d.py new file mode 100755 index 0000000..07408f8 --- /dev/null +++ b/torchsig/models/iq_models/xcit/xcit1d.py @@ -0,0 +1,83 @@ +import timm +from torch import cat +from torch.nn import Module, Conv1d, Linear + +from torchsig.models.model_utils.model_utils_1d.iq_sampling import ConvDownSampler, Chunker + +__all__ = ["XCiT1d"] + +class XCiT1d(Module): + """A 1d implementation of the XCiT architecture from + `"XCiT: Cross-Covariance Image Transformers" `_. + + Args: + + input_channels (int): + Number of 1d input channels; e.g., common practice is to split complex number time-series data into 2 channels, representing the real and imaginary parts respectively + + n_features (int): + Number of output features; should be the number of classes when used directly for classification + + xcit_version (str): + Specifies the version of efficientnet to use. See the timm xcit documentation for details. Examples are 'nano_12_p16_224', and 'xcit_tiny_12_p16_224' + + drop_path_rate (float): + Drop path rate for training + + drop_rate (float): + Dropout rate for training + + ds_method (str): + Specifies the downsampling method to use in the model. Currently convolutional downsampling and chunking are supported, using string arguments 'downsample' and 'chunk' respectively + + ds_rate (int): + Specifies the downsampling rate; e.g., ds_rate=2 will downsample the imput by a factor of 2 + """ + def __init__(self, + input_channels: int, + n_features: int, + xcit_version: str = "nano_12_p16_224", + drop_path_rate: float = 0.0, + drop_rate: float = 0.3, + ds_method: str = "downsample", + ds_rate: int = 2): + + super().__init__() + self.backbone = timm.create_model( + "xcit_" + xcit_version, + num_classes=n_features, + in_chans=input_channels, + drop_path_rate=drop_path_rate, + drop_rate=drop_rate, + ) + + W = self.backbone.num_features + self.grouper = Conv1d(W, n_features, 1) + if ds_method == "downsample": + self.backbone.patch_embed = ConvDownSampler(input_channels, W, ds_rate) + elif ds_method == "chunk": + self.backbone.patch_embed = Chunker(input_channels, W, ds_rate) + else: + raise ValueError(ds_method + " is not a supported downsampling method; currently 'downsample', and 'chunk' are supported") + + self.backbone.head = Linear(self.backbone.head.in_features, n_features) + + def forward(self, x): + mdl = self.backbone + B = x.shape[0] + x = self.backbone.patch_embed(x) + + Hp, Wp = x.shape[-1], 1 + pos_encoding = mdl.pos_embed(B, Hp, Wp).reshape(B, -1, Hp).permute(0, 2, 1).half() + x = x.transpose(1, 2) + pos_encoding + for blk in mdl.blocks: + x = blk(x, Hp, Wp) + cls_tokens = mdl.cls_token.expand(B, -1, -1) + x = cat((cls_tokens, x), dim=1) + for blk in mdl.cls_attn_blocks: + x = blk(x) + x = mdl.norm(x) + x = self.grouper(x.transpose(1, 2)[:, :, :1]).squeeze() + if x.dim() == 2: + x = x.unsqueeze(0) + return x diff --git a/torchsig/models/model_utils/general_layers.py b/torchsig/models/model_utils/general_layers.py new file mode 100755 index 0000000..c4672e1 --- /dev/null +++ b/torchsig/models/model_utils/general_layers.py @@ -0,0 +1,82 @@ +from torch import mean +from torch.nn import Module, LSTM + +class DebugPrintLayer(Module): + """ + A layer for debugging pytorch models; prints out the shape and data type of the input tensor at runtime + returns he input tensor unchanged + """ + def __init__(self): + super().__init__() + + def forward(self, x): + print(x.shape, x.dtype) + return x + +class ScalingLayer(Module): + """ + A layer that given input tensor x outputs scale_val * x + used to linearly scale inputs by a fixed value + """ + def __init__(self, scale_val): + super().__init__() + self.scale_val = scale_val + + def forward(self, x): + return self.scale_val * x + +class DropChannel(Module): + """ + A layer that drops the last color channel of an image [must be in channel-first form] + """ + def __init__(self): + super().__init__() + + def forward(self, x): + return x[:,:-1,:,:] + +class LSTMImageReader(Module): + """ + TODO add some real documentation here + """ + def __init__(self, input_width, lstm_width, img_shape, num_layers=2): + super().__init__() + self.img_shape = img_shape + self.img_height = img_shape[0] + self.img_width = img_shape[1] + self.input_width = input_width + self.lstm_width = lstm_width + self.lstm_model = LSTM(self.input_width,self.lstm_width,num_layers,True,True,0,False,self.img_height) + + def forward(self, x): + output, (h,c) = self.lstm_model(x.transpose(1,2)) + img_tensor = output.transpose(1,2)[:,:self.img_height,:self.img_width] #take only the last img_height entries in the outut sequence + return img_tensor.reshape([x.size(0),1,self.img_height,self.img_width]) + +class Reshape(Module): + """ + A layer that reshapes the input tensor to a tensor of the given shape + if keep_batch_dim is True (defaults to True), the batch dimension is excluded from the reshape operation; otherwise it is included + """ + def __init__(self, shape, keep_batch_dim=True): + super(Reshape, self).__init__() + self.shape = shape + self.keep_batch_dim = keep_batch_dim + + def forward(self, x): + if self.keep_batch_dim: + batch_dim = x.size(0) + shape = [batch_dim] + list(self.shape) + return x.view(shape) + return x.view(self.shape) + +class Mean(Module): + """ + A layer which returns the mean(s) along the dimension specified by dim of the input tensor + """ + def __init__(self, dim): + super(Mean, self).__init__() + self.dim = dim + + def forward(self, x): + return mean(x,self.dim) \ No newline at end of file diff --git a/torchsig/models/model_utils/layer_tools.py b/torchsig/models/model_utils/layer_tools.py new file mode 100755 index 0000000..42b107d --- /dev/null +++ b/torchsig/models/model_utils/layer_tools.py @@ -0,0 +1,111 @@ +def get_layer_list(model): + """ + returns a list of all layers in the input model, including layers in any nested models therein + layers are listed in forward-pass order + """ + arr = [] + final_arr = [] + try: + arr = [m for m in model.modules()] + if len(arr) > 1: + for module in arr[1:]: + final_arr += get_module_list(module) + return final_arr + else: + return arr + except: + raise(NotImplementedError("expected module list to be populated, but no '_modules' field was found")) + +def replace_layer(old_layer, new_layer, model): + """ + search through model until old_layer is found, and replace it with new layer; + returns True is old_layer was found; False otherwise + """ + try: + modules = model._modules + for k in modules.keys(): + if modules[k] == old_layer: + modules[k] = new_layer + return True + else: + if replace_layer(old_layer, new_layer, modules[k]): + return True + return False + except: + raise(NotImplementedError("expected module list to be populated, but no '_modules' field was found")) + +def is_same_type(layer1, layer2): + """ + returns True if layer1 and layer2 are of the same type; false otherwise + if a class is input as layer2 [e.g., is_same_type(my_conv_layer, Conv2d) ], the type defined by the class is used + if a string is input as layer2, the string is matched to the name of the class of layer1 + """ + if type(layer2) == type: + return type(layer1) == layer2 + elif type(layer2) == str: + return type(layer1).__name__ == layer2 + else: + return type(layer1) == type(layer2) + +def same_type_fn(layer1): + """ + curried version of is_same_type; returns a function f such than f(layer2) <-> is_same_type(layer1, layer2) + """ + return lambda x: is_same_type(x, layer1) + + +def replace_layers_on_condition(model, condition_fn, new_layer_factory_fn): + """ + search through model finding all layers L such that conditional_fn(L), and replace them with new_layer_factory_fn(L) + returns true if at least one layer was replaced; false otherwise + """ + has_replaced = False + try: + modules = model._modules + for k in modules.keys(): + if condition_fn(modules[k]): + modules[k] = new_layer_factory_fn(modules[k]) + has_replaced = True + else: + has_replaced = replace_layers_on_condition(modules[k], condition_fn, new_layer_factory_fn) or has_replaced + return has_replaced + except: + raise(NotImplementedError("expected module list to be populated, but no '_modules' field was found")) + +def replace_layers_on_conditions(model, condition_factory_pairs): + """ + search through model finding all layers L such that for some ordered pair [conditional_fn, new_layer_factory_fn] in condition_factory_pairs, + conditional_fn(L), and replace them with new_layer_factory_fn(L) + layers will only be replaced once, so the first conditional for which a layer returns true will be last conditional to which it is compared + returns true if at least one layer was replaced; false otherwise + """ + has_replaced = False + try: + modules = model._modules + for k in modules.keys(): + for (condition_fn, new_layer_factory_fn) in condition_factory_pairs: + if condition_fn(modules[k]): + modules[k] = new_layer_factory_fn(modules[k]) + has_replaced = True + break + else: + has_replaced = replace_layers_on_conditions(modules[k], condition_factory_pairs) or has_replaced + return has_replaced + except: + raise(NotImplementedError("expected module list to be populated, but no '_modules' field was found")) + +def replace_layers_of_type(model, layer_type, new_layer_factory_fn): + """ + search through model finding all layers L of type layer_type and replace with new_layer_factory_fn(L) + returns true if at least one layer was replaced; false otherwise + """ + return replace_layers_on_condition(model, lambda x: is_same_type(x,layer_type), new_layer_factory_fn) + +def replace_layers_of_types(model, type_factory_pairs): + """ + search through model finding all layers L such that for some ordered pair [layer_type, new_layer_factory_fn] in type_factory_pairs, + L is of type layer_type, and replace with new_layer_factory_fn(L) + returns true if at least one layer was replaced; false otherwise + """ + condition_factory_pairs = [(same_type_fn(layer_type), new_layer_factory_fn) for (layer_type, new_layer_factory_fn) in type_factory_pairs] + return replace_layers_on_conditions(model, condition_factory_pairs) \ No newline at end of file diff --git a/torchsig/models/model_utils/model_utils_1d/conversions_to_1d.py b/torchsig/models/model_utils/model_utils_1d/conversions_to_1d.py new file mode 100755 index 0000000..3c8e377 --- /dev/null +++ b/torchsig/models/model_utils/model_utils_1d/conversions_to_1d.py @@ -0,0 +1,63 @@ +from torchsig.models.model_utils.layer_tools import replace_layers_of_types +from torchsig.models.model_utils.model_utils_1d.layers_1d import GBN1d, SqueezeExcite1d, FastGlobalAvgPool1d + +from torch.nn import Conv1d, BatchNorm1d + +def conv2d_to_conv1d(layer_2d): + """ + returns a 1d conv layer corresponding to the input conv2d layer + no muation is performed + """ + return Conv1d(layer_2d.in_channels, layer_2d.out_channels, layer_2d.kernel_size[0], + stride=layer_2d.stride[0], padding=layer_2d.padding[0], dilation=layer_2d.dilation[0], + groups=layer_2d.groups, bias=(True if layer_2d.bias != None else False), padding_mode=layer_2d.padding_mode, device=layer_2d.weight.device, + dtype=layer_2d.weight.dtype + ) + +def batchNorm2d_to_batchNorm1d(layer_2d): + """ + returns a BatchNorm1d layer corresponding to the input BatchNorm2d layer + no muation is performed + """ + return BatchNorm1d(layer_2d.num_features, eps=layer_2d.eps, momentum=layer_2d.momentum, + affine=layer_2d.affine, track_running_stats=layer_2d.track_running_stats, + device=layer_2d.weight.device, dtype=layer_2d.weight.dtype + ) + +def batchNorm2d_to_GBN1d(layer_2d): + """ + returns a GBN1d [Ghost Batch Norm] layer corresponding to the input BatchNorm2d layer + no muation is performed + """ + return GBN1d(layer_2d.num_features) + +def squeezeExcite_to_squeezeExcite1d(layer_2d): + """ + returns a GBN1d [Ghost Batch Norm] layer corresponding to the input BatchNorm2d layer + no muation is performed + """ + return SqueezeExcite1d(layer_2d.conv_reduce.in_channels, reduced_base_chs=layer_2d.conv_reduce.out_channels) + +def make_fast_avg_pooling_layer(layer_2d): + """ + returns a FastGlobalAvgPool1d layer + """ + return FastGlobalAvgPool1d(flatten=True) + +def convert_2d_model_to_1d(model): + """ + converts a 2d model to a corresponding 1d model by replacing convolutional layers and other 2d layers with their 1d equivalents + experimental; may not fully convert models with unrecognized layer types + mutates input model; returns the mutated model + this function is still under development and may not correctly convert all 2d layer types, or may have unexpected behavior on models that perform reshaping internally + """ + type_factory_pairs = [ + ('Conv2d', conv2d_to_conv1d), + ('BatchNorm2d', batchNorm2d_to_GBN1d), + ('SqueezeExcite', squeezeExcite_to_squeezeExcite1d), + ('SelectAdaptivePool2d',make_fast_avg_pooling_layer), + ] + + replace_layers_of_types(model, type_factory_pairs) + return model + diff --git a/torchsig/models/model_utils/model_utils_1d/iq_sampling.py b/torchsig/models/model_utils/model_utils_1d/iq_sampling.py new file mode 100755 index 0000000..6ede218 --- /dev/null +++ b/torchsig/models/model_utils/model_utils_1d/iq_sampling.py @@ -0,0 +1,50 @@ +from torch import cat, split +from torch.nn import Module, Conv1d, BatchNorm1d, SiLU, Sequential + +__all__ = ["ConvDownSampler", "Chunker"] + +class ConvDownSampler(Module): + def __init__(self, in_chans, embed_dim, ds_rate=16): + super().__init__() + ds_rate //= 2 + chan = embed_dim // ds_rate + blocks = [Conv1d(in_chans, chan, 5, 2, 2), BatchNorm1d(chan), SiLU()] + + while ds_rate > 1: + blocks += [ + Conv1d(chan, 2 * chan, 5, 2, 2), + BatchNorm1d(2 * chan), + SiLU(), + ] + ds_rate //= 2 + chan = 2 * chan + + blocks += [ + Conv1d( + chan, + embed_dim, + 1, + ) + ] + self.blocks = Sequential(*blocks) + + def forward(self, X): + return self.blocks(X) + + +class Chunker(Module): + def __init__(self, in_chans, embed_dim, ds_rate=16): + super().__init__() + self.embed = Conv1d(in_chans, embed_dim // ds_rate, 7, padding=3) + self.project = Conv1d((embed_dim // ds_rate) * ds_rate, embed_dim, 1) + self.ds_rate = ds_rate + + def forward(self, X): + X = self.embed(X) + X = cat( + [cat(split(x_i, 1, -1), 1) for x_i in split(X, self.ds_rate, -1)], + -1, + ) + X = self.project(X) + + return X \ No newline at end of file diff --git a/torchsig/models/model_utils/model_utils_1d/layers_1d.py b/torchsig/models/model_utils/model_utils_1d/layers_1d.py new file mode 100755 index 0000000..b919f1b --- /dev/null +++ b/torchsig/models/model_utils/model_utils_1d/layers_1d.py @@ -0,0 +1,82 @@ +import numpy as np +from torch import sigmoid, cat +from torch.nn import Module, functional, SiLU, Conv1d, BatchNorm1d + +class SqueezeExcite1d(Module): + def __init__( + self, + in_chs, + se_ratio=0.25, + reduced_base_chs=None, + act_layer=SiLU, + gate_fn=sigmoid, + divisor=1, + **_, + ): + super(SqueezeExcite1d, self).__init__() + reduced_chs = reduced_base_chs + self.conv_reduce = Conv1d(in_chs, reduced_chs, 1, bias=True) + self.act1 = act_layer(inplace=True) + self.conv_expand = Conv1d(reduced_chs, in_chs, 1, bias=True) + self.gate_fn = gate_fn + + def forward(self, x): + x_se = x.mean((2,), keepdim=True) + x_se = self.conv_reduce(x_se) + x_se = self.act1(x_se) + x_se = self.conv_expand(x_se) + return x * self.gate_fn(x_se) + + +class FastGlobalAvgPool1d(Module): + def __init__(self, flatten=False): + super(FastGlobalAvgPool1d, self).__init__() + self.flatten = flatten + + def forward(self, x): + if self.flatten: + in_size = x.size() + return x.view((in_size[0], in_size[1], -1)).mean(dim=2) + else: + return x.view(x.size(0), x.size(1), -1).mean(-1).view(x.size(0), x.size(1), 1) + + +class GBN1d(Module): + """ + Ghost Batch Normalization + https://arxiv.org/abs/1705.08741 + """ + + def __init__(self, input_dim, virtual_batch_size=128, momentum=0.1): + super(GBN1d, self).__init__() + + self.input_dim = input_dim + self.virtual_batch_size = virtual_batch_size + self.bn = BatchNorm1d(self.input_dim, momentum=momentum) + + def forward(self, x): + chunks = x.chunk(int(np.ceil(x.shape[0] / self.virtual_batch_size)), 0) + res = [self.bn(x_) for x_ in chunks] + + return cat(res, dim=0) + +class ImageFrom1D(Module): + """ + A layer for reshaping (batch_size x n_channels x N) 1d signal data to (batch_size x n_channels x new_y x new_x) + where new_y and new_x are N**0.5 padded to as near a perfect square as can be formed from N points without adding more than a full row of padding + """ + def __init__(self, n_channels = 3): + super().__init__() + self.n_channels = n_channels + + def forward(self, x): + batch_size = x.size(0) + x_flat = x.reshape(batch_size,self.n_channels, -1) + n_values = x_flat.size(2) + y_dim = np.sqrt(n_values).astype(np.int32) + x_dim = (n_values // y_dim) + n_pad = ((batch_size * self.n_channels * y_dim * x_dim) - (x_flat.size(0) * x_flat.size(1) * x_flat.size(2)))//(x_flat.size(0) * x_flat.size(1)) + x_flat = functional.pad(x_flat, [0,n_pad]) + img_tensor = x_flat.reshape(batch_size,self.n_channels,y_dim,x_dim) + return img_tensor + diff --git a/torchsig/models/model_utils/simple_models.py b/torchsig/models/model_utils/simple_models.py new file mode 100755 index 0000000..e0c2862 --- /dev/null +++ b/torchsig/models/model_utils/simple_models.py @@ -0,0 +1,105 @@ +#A script for defining some useful functions and classes for building models + +import numpy as np + +from torch.nn import ELU, Sequential, Conv2d, Conv1d, ConvTranspose2d, BatchNorm1d, BatchNorm2d, Linear + +STANDARD_ACTIVATION_FUNCTION = ELU + +def convnet_block_2d(in_width, out_width, kernel_shape = [5,5], activation_fn = None): + """ + returns a block of layers consisting of a 2d convolution, batch normalization, and activation function, + with the input and output channels given by in_width and out_with, a kernel given by kernel_shape, and using the given activation_fn + if no activation function is provided, this defaults to ELU + """ + if not activation_fn: + activation_fn = STANDARD_ACTIVATION_FUNCTION + return Sequential( + Conv2d(in_width, out_width, kernel_shape, padding=[kernel_dim//2 for kernel_dim in kernel_shape]), + BatchNorm2d(out_width), + activation_fn() + ) + +def convnet_block_1d(in_width, out_width, kernel_dim = 5, activation_fn = None): + """ + 1d version of convnet_block_2d above + """ + if not activation_fn: + activation_fn = STANDARD_ACTIVATION_FUNCTION + return Sequential( + Conv1d(in_width, out_width, kernel_dim, padding=kernel_dim//2), + BatchNorm1d(out_width), + activation_fn() + ) + +def dense_block(in_width, out_width, activation_fn = None): + """ + returns a block of layers consisting of a 2d convolution, batch normalization, and activation function, + with the input and output channels given by in_width and out_with, a kernel given by kernel_shape, and using the given activation_fn + if no activation function is provided, this defaults to ELU + """ + if not activation_fn: + activation_fn = STANDARD_ACTIVATION_FUNCTION + return Sequential( + Linear(in_width, out_width), + BatchNorm1d(out_width), + activation_fn() + ) + +def simple_convnet_2d(layer_width_list): + """ + takes in a list or tuple of convoluional channel widths and returns a sequential model with those widths + used to quickly prototype convolutional neural nets; + for example, simple_convnet_2d([3,8,32,64,64,1]) would return a model with 5 convolutional layers that takes in + an X by Y image with 3 color channels and outputs an X by Y image with a single channel. + Because the returned model doesnt include pooling, striding, or dilation, etc., it does not reduce the scale of the input except possibly in the channel dimension + As such, it can take up a great deal of memory, and should not be used by itself to perform complicated tasks on large images + """ + layers = [] + prev_width = layer_width_list[0] + for layer_width in layer_width_list[1:]: + layers += [convnet_block_2d(prev_width, layer_width)] + prev_width = layer_width + layers += [Conv2d(prev_width, prev_width, [1,1])] + return Sequential(*layers) + +def simple_convnet_1d(layer_width_list): + """ + 1d version of simple_convnet_2d above + """ + layers = [] + prev_width = layer_width_list[0] + for layer_width in layer_width_list[1:]: + layers += [convnet_block_1d(prev_width, layer_width)] + prev_width = layer_width + layers += [Conv1d(prev_width, prev_width, 1)] + return Sequential(*layers) + +def simple_densenet(layer_width_list): + """ + takes in a list or tuple of dense layer widths and returns a sequential model with those widths + used to quickly prototype simple feed-forward neural nets; + for example, simple_densenet([6,8,32,64,64,1]) would return a model with fully-connected linear layers that takes in + a vecor of length 6 and outputs a single value. + """ + layers = [] + prev_width = layer_width_list[0] + for layer_width in layer_width_list[1:]: + layers += [dense_block(prev_width, layer_width)] + prev_width = layer_width + layers += [Linear(prev_width, prev_width)] + return Sequential(*layers) + +def double_image_scale_2d(width = 1, kernel_shape = [5,5], activation_fn = None): + """ + doubles the scale dimensions of an image of channel width width using a transposd convolution of kernel shape kernel_shape + calls batch norm and an activation function provided by activation_fn on the result + if no function is provided, this defaults to ELU + """ + if not activation_fn: + activation_fn = ELU + return Sequential( + ConvTranspose2d(width, width, kernel_shape, padding=[kernel_dim//2 for kernel_dim in kernel_shape], stride=2, output_padding=1), + BatchNorm2d(width), + activation_fn() + ) \ No newline at end of file