luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

prototype conditional pipeline

Browse Source
This commit is contained in:
luni
2022-06-10 06:59:58 +00:00
parent 03c22c1ba6
commit 31607cc754
24 changed files with 1093 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
convnext/LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) Meta Platforms, Inc. and affiliates.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
convnext/README.md
+42
View File
@@ -0,0 +1,42 @@
# ConvNeXt Tensorflow
This is unofficial tensorflow keras implementation of ConvNeXt.
Its based on official [PyTorch implementation](https://github.com/facebookresearch/ConvNeXt).
## Pre-trained Models
| name | resolution | pretrain | acc@1 | #params | FLOPs | model |
|:---:|:---:|:---:|:---:| :---:|:---:|:---:|
| `convnext_tiny_224` | 224x224 |ImageNet-1K | 82.1 | 28M | 4.5G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth) |
| `convnext_small_224` | 224x224 |ImageNet-1K | 83.1 | 50M | 8.7G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth) |
| `convnext_base_224` | 224x224 |ImageNet-21K-1K | 85.8 | 89M | 15.4G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth) |
| `convnext_base_384` | 384x384 |ImageNet-21K-1K | 86.8 | 89M | 45.0G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth) |
| `convnext_large_224` | 224x224 |ImageNet-21K-1K | 86.6 | 198M | 34.4G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth) |
| `convnext_large_384` | 384x384 |ImageNet-21K-1K | 87.5 | 198M | 101.0G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth) |
| `convnext_xlarge_224` | 224x224 |ImageNet-21K-1K | 87.0 | 350M | 60.9G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth) |
| `convnext_xlarge_384` | 384x384 |ImageNet-21K-1K | 87.8 | 350M | 179.0G | [github](https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth) |
## Note
I've ported only ImageNet-21K-1K weights for base, large and xlarge models.
If you want to convert another pretrained weight in official repo, you can refer to [this script](https://github.com/bamps53/convnext-tf/blob/master/convert_weights.py) or just let me know.
## Examples
```python
import tensorflow as tf
from models.convnext_tf import create_model
x = tf.zeros((1, 224, 224, 3), dtype=tf.float32)
model = create_model('convnext_tiny_224', input_shape=(224, 224), pretrained=True)
out = model(x) # (1, 1000)
model = create_model('convnext_tiny_224', input_shape=(224, 224), num_classes=1, pretrained=True)
out = model(x) # (1, 1)
model = create_model('convnext_tiny_224', input_shape=(224, 224), include_top=False, pretrained=True)
out = model(x) # (1, 16, 16, 768)
```
## Reference
https://github.com/facebookresearch/ConvNeXt
https://github.com/rishigami/Swin-Transformer-TF
convnext/convert_weights.py
+118
View File
@@ -0,0 +1,118 @@
import os
import cv2
import numpy as np
import tensorflow as tf
import timm
import torch
from models.convnext import * # just to register model
from models.convnext_tf import create_model
model_urls = [
'https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth',
'https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth',
'https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth',
'https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth',
'https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth',
'https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth',
'https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth',
'https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth',
]
def load_test_image(image_size=(224, 224)):
from skimage.data import chelsea
img = chelsea() # Chelsea the cat
img = cv2.resize(img, image_size)
img = img / 255
img = (img - np.array([0.485, 0.456, 0.406])) / \
np.array([0.229, 0.224, 0.225])
return img
def load_from_url(model, url):
checkpoint = torch.hub.load_state_dict_from_url(
url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint["model"])
return model
def get_pt_state_dict(pt_model):
pt_state_dict = {}
for n, p in pt_model.named_parameters():
pt_state_dict[n] = p.detach().numpy()
return pt_state_dict
def get_tf_state_dict(tf_model):
tf_state_dict = {}
for layer in tf_model.layers:
for var in layer.weights:
if ('dense' in var.name) and ('kernel' in var.name):
w = tf.transpose(var, [1, 0]).numpy()
elif ('depthwise_conv2d' in var.name) and ('kernel' in var.name):
w = tf.transpose(var, [2, 3, 0, 1]).numpy()
elif ('conv2d' in var.name) and ('kernel' in var.name):
w = tf.transpose(var, [3, 2, 0, 1]).numpy()
else:
w = var.numpy()
tf_state_dict[var.name] = w
return tf_state_dict
def convert_weights(tf_model, pt_state_dict, key_map):
for layer in tf_model.layers:
for var in layer.weights:
pt_key = key_map[var.name]
pt_w = pt_state_dict[pt_key]
if ('dense' in var.name) and ('kernel' in var.name):
tf_w = np.transpose(pt_w, [1, 0])
elif ('depthwise_conv2d' in var.name) and ('kernel' in var.name):
tf_w = np.transpose(pt_w, [2, 3, 0, 1])
elif ('conv2d' in var.name) and ('kernel' in var.name):
tf_w = np.transpose(pt_w, [2, 3, 1, 0])
else:
tf_w = pt_w
var.assign(tf_w)
def test_models(pt_model, tf_model):
img = load_test_image()
pt_x = torch.tensor(img[None, ]).permute(0, 3, 1, 2).to(torch.float32)
with torch.no_grad():
pt_y = torch.softmax(pt_model(pt_x), -1).numpy()
tf_x = tf.convert_to_tensor(img[None, ])
tf_y = tf.nn.softmax(tf_model(tf_x)).numpy()
np.testing.assert_allclose(pt_y, tf_y, rtol=1e-5)
def main():
os.makedirs('weights', exist_ok=True)
for model_url in model_urls:
file_name = os.path.basename(model_url)
model_name = '_'.join(file_name.split('_')[:2])
save_name = file_name.replace("pth", "h5")
if os.path.exists(f'weights/{save_name}'):
continue
pt_model = timm.create_model(model_name)
pt_model = load_from_url(pt_model, model_url)
pt_model.eval()
pt_state_dict = get_pt_state_dict(pt_model)
tf_model = create_model(model_name)
tf_state_dict = get_tf_state_dict(tf_model)
key_map = {t: p for t, p in zip(
tf_state_dict.keys(), pt_state_dict.keys())}
convert_weights(tf_model, pt_state_dict, key_map)
test_models(pt_model, tf_model)
print(f'successfully converted {model_name}!')
tf_model.save_weights(f'weights/{save_name}')
if __name__ == '__main__':
main()
convnext/models/convnext.py
+201
View File
@@ -0,0 +1,201 @@
# copied from https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import trunc_normal_, DropPath
from timm.models.registry import register_model
class Block(nn.Module):
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, drop_path=0., layer_scale_init_value=1e-6):
super().__init__()
self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
self.norm = LayerNorm(dim, eps=1e-6)
self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
self.act = nn.GELU()
self.pwconv2 = nn.Linear(4 * dim, dim)
self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
requires_grad=True) if layer_scale_init_value > 0 else None
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
def forward(self, x):
input = x
x = self.dwconv(x)
x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
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)
return x
class ConvNeXt(nn.Module):
r""" ConvNeXt
A PyTorch impl of : `A ConvNet for the 2020s` -
https://arxiv.org/pdf/2201.03545.pdf
Args:
in_chans (int): Number of input image channels. Default: 3
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]
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, in_chans=3, num_classes=1000,
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0.,
layer_scale_init_value=1e-6, head_init_scale=1.,
):
super().__init__()
self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
stem = nn.Sequential(
nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
)
self.downsample_layers.append(stem)
for i in range(3):
downsample_layer = nn.Sequential(
LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
)
self.downsample_layers.append(downsample_layer)
self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(4):
stage = nn.Sequential(
*[Block(dim=dims[i], drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
)
self.stages.append(stage)
cur += depths[i]
self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
self.head = nn.Linear(dims[-1], num_classes)
self.apply(self._init_weights)
self.head.weight.data.mul_(head_init_scale)
self.head.bias.data.mul_(head_init_scale)
def _init_weights(self, m):
if isinstance(m, (nn.Conv2d, nn.Linear)):
trunc_normal_(m.weight, std=.02)
nn.init.constant_(m.bias, 0)
def forward_features(self, x):
for i in range(4):
x = self.downsample_layers[i](x)
x = self.stages[i](x)
return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C)
def forward(self, x):
x = self.forward_features(x)
x = self.head(x)
return x
class LayerNorm(nn.Module):
r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
with shape (batch_size, channels, height, width).
"""
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
super().__init__()
self.weight = nn.Parameter(torch.ones(normalized_shape))
self.bias = nn.Parameter(torch.zeros(normalized_shape))
self.eps = eps
self.data_format = data_format
if self.data_format not in ["channels_last", "channels_first"]:
raise NotImplementedError
self.normalized_shape = (normalized_shape, )
def forward(self, x):
if self.data_format == "channels_last":
return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
elif self.data_format == "channels_first":
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None, None] * x + self.bias[:, None, None]
return x
model_urls = {
"convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
"convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
"convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
"convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
"convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
"convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
"convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
}
@register_model
def convnext_tiny(pretrained=False, **kwargs):
model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
if pretrained:
url = model_urls['convnext_tiny_1k']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint["model"])
return model
@register_model
def convnext_small(pretrained=False, **kwargs):
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
if pretrained:
url = model_urls['convnext_small_1k']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint["model"])
return model
@register_model
def convnext_base(pretrained=False, in_22k=False, **kwargs):
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
if pretrained:
url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint["model"])
return model
@register_model
def convnext_large(pretrained=False, in_22k=False, **kwargs):
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
if pretrained:
url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint["model"])
return model
@register_model
def convnext_xlarge(pretrained=False, in_22k=False, **kwargs):
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
if pretrained:
url = model_urls['convnext_xlarge_22k'] if in_22k else model_urls['convnext_xlarge_1k']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint["model"])
return model
convnext/models/convnext_tf.py
+207
View File
@@ -0,0 +1,207 @@
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import (Conv2D, Dense, DepthwiseConv2D,
GlobalAveragePooling2D, Layer,
LayerNormalization)
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
# Compute keep_prob
keep_prob = 1.0 - drop_prob
# 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
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):
super().__init__()
self.drop_prob = drop_prob
def call(self, x, training=None):
return drop_path(x, self.drop_prob, training)
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, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
super().__init__()
self.dwconv = DepthwiseConv2D(
kernel_size=7, 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
def call(self, x):
input = x
x = self.dwconv(x)
# x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
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)
return 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,
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.,
):
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)]
)
self.downsample_layers.append(stem)
for i in range(3):
downsample_layer = tf.keras.Sequential([
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(4):
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]
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
def forward_features(self, x):
for i in range(4):
x = self.downsample_layers[i](x)
x = self.stages[i](x)
return x
def call(self, x):
x = self.forward_features(x)
if self.include_top:
x = self.avg(x)
x = self.norm(x)
x = self.head(x)
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]
),
)
def create_model(model_name='convnext_tiny_1k', input_shape=(224, 224), num_classes=1000, include_top=True, pretrained=True, use_tpu=False, **kwargs):
cfg = model_configs['_'.join(model_name.split('_')[:2])]
net = ConvNeXt(num_classes, cfg['depths'],
cfg['dims'], include_top, **kwargs)
net(tf.keras.Input(shape=(*input_shape, 3)))
if pretrained is True:
url = model_urls[model_name]
pretrained_ckpt = tf.keras.utils.get_file(
f'{model_name}.h5', url, untar=False)
if use_tpu:
load_locally = tf.saved_model.LoadOptions(
experimental_io_device='/job:localhost')
net.load_weights(
pretrained_ckpt, options=load_locally, skip_mismatch=True)
else:
net.load_weights(pretrained_ckpt, skip_mismatch=True, by_name=True)
return net