mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
64 lines
2.5 KiB
64 lines
2.5 KiB
3 years ago
|
# modified from https://github.com/lucidrains/mlp-mixer-pytorch/blob/main/mlp_mixer_pytorch/mlp_mixer_pytorch.py
|
||
|
|
from functools import partial
|
||
|
|
from colossalai.context import ParallelMode
|
||
|
|
from colossalai.registry import MODELS
|
||
|
|
from torch import nn
|
||
|
|
from colossalai import nn as col_nn
|
||
|
|
from colossalai.nn.layer.parallel_3d._utils import get_depth_from_env
|
||
|
|
from einops.layers.torch import Rearrange, Reduce
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
'MLPMixer',
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
class PreNormResidual(nn.Module):
|
||
|
|
def __init__(self, dim, fn, depth_3d):
|
||
|
|
super().__init__()
|
||
|
|
self.fn = fn
|
||
|
|
self.norm = col_nn.LayerNorm3D(
|
||
|
|
dim, depth_3d, ParallelMode.PARALLEL_3D_INPUT, ParallelMode.PARALLEL_3D_WEIGHT)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
return self.fn(self.norm(x)) + x
|
||
|
|
|
||
|
|
|
||
|
|
def FeedForward(dim, depth_3d, expansion_factor=4, dropout=0., dense=None):
|
||
|
|
if dense is None:
|
||
|
|
dense = partial(col_nn.Linear3D, depth=depth_3d, input_parallel_mode=ParallelMode.PARALLEL_3D_INPUT,
|
||
|
|
weight_parallel_mode=ParallelMode.PARALLEL_3D_WEIGHT)
|
||
|
|
return nn.Sequential(
|
||
|
|
dense(dim, dim * expansion_factor),
|
||
|
|
nn.GELU(),
|
||
|
|
nn.Dropout(dropout),
|
||
|
|
dense(dim * expansion_factor, dim),
|
||
|
|
nn.Dropout(dropout)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@MODELS.register_module
|
||
|
|
def MLPMixer(image_size, channels, patch_size, dim, depth, num_classes, expansion_factor=4, dropout=0.):
|
||
|
|
assert (image_size % patch_size) == 0, 'image must be divisible by patch size'
|
||
|
|
num_patches = (image_size // patch_size) ** 2
|
||
|
|
depth_3d = get_depth_from_env()
|
||
|
|
linear = partial(col_nn.Linear3D, depth=depth_3d, input_parallel_mode=ParallelMode.PARALLEL_3D_INPUT,
|
||
|
|
weight_parallel_mode=ParallelMode.PARALLEL_3D_WEIGHT)
|
||
|
|
norm_layer = partial(col_nn.LayerNorm3D, depth=depth_3d, input_parallel_mode=ParallelMode.PARALLEL_3D_INPUT,
|
||
|
|
weight_parallel_mode=ParallelMode.PARALLEL_3D_WEIGHT)
|
||
|
|
chan_first, chan_last = partial(nn.Conv1d, kernel_size=1), linear
|
||
|
|
|
||
|
|
return nn.Sequential(
|
||
|
|
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)',
|
||
|
|
p1=patch_size, p2=patch_size),
|
||
|
|
linear((patch_size ** 2) * channels, dim),
|
||
|
|
*[nn.Sequential(
|
||
|
|
PreNormResidual(dim, FeedForward(
|
||
|
|
num_patches, expansion_factor, dropout, chan_first)),
|
||
|
|
PreNormResidual(dim, FeedForward(
|
||
|
|
dim, expansion_factor, dropout, chan_last))
|
||
|
|
) for _ in range(depth)],
|
||
|
|
norm_layer(dim),
|
||
|
|
Reduce('b n c -> b c', 'mean'),
|
||
|
|
linear(dim, num_classes)
|
||
|
|
)
|