static encoder runs, but doesn't work yet
This commit is contained in:
+37
-149
@@ -1,134 +1,21 @@
|
||||
import tensorflow as tf
|
||||
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
|
||||
import torch
|
||||
from torch import nn
|
||||
from models.convnext import ConvNeXt, CNBlock, CNBlockConfig
|
||||
from typing import Any, Callable, List, Optional, Sequence
|
||||
|
||||
class Bottleneck(tf.keras.Model):
|
||||
|
||||
class VariationalEncoder(nn.Module):
|
||||
|
||||
def __init__(self, dim, latent_dim, kernel_size=7):
|
||||
super(Bottleneck, self).__init__()
|
||||
|
||||
self.bottleneck = tf.keras.Sequential(
|
||||
[Block(dim=dim, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.dense = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
x = self.flatten(x)
|
||||
x = self.dense(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class VariationalEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, kernel_size=7):
|
||||
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7):
|
||||
super(VariationalEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.mean = tf.keras.layers.Dense(latent_dim)
|
||||
self.log_var = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.flatten(x)
|
||||
mean = self.mean(x)
|
||||
log_var = self.log_var(x)
|
||||
|
||||
return mean, log_var
|
||||
|
||||
|
||||
class AEEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
|
||||
super(AEEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.bottleneck(x, training=training)
|
||||
|
||||
if len(inputs) == 2:
|
||||
return x, inputs[1]
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
|
||||
self.reshape = tf.keras.layers.Reshape(input_size)
|
||||
self.upscale = upscale
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.dense(x)
|
||||
x = self.reshape(x)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
if self.upscale:
|
||||
x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
|
||||
if len(inputs) == 2:
|
||||
if len(mask.shape) == 2:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class VariationalEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7):
|
||||
super(VariationalEncoderNew, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
|
||||
self.mean_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.var_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.flatten_mean = tf.keras.layers.Flatten()
|
||||
self.flatten_var = tf.keras.layers.Flatten()
|
||||
self.mean_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
|
||||
self.var_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
|
||||
|
||||
self.flatten_mean = nn.Flatten()
|
||||
self.flatten_var = nn.Flatten()
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
@@ -144,23 +31,27 @@ class VariationalEncoderNew(tf.keras.Model):
|
||||
return mean, log_var
|
||||
|
||||
|
||||
|
||||
class AEEncoderNew(tf.keras.Model):
|
||||
class AEEncoder(nn.Module):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
|
||||
super(AEEncoderNew, self).__init__()
|
||||
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', device='cuda'):
|
||||
super(AEEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.act = tf.keras.layers.Activation(activation)
|
||||
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device)
|
||||
if activation == 'tanh':
|
||||
self.act = nn.Tanh()
|
||||
elif activation == 'sigmoid':
|
||||
self.act = nn.Sigmoid()
|
||||
else:
|
||||
raise ValueError('Activation function not supported')
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
def forward(self, inputs):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.encoder(x)
|
||||
x = self.act(x)
|
||||
|
||||
if len(inputs) == 2:
|
||||
@@ -169,39 +60,36 @@ class AEEncoderNew(tf.keras.Model):
|
||||
return x
|
||||
|
||||
|
||||
class DecoderNew(tf.keras.Model):
|
||||
class Decoder(nn.Module):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, is_label=False):
|
||||
super(DecoderNew, self).__init__()
|
||||
def __init__(self, block_setting: CNBlockConfig, kernel_size=7, is_label=False, nr_of_classes=None, device='cuda'):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
self.is_label = is_label
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
|
||||
self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_classes, device=device, reversed=True)
|
||||
|
||||
self.softmax = nn.Softmax(dim=1)
|
||||
|
||||
def call(self, inputs, training):
|
||||
def forward(self, inputs):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.decoder(x, training=training)
|
||||
x = self.decoder(x)
|
||||
|
||||
if self.is_label:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
x = x * mask[:,:,:,None]
|
||||
|
||||
background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
background = torch.ones((mask.shape[0],512,256,1), device=mask.get_device()) * (1 - mask[:,:,:,None])
|
||||
x = torch.concat((background, x), axis=-1)
|
||||
x = self.softmax(x)
|
||||
|
||||
if len(inputs) == 2:
|
||||
if len(mask.shape) == 2:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
x = x * mask[None,None,:,:]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
x = x * mask[:,None,:,:]
|
||||
|
||||
return x
|
||||
+10
-27
@@ -1,38 +1,21 @@
|
||||
import tensorflow as tf
|
||||
from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
|
||||
import torch
|
||||
from torch import nn
|
||||
from models.base import AEEncoder, Decoder
|
||||
|
||||
|
||||
class AutoEncoder(tf.keras.Model):
|
||||
class AutoEncoder(nn.Module):
|
||||
|
||||
def __init__(self, config):
|
||||
super(AutoEncoder, self).__init__()
|
||||
|
||||
self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
|
||||
self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
|
||||
self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], device=config['device'])
|
||||
self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'])
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
def forward(self, input):
|
||||
x = input
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class AutoEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(AutoEncoderNew, self).__init__()
|
||||
|
||||
self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
|
||||
self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.decoder(x, training=training)
|
||||
x = self.encoder(x)
|
||||
x = self.decoder(x)
|
||||
|
||||
return x
|
||||
+394
-276
@@ -1,306 +1,424 @@
|
||||
from unittest import skip
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
|
||||
GlobalAveragePooling2D, Layer,
|
||||
LayerNormalization, Flatten, Reshape, UpSampling2D)
|
||||
# taken from torchvision.models.convnext
|
||||
|
||||
def drop_path(inputs, drop_prob, is_training):
|
||||
# borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
|
||||
if (not is_training) or (drop_prob == 0.):
|
||||
return inputs
|
||||
from functools import partial
|
||||
from typing import Any, Callable, List, Optional, Sequence
|
||||
|
||||
# Compute keep_prob
|
||||
keep_prob = 1.0 - drop_prob
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
from torch.nn import functional as F
|
||||
|
||||
# Compute drop_connect tensor
|
||||
random_tensor = keep_prob
|
||||
shape = (tf.shape(inputs)[0],) + (1,) * \
|
||||
(len(tf.shape(inputs)) - 1)
|
||||
random_tensor += tf.random.uniform(shape, dtype=inputs.dtype)
|
||||
binary_tensor = tf.floor(random_tensor)
|
||||
output = tf.math.divide(inputs, keep_prob) * binary_tensor
|
||||
return output
|
||||
from torchvision.ops.misc import Conv2dNormActivation, Permute
|
||||
# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation
|
||||
# from torchvision.ops import Permute
|
||||
from torchvision.ops.stochastic_depth import StochasticDepth
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
|
||||
class DropPath(tf.keras.layers.Layer):
|
||||
# borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
|
||||
def __init__(self, drop_prob=None):
|
||||
__all__ = [
|
||||
"ConvNeXt",
|
||||
"ConvNeXt_Tiny_Weights",
|
||||
"ConvNeXt_Small_Weights",
|
||||
"ConvNeXt_Base_Weights",
|
||||
"ConvNeXt_Large_Weights",
|
||||
"convnext_tiny",
|
||||
"convnext_small",
|
||||
"convnext_base",
|
||||
"convnext_large",
|
||||
]
|
||||
|
||||
|
||||
class LayerNorm2d(nn.LayerNorm):
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = x.permute(0, 2, 3, 1)
|
||||
x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
|
||||
x = x.permute(0, 3, 1, 2)
|
||||
return x
|
||||
|
||||
|
||||
class CNBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
layer_scale: float = 1e-6,
|
||||
stochastic_depth_prob: float = 0.0,
|
||||
kernel_size: int = 7,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.drop_prob = drop_prob
|
||||
if norm_layer is None:
|
||||
norm_layer = partial(nn.LayerNorm, eps=1e-6)
|
||||
|
||||
def call(self, x, training):
|
||||
return drop_path(x, self.drop_prob, training)
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(dim, dim, kernel_size=kernel_size, padding='same', groups=dim, bias=True),
|
||||
Permute([0, 2, 3, 1]),
|
||||
norm_layer(dim),
|
||||
nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
|
||||
nn.GELU(),
|
||||
nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
|
||||
Permute([0, 3, 1, 2]),
|
||||
)
|
||||
self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
|
||||
self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
result = self.layer_scale * self.block(input)
|
||||
result = self.stochastic_depth(result)
|
||||
result += input
|
||||
return result
|
||||
|
||||
|
||||
class Block(Layer):
|
||||
r""" ConvNeXt Block. There are two equivalent implementations:
|
||||
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
|
||||
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
|
||||
We use (2) as we find it slightly faster in PyTorch
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
drop_path (float): Stochastic depth rate. Default: 0.0
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
"""
|
||||
|
||||
def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
|
||||
class ConvNeXtBase(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block_setting: List[CNBlockConfig],
|
||||
stochastic_depth_prob: float = 0.0,
|
||||
layer_scale: float = 1e-6,
|
||||
block: Optional[Callable[..., nn.Module]] = None,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
device: str = 'cuda',
|
||||
kernel_size: int = 7,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.dwconv = DepthwiseConv2D(
|
||||
kernel_size=kernel_size, padding='same') # depthwise conv
|
||||
self.norm = LayerNormalization(epsilon=1e-6)
|
||||
# pointwise/1x1 convs, implemented with linear layers
|
||||
self.pwconv1 = Dense(4 * dim)
|
||||
self.act = tf.keras.activations.gelu
|
||||
self.pwconv2 = Dense(dim)
|
||||
self.drop_path = DropPath(drop_path)
|
||||
self.dim = dim
|
||||
self.layer_scale_init_value = layer_scale_init_value
|
||||
self.prefix = prefix
|
||||
|
||||
def build(self, input_shape):
|
||||
self.gamma = tf.Variable(
|
||||
initial_value=self.layer_scale_init_value * tf.ones((self.dim)),
|
||||
trainable=True,
|
||||
name=f'{self.prefix}/gamma')
|
||||
self.built = True
|
||||
if not block_setting:
|
||||
raise ValueError("The block_setting should not be empty")
|
||||
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
|
||||
raise TypeError("The block_setting should be List[CNBlockConfig]")
|
||||
|
||||
def call(self, x, training):
|
||||
input = x
|
||||
x = self.dwconv(x, training=training)
|
||||
# x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
|
||||
x = self.norm(x, training=training)
|
||||
x = self.pwconv1(x, training=training)
|
||||
if block is None:
|
||||
block = CNBlock
|
||||
|
||||
if norm_layer is None:
|
||||
norm_layer = partial(LayerNorm2d, eps=1e-6)
|
||||
|
||||
# self.permute_cf = Permute([0, 3, 1, 2])
|
||||
# self.permute_cl = Permute([0, 2, 3, 1])
|
||||
|
||||
self.stages: List[nn.Module] = []
|
||||
self.downsamples: List[nn.Module] = []
|
||||
|
||||
# Stem
|
||||
firstconv_output_channels = block_setting[0].input_channels
|
||||
self.stem = Conv2dNormActivation(
|
||||
3,
|
||||
firstconv_output_channels,
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0,
|
||||
norm_layer=norm_layer,
|
||||
activation_layer=None,
|
||||
bias=True,
|
||||
).to(device)
|
||||
|
||||
|
||||
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
|
||||
stage_block_id = 0
|
||||
for cnf in block_setting:
|
||||
# Bottlenecks
|
||||
stage: List[nn.Module] = []
|
||||
for _ in range(cnf.num_layers):
|
||||
# adjust stochastic depth probability based on the depth of the stage block
|
||||
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
|
||||
stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
|
||||
stage_block_id += 1
|
||||
self.stages.append(nn.Sequential(*stage).to(device))
|
||||
if cnf.out_channels is not None:
|
||||
# Downsampling
|
||||
self.downsamples.append(
|
||||
nn.Sequential(
|
||||
norm_layer(cnf.input_channels),
|
||||
nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
||||
).to(device)
|
||||
)
|
||||
|
||||
self.act = nn.Tanh()
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def _forward_impl(self, x: Tensor) -> Tensor:
|
||||
# x = self.permute_cf(x)
|
||||
x = self.stem(x)
|
||||
# features: List[Tensor] = []
|
||||
for i in range(len(self.stages) - 1):
|
||||
x = self.stages[i](x)
|
||||
# features.append(x)
|
||||
x = self.downsamples[i](x)
|
||||
x = self.stages[-1](x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x, training=training)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
# x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
|
||||
|
||||
x = input + self.drop_path(x, training=training)
|
||||
# x = self.permute_cl(x)
|
||||
return x
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self._forward_impl(x)
|
||||
|
||||
class ConvNeXt(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes=1000, kernel_size=7,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
):
|
||||
class ConvNeXtReverseBase2(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block_setting: List[CNBlockConfig],
|
||||
stochastic_depth_prob: float = 0.0,
|
||||
layer_scale: float = 1e-6,
|
||||
block: Optional[Callable[..., nn.Module]] = None,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
device: str = 'cuda',
|
||||
use_skip: bool = False,
|
||||
is_label: bool = False,
|
||||
nr_of_labels: Optional[int] = None,
|
||||
kernel_size: int = 7,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.include_top = include_top
|
||||
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
# stem = tf.keras.Sequential([
|
||||
# Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
# LayerNormalization(epsilon=1e-6)]
|
||||
# )
|
||||
stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)]
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
# downsample_layer = tf.keras.Sequential([
|
||||
# LayerNormalization(epsilon=1e-6),
|
||||
# Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
# )
|
||||
downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
self.use_skip = use_skip
|
||||
|
||||
if self.include_top:
|
||||
self.avg = GlobalAveragePooling2D()
|
||||
self.norm = LayerNormalization(epsilon=1e-6) # final norm layer
|
||||
self.head = Dense(num_classes)
|
||||
else:
|
||||
self.avg = None
|
||||
self.norm = None
|
||||
self.head = None
|
||||
if not block_setting:
|
||||
raise ValueError("The block_setting should not be empty")
|
||||
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
|
||||
raise TypeError("The block_setting should be List[CNBlockConfig]")
|
||||
|
||||
def forward_features(self, x, training):
|
||||
for i in range(len(self.downsample_layers)):
|
||||
# x = self.downsample_layers[i](x)
|
||||
x = self.downsample_layers[i][0](x, training=training)
|
||||
x = self.downsample_layers[i][1](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
if block is None:
|
||||
block = CNBlock
|
||||
|
||||
if norm_layer is None:
|
||||
norm_layer = partial(LayerNorm2d, eps=1e-6)
|
||||
|
||||
# self.permute_cf = Permute([0, 3, 1, 2])
|
||||
# self.permute_cl = Permute([0, 2, 3, 1])
|
||||
|
||||
# Stem
|
||||
# firstconv_output_channels = block_setting[0].input_channels
|
||||
# self.first = nn.Sequential(
|
||||
# nn.ConvTranspose2d(
|
||||
# firstconv_output_channels,
|
||||
# firstconv_output_channels,
|
||||
# kernel_size=2,
|
||||
# stride=2,
|
||||
# padding=0,
|
||||
# ),
|
||||
# norm_layer(firstconv_output_channels),
|
||||
# ).to(device)
|
||||
self.first = nn.Sequential(
|
||||
nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
|
||||
norm_layer(block_setting[0].out_channels),
|
||||
).to(device)
|
||||
|
||||
|
||||
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
|
||||
stage_block_id = 0
|
||||
self.stages: List[nn.Module] = []
|
||||
self.upsamples: List[nn.Module] = []
|
||||
for cnf in block_setting[1:]:
|
||||
if cnf.input_channels is not None:
|
||||
# Bottlenecks
|
||||
stage: List[nn.Module] = []
|
||||
for _ in range(cnf.num_layers):
|
||||
# adjust stochastic depth probability based on the depth of the stage block
|
||||
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
|
||||
stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
|
||||
stage_block_id += 1
|
||||
self.stages.append(nn.Sequential(*stage).to(device))
|
||||
else:
|
||||
self.stages.append(nn.Identity().to(device))
|
||||
|
||||
if cnf.out_channels is not None:
|
||||
# Upsampling
|
||||
self.upsamples.append(
|
||||
nn.Sequential(
|
||||
norm_layer(cnf.input_channels),
|
||||
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
|
||||
).to(device)
|
||||
)
|
||||
|
||||
self.last = nn.Conv2d(
|
||||
block_setting[-1].input_channels,
|
||||
3 if not is_label else nr_of_labels,
|
||||
kernel_size=3,
|
||||
padding='same',
|
||||
).to(device)
|
||||
self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1)
|
||||
|
||||
# self.last = nn.Sequential(
|
||||
# Permute([0, 2, 3, 1]),
|
||||
# nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
|
||||
# Permute([0, 3, 1, 2]),
|
||||
# ).to(device)
|
||||
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
|
||||
def _forward_impl(self, x) -> Tensor:
|
||||
# input, features = x
|
||||
# x = self.permute_cf(x)
|
||||
|
||||
x = self.first(x)
|
||||
for i in range(len(self.stages)):
|
||||
# if self.use_skip:
|
||||
# x = self.stages[i](torch.concat((x, features[2-i]), dim=1))
|
||||
# else:
|
||||
x = self.stages[i](x)
|
||||
x = self.upsamples[i](x)
|
||||
|
||||
# x = self.stages[-2](x)
|
||||
# x = self.upsamples[-1](x)
|
||||
|
||||
# x = self.stages[-1](x)
|
||||
x = self.last(x)
|
||||
x = self.activation(x)
|
||||
# x = self.permute_cl(x)
|
||||
|
||||
return x
|
||||
|
||||
def call(self, x, training):
|
||||
x = self.forward_features(x, training=training)
|
||||
def forward(self, x) -> Tensor:
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
class ConvNeXtReverseBase(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block_setting: List[CNBlockConfig],
|
||||
stochastic_depth_prob: float = 0.0,
|
||||
layer_scale: float = 1e-6,
|
||||
block: Optional[Callable[..., nn.Module]] = None,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
if not block_setting:
|
||||
raise ValueError("The block_setting should not be empty")
|
||||
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
|
||||
raise TypeError("The block_setting should be List[CNBlockConfig]")
|
||||
|
||||
if block is None:
|
||||
block = CNBlock
|
||||
|
||||
if norm_layer is None:
|
||||
norm_layer = partial(LayerNorm2d, eps=1e-6)
|
||||
|
||||
layers: List[nn.Module] = []
|
||||
|
||||
|
||||
# add bottleneck which will be executed at the end, to process after the final larger upsampling
|
||||
|
||||
|
||||
|
||||
# Stem
|
||||
firstconv_output_channels = block_setting[0].input_channels
|
||||
|
||||
layers.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True),
|
||||
|
||||
)
|
||||
)
|
||||
|
||||
stage: List[nn.Module] = []
|
||||
for _ in range(3):
|
||||
stage.append(block(16, layer_scale, 0.0))
|
||||
layers.append(nn.Sequential(*stage))
|
||||
|
||||
layers.append(
|
||||
# Convtr NormActivation(
|
||||
# firstconv_output_channels,
|
||||
# 3,
|
||||
# kernel_size=4,
|
||||
# stride=4,
|
||||
# padding=0,
|
||||
# norm_layer=None,
|
||||
# activation_layer=None,
|
||||
# bias=True,
|
||||
# )
|
||||
nn.Sequential(
|
||||
nn.ConvTranspose2d(
|
||||
firstconv_output_channels,
|
||||
16,
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
|
||||
stage_block_id = 0
|
||||
for cnf in block_setting:
|
||||
# Bottlenecks
|
||||
stage: List[nn.Module] = []
|
||||
for _ in range(cnf.num_layers):
|
||||
# adjust stochastic depth probability based on the depth of the stage block
|
||||
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
|
||||
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
|
||||
stage_block_id += 1
|
||||
layers.append(nn.Sequential(*stage))
|
||||
if cnf.out_channels is not None:
|
||||
# Upsampling
|
||||
layers.append(
|
||||
nn.Sequential(
|
||||
norm_layer(cnf.out_channels),
|
||||
nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
|
||||
)
|
||||
)
|
||||
|
||||
self.features = nn.Sequential(*reversed(layers))
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def _forward_impl(self, x: Tensor) -> Tensor:
|
||||
x = self.features(x)
|
||||
# x = torch.tanh(x)
|
||||
return x
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
def _convnext(
|
||||
block_setting: List[CNBlockConfig],
|
||||
stochastic_depth_prob: float,
|
||||
reversed: bool,
|
||||
is_label: bool,
|
||||
nr_of_labels: Optional[int],
|
||||
kernel_size: int,
|
||||
device='cuda',
|
||||
**kwargs: Any,
|
||||
) -> ConvNeXtBase:
|
||||
|
||||
if reversed:
|
||||
model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
|
||||
else:
|
||||
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)
|
||||
|
||||
|
||||
return model
|
||||
|
||||
|
||||
class ConvNeXt(nn.Module):
|
||||
def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, device='cuda', **kwargs) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = self.model(x)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtReverse(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
):
|
||||
super().__init__()
|
||||
self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
stem = tf.keras.Sequential([
|
||||
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
# UpSampling2D((4, 4), interpolation='nearest'),
|
||||
# Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
|
||||
# Dense(dims[0]),
|
||||
LayerNormalization(epsilon=1e-6)]
|
||||
)
|
||||
self.upsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
downsample_layer = tf.keras.Sequential([
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
|
||||
# UpSampling2D((2,2), interpolation='nearest'),
|
||||
# Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
|
||||
# Dense(dims[i+1]),
|
||||
])
|
||||
self.upsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
self.last = Dense(3, activation='tanh')
|
||||
|
||||
def call(self, x, training):
|
||||
for i in reversed(range(len(self.upsample_layers))):
|
||||
x = self.upsample_layers[i](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
x = self.last(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtReverseLabels(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
nr_of_classes=7):
|
||||
super().__init__()
|
||||
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
stem = tf.keras.Sequential([
|
||||
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
LayerNormalization(epsilon=1e-6)]
|
||||
)
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
downsample_layer = tf.keras.Sequential([
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
)
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
# self.last = Dense(nr_of_classes, activation='softmax')
|
||||
self.last = Dense(nr_of_classes - 1)
|
||||
|
||||
def call(self, x, training):
|
||||
for i in reversed(range(len(self.downsample_layers))):
|
||||
x = self.downsample_layers[i](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
x = self.last(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
model_urls = {
|
||||
"convnext_tiny_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_tiny_1k_224_ema.h5",
|
||||
"convnext_small_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_small_1k_224_ema.h5",
|
||||
"convnext_base_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_224.h5",
|
||||
"convnext_base_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_384.h5",
|
||||
"convnext_large_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_224.h5",
|
||||
"convnext_large_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_384.h5",
|
||||
"convnext_xlarge_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_224_ema.h5",
|
||||
"convnext_xlarge_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_384_ema.h5",
|
||||
}
|
||||
|
||||
|
||||
model_configs = dict(
|
||||
convnext_tiny=dict(
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[8, 16, 32, 512]
|
||||
),
|
||||
convnext_small=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[96, 192, 384, 768]
|
||||
),
|
||||
convnext_base=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[128, 256, 512, 1024]
|
||||
),
|
||||
convnext_large=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[192, 384, 768, 1536]
|
||||
),
|
||||
convnext_xlarge=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[256, 512, 1024, 2048]
|
||||
),
|
||||
)
|
||||
if __name__ == '__main__':
|
||||
model = ConvNeXt().to('cuda')
|
||||
test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
|
||||
|
||||
pass
|
||||
Reference in New Issue
Block a user