fix format (#374)

model_zoo/mlp_mixer/parallel_3d/__init__.py

-from .mlp_mixer import *

model_zoo/mlp_mixer/parallel_3d/mlp_mixer.py

-# 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)
-    )