Tutorial 3: Using a custom architecture¶
Out of the box, Slideflow includes support for 21 model architectures in the Tensorflow backend and 17 with the PyTorch backend. In this tutorial, we will demonstrate how to train a custom model architecture (ViT) in either backend.
Custom Tensorflow model¶
Any Tensorflow/Keras model (tf.keras.Model) can be trained in Slideflow by setting the model parameter of a slideflow.ModelParams object to a function which initalizes the model.
First, define the model in a file that can be imported. In this example, we will define a vision transformer (ViT) model in a file vit_tensorflow.py:
# From:
# https://github.com/ashishpatel26/Vision-Transformer-Keras-Tensorflow-Pytorch-Examples/blob/main/Vision_Transformer_with_tf2.ipynb
import math
import six
import tensorflow as tf
from einops.layers.tensorflow import Rearrange
def gelu(x):
"""Gaussian Error Linear Unit.
This is a smoother version of the RELU.
Original paper: https://arxiv.org/abs/1606.08415
Args:
x: float Tensor to perform activation.
Returns:
`x` with the GELU activation applied.
"""
cdf = 0.5 * (1.0 + tf.tanh(
(math.sqrt(2 / math.pi) * (x + 0.044715 * tf.pow(x, 3)))))
return x * cdf
def get_activation(identifier):
"""Maps a identifier to a Python function, e.g., "relu" => `tf.nn.relu`.
It checks string first and if it is one of customized activation not in TF,
the corresponding activation will be returned. For non-customized activation
names and callable identifiers, always fallback to tf.keras.activations.get.
Args:
identifier: String name of the activation function or callable.
Returns:
A Python function corresponding to the activation function.
"""
if isinstance(identifier, six.string_types):
name_to_fn = {"gelu": gelu}
identifier = str(identifier).lower()
if identifier in name_to_fn:
return tf.keras.activations.get(name_to_fn[identifier])
return tf.keras.activations.get(identifier)
class Residual(tf.keras.Model):
def __init__(self, fn):
super().__init__()
self.fn = fn
def call(self, x):
return self.fn(x) + x
class PreNorm(tf.keras.Model):
def __init__(self, dim, fn):
super().__init__()
self.norm = tf.keras.layers.LayerNormalization(epsilon=1e-5)
self.fn = fn
def call(self, x):
return self.fn(self.norm(x))
class FeedForward(tf.keras.Model):
def __init__(self, dim, hidden_dim):
super().__init__()
self.net = tf.keras.Sequential([
tf.keras.layers.Dense(
hidden_dim,
activation=get_activation('gelu')
),
tf.keras.layers.Dense(dim)]
)
def call(self, x):
return self.net(x)
class Attention(tf.keras.Model):
def __init__(self, dim, heads=8):
super().__init__()
self.heads = heads
self.scale = dim ** -0.5
self.to_qkv = tf.keras.layers.Dense(dim * 3, use_bias=False)
self.to_out = tf.keras.layers.Dense(dim)
self.rearrange_qkv = Rearrange(
'b n (qkv h d) -> qkv b h n d',
qkv=3,
h=self.heads
)
self.rearrange_out = Rearrange('b h n d -> b n (h d)')
def call(self, x):
qkv = self.to_qkv(x)
qkv = self.rearrange_qkv(qkv)
q = qkv[0]
k = qkv[1]
v = qkv[2]
dots = tf.einsum('bhid,bhjd->bhij', q, k) * self.scale
attn = tf.nn.softmax(dots, axis=-1)
out = tf.einsum('bhij,bhjd->bhid', attn, v)
out = self.rearrange_out(out)
out = self.to_out(out)
return out
class Transformer(tf.keras.Model):
def __init__(self, dim, depth, heads, mlp_dim):
super().__init__()
layers = []
for _ in range(depth):
layers.extend([
Residual(PreNorm(dim, Attention(dim, heads=heads))),
Residual(PreNorm(dim, FeedForward(dim, mlp_dim)))
])
self.net = tf.keras.Sequential(layers)
def call(self, x):
return self.net(x)
class ViT(tf.keras.Model):
def __init__(self, *, image_size, patch_size, num_classes,
dim, depth, heads, mlp_dim):
super().__init__()
if not image_size % patch_size == 0:
raise ValueError('image dimensions must be divisible by the '
'patch size')
num_patches = (image_size // patch_size) ** 2
self.patch_size = patch_size
self.dim = dim
self.pos_embedding = self.add_weight(
"position_embeddings",
shape=[num_patches + 1, dim],
initializer=tf.keras.initializers.RandomNormal(),
dtype=tf.float32
)
self.patch_to_embedding = tf.keras.layers.Dense(dim)
self.cls_token = self.add_weight(
"cls_token",
shape=[1, 1, dim],
initializer=tf.keras.initializers.RandomNormal(),
dtype=tf.float32
)
self.rearrange = Rearrange(
'b (h p1) (w p2) c -> b (h w) (p1 p2 c)',
p1=self.patch_size,
p2=self.patch_size
)
self.transformer = Transformer(dim, depth, heads, mlp_dim)
self.to_cls_token = tf.identity
self.mlp_head = tf.keras.Sequential([
tf.keras.layers.Dense(mlp_dim, activation=get_activation('gelu')),
tf.keras.layers.Dense(num_classes)
])
@tf.function
def call(self, img):
shapes = tf.shape(img)
x = self.rearrange(img)
x = self.patch_to_embedding(x)
cls_tokens = tf.broadcast_to(self.cls_token, (shapes[0], 1, self.dim))
x = tf.concat((cls_tokens, x), axis=1)
x += self.pos_embedding
x = self.transformer(x)
x = self.to_cls_token(x[:, 0])
return self.mlp_head(x)
Next, define a function that accepts any combination of the keyword arguments input_shape, include_top, pooling, and/or weights and returns an instanced model.
from vit_tensorflow import ViT
def vit_model(image_shape, **kwargs):
return ViT(
image_size=input_shape[0],
patch_size=23,
num_classes=1000,
dim=1024,
depth=6,
heads=16,
mlp_dim=2048
)
Then, create a slideflow.ModelParams object with your training parameters, setting the model argument equal to the function you just defined:
import slideflow as sf
from vit_tensorflow impport ViT
def vit_model(image_shape, **kwargs):
...
hp = ModelParams(
tile_px=299,
tile_um=302,
batch_size=32,
model=vit_model,
...
)
You can now train the model as described in Tutorial 1: Model training (simple).
Custom PyTorch model¶
The process is very similar when using PyTorch. In this example, instead of defining the architecture in a separate file, we will use an implementation of ViT available via PyPI:
pip3 install vit-pytorch
Next, define a function which accepts any combination of the keyword arguments image_size and/or pretrained and returns an instanced model.
import slideflow as sf
from vit_pytorch impport ViT
def vit_model(image_shape, **kwargs):
model = ViT(
image_size=image_size,
patch_size=23,
num_classes=1000,
dim=1024,
depth=6,
heads=16,
mlp_dim=2048,
dropout=0.1,
emb_dropout=0.1
)
model.out_features = 1000
return model
Finally, set the model argument of a slideflow.ModelParams object equal to this function:
import slideflow as sf
from vit_pytorch impport ViT
def vit_model(image_shape, **kwargs):
...
hp = ModelParams(
tile_px=299,
tile_um=302,
batch_size=32,
model=vit_model,
...
)
You can now train the model as described in Tutorial 1: Model training (simple).