Sorobot/Document-modify-Toolkit

1import paddle.nn as nn2 3# from weight_init import weight_init_4from .weight_init import weight_init_5 6class ResidualBlock(nn.Layer):7    """Residual Block with custom normalization."""8 9    def __init__(self, in_planes, planes, norm_fn="group", stride=1):10        super(ResidualBlock, self).__init__()11 12        self.conv1 = nn.Conv2D(in_planes, planes, 3, padding=1, stride=stride)13        self.conv2 = nn.Conv2D(planes, planes, 3, padding=1)14        self.relu = nn.ReLU()15 16        if norm_fn == "group":17            num_groups = planes // 818            self.norm1 = nn.GroupNorm(num_groups, planes)19            self.norm2 = nn.GroupNorm(num_groups, planes)20            if not stride == 1:21                self.norm3 = nn.GroupNorm(num_groups, planes)22        elif norm_fn == "batch":23            self.norm1 = nn.BatchNorm2D(planes)24            self.norm2 = nn.BatchNorm2D(planes)25            if not stride == 1:26                self.norm3 = nn.BatchNorm2D(planes)27        elif norm_fn == "instance":28            self.norm1 = nn.InstanceNorm2D(planes)29            self.norm2 = nn.InstanceNorm2D(planes)30            if not stride == 1:31                self.norm3 = nn.InstanceNorm2D(planes)32        elif norm_fn == "none":33            self.norm1 = nn.Sequential()34            self.norm2 = nn.Sequential()35            if not stride == 1:36                self.norm3 = nn.Sequential()37 38        if stride == 1:39            self.downsample = None40        else:41            self.downsample = nn.Sequential(42                nn.Conv2D(in_planes, planes, 1, stride=stride), self.norm343            )44 45    def forward(self, x):46        y = x47        y = self.relu(self.norm1(self.conv1(y)))48        y = self.relu(self.norm2(self.conv2(y)))49 50        if self.downsample is not None:51            x = self.downsample(x)52 53        return self.relu(x + y)54 55 56class BasicEncoder(nn.Layer):57    """Basic Encoder with custom normalization."""58 59    def __init__(self, output_dim=128, norm_fn="batch"):60        super(BasicEncoder, self).__init__()61 62        self.norm_fn = norm_fn63 64        if self.norm_fn == "group":65            self.norm1 = nn.GroupNorm(8, 64)66        elif self.norm_fn == "batch":67            self.norm1 = nn.BatchNorm2D(64)68        elif self.norm_fn == "instance":69            self.norm1 = nn.InstanceNorm2D(64)70        elif self.norm_fn == "none":71            self.norm1 = nn.Sequential()72 73        self.conv1 = nn.Conv2D(3, 64, 7, stride=2, padding=3)74        self.relu1 = nn.ReLU()75 76        self.in_planes = 6477        self.layer1 = self._make_layer(64, stride=1)78        self.layer2 = self._make_layer(128, stride=2)79        self.layer3 = self._make_layer(192, stride=2)80 81        self.conv2 = nn.Conv2D(192, output_dim, 1)82 83        for m in self.sublayers():84            if isinstance(m, nn.Conv2D):85                weight_init_(86                    m.weight, "kaiming_normal_", mode="fan_out", nonlinearity="relu"87                )88            elif isinstance(m, (nn.BatchNorm2D, nn.InstanceNorm2D, nn.GroupNorm)):89                weight_init_(m, "Constant", value=1, bias_value=0.0)90 91    def _make_layer(self, dim, stride=1):92        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)93        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)94        layers = layer1, layer295 96        self.in_planes = dim97        return nn.Sequential(*layers)98 99    def forward(self, x):100        x = self.conv1(x)101        x = self.norm1(x)102        x = self.relu1(x)103 104        x = self.layer1(x)105        x = self.layer2(x)106        x = self.layer3(x)107 108        x = self.conv2(x)109 110        return x111