| import tensorflow as tf |
| from tensorflow.keras.layers import Conv2d,LayerNormalization,ZeroPadding2D,UpSampling2D,Activation |
| from tensorflow.keras import Model |
| from einops import rearrange |
| from math import sqrt |
| from functools import partial |
|
|
| |
|
|
| def exists(val): |
| return val is not None |
|
|
| def cast_tuple(val, depth): |
| return val if isinstance(val, tuple) else (val,) * depth |
|
|
| |
|
|
| class DsConv2d: |
| def __init__(self, dim_in, dim_out, kernel_size, padding, stride = 1, bias = True): |
| self.net = tf.keras.Sequential() |
| self.net.add(Conv2d(dim_in, kernel_size = kernel_size, strides = stride, use_bias = bias)) |
| self.net.add(ZeroPadding2D(padding)) |
| self.net.add(Conv2d(dim_out, kernel_size = 1, use_bias = bias)) |
| |
| def __call__(self, x): |
| return self.net(x) |
|
|
| class LayerNorm(tf.keras.layers.Layer): |
| def __init__(self, dim, eps = 1e-5): |
| self.eps = eps |
| self.g = self.add_weight( |
| name='g', |
| shape=(1, dim, 1, 1), |
| initializer=tf.keras.initializers.Ones(), |
| trainable=True |
| ) |
| self.b = self.add_weight( |
| name='b', |
| shape=(1, dim, 1, 1), |
| initializer=tf.keras.initializers.Zeros(), |
| trainable=True |
| ) |
|
|
| def __call__(self, x): |
| std = tf.math.sqrt(tf.math.reduce_variance(x, axis=1, keepdims=True)) |
| mean = tf.reduce_mean(x, axis= 1, keepdim = True) |
| return (x - mean) / (std + self.eps) * self.g + self.b |
|
|
| class PreNorm: |
| def __init__(self, dim, fn): |
| self.fn = fn |
| self.norm = LayerNormalization() |
|
|
| def __call__(self, x): |
| return self.fn(self.norm(x)) |
|
|
| class EfficientSelfAttention: |
| def __init__( |
| self, |
| dim, |
| heads, |
| reduction_ratio |
| ): |
| self.scale = (dim // heads) ** -0.5 |
| self.heads = heads |
|
|
| self.to_q = Conv2d(dim, 1, use_bias = False) |
| self.to_kv = Conv2d(dim * 2, reduction_ratio, strides = reduction_ratio, use_bias = False) |
| self.to_out = Conv2d(dim, 1, use_bias = False) |
|
|
| def __call__(self, x): |
| h, w = x.shape[1], x.shape[2] |
| heads = self.heads |
|
|
| q, k, v = (self.to_q(x), *tf.split(self.to_kv(x), num_or_size_splits=2, axis=-1)) |
| q, k, v = map(lambda t: rearrange(t, 'b x y (h c) -> (b h) (x y) c', h = heads), (q, k, v)) |
| |
| sim = tf.einsum('b i d, b j d -> b i j', q, k) * self.scale |
| attn = tf.nn.softmax(sim) |
|
|
| out = tf.einsum('b i j, b j d -> b i d', attn, v) |
| out = rearrange(out, '(b h) (x y) c -> b x y (h c)', h = heads, x = h, y = w) |
| return self.to_out(out) |
|
|
| class MixFeedForward: |
| def __init__( |
| self, |
| dim, |
| expansion_factor |
| ): |
| hidden_dim = dim * expansion_factor |
| self.net = tf.keras.Sequential() |
| self.net.add(Conv2d(hidden_dim, 1)) |
| self.net.add(DsConv2d(hidden_dim, hidden_dim, 3, padding = 1)) |
| self.net.add(Activation('gelu')) |
| self.net.add(Conv2d(dim, 1)) |
|
|
| def __call__(self, x): |
| return self.net(x) |
|
|
| class Unfold: |
| def __init__(self, kernel, stride, padding): |
| self.kernel = kernel |
| self.stride = stride |
| self.padding = padding |
| self.zeropadding2d = ZeroPadding2D(padding) |
| |
| def __call__(self, x): |
| x = self.zeropadding2d(x) |
| x = tf.image.extract_patches(x, sizes=[1, self.kernel, self.kernel, 1], strides=[1, self.stride, self.stride, 1], rates=[1, 1, 1, 1], padding='VALID') |
| x = tf.reshape(x, (x.shape[0], -1, x.shape[-1])) |
| return x |
|
|
| class MiT: |
| def __init__( |
| self, |
| channels, |
| dims, |
| heads, |
| ff_expansion, |
| reduction_ratio, |
| num_layers |
| ): |
| stage_kernel_stride_pad = ((7, 4, 3), (3, 2, 1), (3, 2, 1), (3, 2, 1)) |
|
|
| dims = (channels, *dims) |
| dim_pairs = list(zip(dims[:-1], dims[1:])) |
|
|
| self.stages = [] |
| |
| for (dim_in, dim_out), (kernel, stride, padding), num_layers, ff_expansion, heads, reduction_ratio in zip(dim_pairs, stage_kernel_stride_pad, num_layers, ff_expansion, heads, reduction_ratio): |
| get_overlap_patches = Unfold(kernel, stride, padding) |
| overlap_patch_embed = Conv2d(dim_out, 1) |
|
|
| layers = [] |
|
|
| for _ in range(num_layers): |
| layers.append([ |
| PreNorm(dim_out, EfficientSelfAttention(dim = dim_out, heads = heads, reduction_ratio = reduction_ratio)), |
| PreNorm(dim_out, MixFeedForward(dim = dim_out, expansion_factor = ff_expansion)), |
| ]) |
|
|
| self.stages.append([ |
| get_overlap_patches, |
| overlap_patch_embed, |
| layers |
| ]) |
|
|
| def __call__( |
| self, |
| x, |
| return_layer_outputs = False |
| ): |
| h, w = x.shape[1], x.shape[2] |
|
|
| layer_outputs = [] |
| for (get_overlap_patches, overlap_embed, layers) in self.stages: |
| x = get_overlap_patches(x) |
|
|
| num_patches = x.shape[-2] |
| ratio = int(sqrt((h * w) / num_patches)) |
| x = rearrange(x, 'b (h w) c -> b h w c', h = h // ratio) |
|
|
| x = overlap_embed(x) |
| for (attn, ff) in layers: |
| x = attn(x) + x |
| x = ff(x) + x |
|
|
| layer_outputs.append(x) |
|
|
| ret = x if not return_layer_outputs else layer_outputs |
| return ret |
|
|
| class Segformer(Model): |
| def __init__( |
| self, |
| dims = (32, 64, 160, 256), |
| heads = (1, 2, 5, 8), |
| ff_expansion = (8, 8, 4, 4), |
| reduction_ratio = (8, 4, 2, 1), |
| num_layers = 2, |
| channels = 3, |
| decoder_dim = 256, |
| num_classes = 4 |
| ): |
| super(Segformer, self).__init__() |
| dims, heads, ff_expansion, reduction_ratio, num_layers = map(partial(cast_tuple, depth = 4), (dims, heads, ff_expansion, reduction_ratio, num_layers)) |
| assert all([*map(lambda t: len(t) == 4, (dims, heads, ff_expansion, reduction_ratio, num_layers))]), 'only four stages are allowed, all keyword arguments must be either a single value or a tuple of 4 values' |
|
|
| self.mit = MiT( |
| channels = channels, |
| dims = dims, |
| heads = heads, |
| ff_expansion = ff_expansion, |
| reduction_ratio = reduction_ratio, |
| num_layers = num_layers |
| ) |
| |
| self.to_fused = [] |
| for i, dim in enumerate(dims): |
| to_fused = tf.keras.Sequential() |
| to_fused.add(Conv2d(decoder_dim, 1)) |
| to_fused.add(UpSampling2D(2 ** i)) |
| self.to_fused.append(to_fused) |
|
|
| self.to_segmentation = tf.keras.Sequential() |
| self.to_segmentation.add(Conv2d(decoder_dim, 1)) |
| self.to_segmentation.add(Conv2d(num_classes, 1)) |
|
|
| def __call__(self, x): |
| layer_outputs = self.mit(x, return_layer_outputs = True) |
|
|
| fused = [to_fused(output) for output, to_fused in zip(layer_outputs, self.to_fused)] |
| fused = tf.concat(fused, axis = -1) |
| return self.to_segmentation(fused) |
