Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| # Conv Layer | |
| class ConvLayer(nn.Module): | |
| def __init__(self, in_channels, out_channels, kernel_size, stride, groups=1): | |
| super(ConvLayer, self).__init__() | |
| paddings = kernel_size // 2 | |
| self.reflection_pad = nn.ReflectionPad2d(paddings) | |
| self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, groups=groups) #, padding) | |
| # self.in_d = nn.InstanceNorm2d(out_channels, affine=True) | |
| def forward(self, x): | |
| out = self.reflection_pad(x) | |
| out = self.conv2d(out) | |
| return out | |
| class ConvLayer_dpws(nn.Module): | |
| def __init__(self, in_channels, out_channels, kernel_size, stride): | |
| super(ConvLayer_dpws, self).__init__() | |
| self.conv1 = ConvLayer(in_channels, in_channels, kernel_size, stride=stride, groups=in_channels) | |
| self.in_1d = nn.InstanceNorm2d(in_channels, affine=True) | |
| self.conv2 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1) | |
| self.in_2d = nn.InstanceNorm2d(out_channels, affine=True) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| out = self.in_1d(self.conv1(x)) | |
| out = self.relu(self.in_2d(self.conv2(out))) | |
| return out | |
| class ConvLayer_dpws_last(nn.Module): | |
| def __init__(self, in_channels, out_channels, kernel_size, stride): | |
| super(ConvLayer_dpws_last, self).__init__() | |
| self.conv1 = ConvLayer(in_channels, in_channels, kernel_size, stride=stride, groups=in_channels) | |
| self.in_1d = nn.InstanceNorm2d(in_channels, affine=True) | |
| self.conv2 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1) | |
| # self.in_2d = nn.InstanceNorm2d(out_channels, affine=True) | |
| # self.relu = nn.ReLU() | |
| def forward(self, x): | |
| out = self.in_1d(self.conv1(x)) | |
| # out = self.relu(self.in_2d(self.conv2(out))) | |
| out = self.conv2(out) | |
| return out | |
| # Upsample Conv Layer | |
| class UpsampleConvLayer(nn.Module): | |
| def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None): | |
| super(UpsampleConvLayer, self).__init__() | |
| self.upsample = upsample | |
| if upsample: | |
| self.upsample = nn.Upsample(scale_factor=upsample, mode='nearest') | |
| reflection_padding = kernel_size // 2 | |
| self.reflection_pad = nn.ReflectionPad2d(reflection_padding) | |
| self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride) | |
| # self.in_d = nn.InstanceNorm2d(out_channels, affine=True) | |
| def forward(self, x): | |
| if self.upsample: | |
| x = self.upsample(x) | |
| out = self.reflection_pad(x) | |
| out = self.conv2d(out) | |
| return out | |
| class UpsampleConvLayer_dpws(nn.Module): | |
| def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None): | |
| super(UpsampleConvLayer_dpws, self).__init__() | |
| self.upsample = upsample | |
| if upsample: | |
| self.upsample = nn.Upsample(scale_factor=upsample, mode='nearest') | |
| self.conv1 = ConvLayer(in_channels, in_channels, kernel_size, stride, groups=in_channels) | |
| self.in1 = nn.InstanceNorm2d(in_channels, affine=True ) | |
| self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1) | |
| self.in2 = nn.InstanceNorm2d(out_channels, affine=True ) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| if self.upsample: | |
| x = self.upsample(x) | |
| # out = self.reflection_pad(x) | |
| # out = self.conv2d(out) | |
| out = self.relu(self.in1(self.conv1(x))) | |
| out = self.in2(self.conv2(out)) | |
| return out | |
| class DeConvLayer(nn.Module): | |
| def __init__(self, in_channels, out_channels, kernel_size, stride): | |
| super(DeConvLayer, self).__init__() | |
| # reflection_padding = kernel_size // 2 | |
| # self.reflection_pad = nn.ReflectionPad2d(reflection_padding) | |
| self.deconv2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding=1, output_padding=1) | |
| def forward(self, x): | |
| # out = self.reflection_pad(x) | |
| out = self.deconv2d(x) | |
| return out | |
| # Residual Block | |
| # adapted from pytorch tutorial | |
| # https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/02- | |
| # intermediate/deep_residual_network/main.py | |
| class ResidualBlock(nn.Module): | |
| def __init__(self, channels): | |
| super(ResidualBlock, self).__init__() | |
| self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1) | |
| self.in1 = nn.InstanceNorm2d(channels, affine=True) | |
| self.relu = nn.ReLU() | |
| self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1) | |
| self.in2 = nn.InstanceNorm2d(channels, affine=True) | |
| def forward(self, x): | |
| residual = x | |
| out = self.relu(self.in1(self.conv1(x))) | |
| # out = self.relu(self.in2(self.conv2(out))) | |
| out = self.in2(self.conv2(out)) | |
| # out = self.relu(self.conv2(out)) | |
| # out = self.conv2(out) | |
| out = out + residual | |
| # out = self.relu(out) | |
| return out | |
| class ResidualBlock_depthwise(nn.Module): | |
| def __init__(self, channels): | |
| super(ResidualBlock_depthwise, self).__init__() | |
| # ########################## deptwise ########################################### | |
| self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1, groups=channels) | |
| self.in1 = nn.InstanceNorm2d(channels, affine=True ) | |
| self.conv2 = nn.Conv2d(channels, channels, kernel_size=1, stride=1) | |
| self.in2 = nn.InstanceNorm2d(channels, affine=True ) | |
| self.conv3 = ConvLayer(channels, channels, kernel_size=3, stride=1, groups=channels) | |
| self.in3 = nn.InstanceNorm2d(channels, affine=True ) | |
| self.conv4 = nn.Conv2d(channels, channels, kernel_size=1, stride=1) | |
| self.in4 = nn.InstanceNorm2d(channels, affine=True ) | |
| self.relu = nn.ReLU() | |
| self.prelu = nn.PReLU() | |
| def forward(self, x): | |
| # ############### DEPTWISE ################### | |
| # residual = x | |
| # out = self.relu(self.in1(self.conv1(x))) | |
| # out = self.relu(self.in2(self.conv2(out))) | |
| # out = self.relu(self.in3(self.conv3(out))) | |
| # out = self.relu(self.in4(self.conv4(out))) | |
| # out = out + residual | |
| # # ################## v1 #################### | |
| # residual = x | |
| # out = self.in1(self.conv1(x)) | |
| # out = self.relu(self.in2(self.conv2(out))) | |
| # out = self.in3(self.conv3(out)) | |
| # out = self.in4(self.conv4(out)) | |
| # out = out + residual | |
| # out = self.relu(out) | |
| # ################## v2 #################### √ | |
| residual = x | |
| out = self.in1(self.conv1(x)) | |
| out = self.relu(self.in2(self.conv2(out))) | |
| out = self.in3(self.conv3(out)) | |
| out = self.in4(self.conv4(out)) | |
| out = out + residual | |
| # ################## v3 #################### | |
| # residual = x | |
| # out = self.conv1(x) | |
| # out = self.relu(self.in2(self.conv2(out))) | |
| # out = self.conv3(out) | |
| # out = self.in4(self.conv4(out)) | |
| # out = out + residual | |
| # ################## v4 #################### | |
| # residual = x | |
| # out = self.in1(self.conv1(x)) | |
| # out = self.relu(self.in2(self.conv2(out))) | |
| # out = self.in3(self.conv3(out)) | |
| # out = self.relu(self.in4(self.conv4(out))) | |
| # out = out + residual | |
| return out | |
| # Image Transform Network | |
| class ImageTransformNet(nn.Module): | |
| def __init__(self): | |
| super(ImageTransformNet, self).__init__() | |
| # nonlineraity | |
| self.relu = nn.ReLU() | |
| self.tanh = nn.Tanh() | |
| # encoding layers | |
| self.conv1 = ConvLayer(3, 32, kernel_size=9, stride=1) | |
| self.in1_e = nn.InstanceNorm2d(32, affine=True ) | |
| self.conv2 = ConvLayer(32, 64, kernel_size=3, stride=2) | |
| self.in2_e = nn.InstanceNorm2d(64, affine=True ) | |
| self.conv3 = ConvLayer(64, 128, kernel_size=3, stride=2) | |
| self.in3_e = nn.InstanceNorm2d(128, affine=True ) | |
| # residual layers | |
| self.res1 = ResidualBlock(128) | |
| self.res2 = ResidualBlock(128) | |
| self.res3 = ResidualBlock(128) | |
| self.res4 = ResidualBlock(128) | |
| self.res5 = ResidualBlock(128) | |
| # self.res6 = ResidualBlock(128) | |
| # decoding layers | |
| # TODO: | |
| # self.deconv3 = DeConvLayer(128, 64, kernel_size=3, stride=2) | |
| self.deconv3 = UpsampleConvLayer(128, 64, kernel_size=3, stride=1, upsample=2) | |
| self.in3_d = nn.InstanceNorm2d(64, affine=True ) | |
| # self.deconv2 = DeConvLayer(64, 32, kernel_size=3, stride=2) | |
| self.deconv2 = UpsampleConvLayer(64, 32, kernel_size=3, stride=1, upsample=2) | |
| self.in2_d = nn.InstanceNorm2d(32, affine=True ) | |
| self.deconv1 = ConvLayer(32, 3, kernel_size=9, stride=1) | |
| self.in1_d = nn.InstanceNorm2d(3, affine=True ) | |
| def forward(self, x): | |
| # encode | |
| y = self.relu(self.in1_e(self.conv1(x))) | |
| y = self.relu(self.in2_e(self.conv2(y))) | |
| y = self.relu(self.in3_e(self.conv3(y))) | |
| # y = self.relu(self.conv1(x)) | |
| # y = self.relu(self.conv2(y)) | |
| # y = self.relu(self.conv3(y)) | |
| y_downsample = y | |
| # residual layers | |
| y = self.res1(y) | |
| y = self.res2(y) | |
| y = self.res3(y) | |
| y = self.res4(y) | |
| y = self.res5(y) | |
| y_upsample = y | |
| # decode | |
| y = self.relu(self.in3_d(self.deconv3(y))) | |
| y = self.relu(self.in2_d(self.deconv2(y))) | |
| # y = self.relu(self.deconv3(y)) | |
| # y = self.relu(self.deconv2(y)) | |
| # y = self.tanh(self.in1_d(self.deconv1(y))) | |
| y = self.deconv1(y) | |
| # return y, y_downsample, y_upsample | |
| return y | |
| ALAPHA_1 = 0.25 | |
| ALAPHA_2 = 0.25 | |
| # ALAPHA_1 = 0.5 | |
| # ALAPHA_2 = 0.5 | |
| class ImageTransformNet_dpws(nn.Module): | |
| def __init__(self): | |
| super(ImageTransformNet_dpws, self).__init__() | |
| # nonlineraity | |
| self.relu = nn.ReLU() | |
| self.tanh = nn.Tanh() | |
| # encoding layers | |
| self.conv1 = ConvLayer_dpws(3, int(32*ALAPHA_1), kernel_size=9, stride=1) | |
| # self.in1_e = nn.InstanceNorm2d(int(32*ALAPHA_1), affine=True ) | |
| self.conv2 = ConvLayer_dpws(int(32*ALAPHA_1), int(64*ALAPHA_1), kernel_size=3, stride= 2) | |
| self.conv3 = ConvLayer_dpws(int(64*ALAPHA_1), int(128*ALAPHA_2), kernel_size=3, stride= 2) | |
| # residual layers | |
| self.res1 = ResidualBlock_depthwise(int(128*ALAPHA_2)) | |
| self.res2 = ResidualBlock_depthwise(int(128*ALAPHA_2)) | |
| self.res3 = ResidualBlock_depthwise(int(128*ALAPHA_2)) | |
| self.res4 = ResidualBlock_depthwise(int(128*ALAPHA_2)) | |
| self.res5 = ResidualBlock_depthwise(int(128*ALAPHA_2)) | |
| # self.res6 = ResidualBlock_depthwise(128) | |
| # decoding layers | |
| # TODO: | |
| # self.deconv3 = DeConvLayer(128, 64, kernel_size=3, stride=2) | |
| self.deconv3 = UpsampleConvLayer_dpws(int(128*ALAPHA_2), int(64*ALAPHA_1), kernel_size=3, stride=1, upsample=2) | |
| # self.in3_d = nn.InstanceNorm2d(int(64*ALAPHA_1), affine=True ) | |
| # self.deconv2 = DeConvLayer(64, 32, kernel_size=3, stride=2) | |
| self.deconv2 = UpsampleConvLayer_dpws(int(64*ALAPHA_1), int(32*ALAPHA_1), kernel_size=3, stride=1, upsample=2) | |
| # self.in2_d = nn.InstanceNorm2d(32, affine=True ) | |
| self.deconv1 = ConvLayer_dpws_last(int(32*ALAPHA_1), 3, kernel_size=9, stride=1) | |
| # self.deconv1 = ConvLayer_dpws_last(int(32*ALAPHA_1), 3, kernel_size=9, stride=1) | |
| # self.in1_d = nn.InstanceNorm2d(3, affine=True ) | |
| def forward(self, x): | |
| # encode | |
| # y = self.relu(self.in1_e(self.conv1(x))) | |
| y = self.conv1(x) | |
| y = self.conv2(y) | |
| y = self.conv3(y) | |
| y_downsample = y | |
| # y = self.relu(self.in2_e(self.conv2(y))) | |
| # y = self.relu(self.in3_e(self.conv3(y))) | |
| # residual layers | |
| y = self.res1(y) | |
| y = self.res2(y) | |
| y = self.res3(y) | |
| y = self.res4(y) | |
| y = self.res5(y) | |
| # y = self.res6(y) | |
| y_upsample = y | |
| # decode | |
| y = self.deconv3(y) | |
| y = self.deconv2(y) | |
| y = self.deconv1(y) | |
| # return y, y_downsample, y_upsample | |
| return y | |
| class distiller_1(nn.Module): | |
| def __init__(self): | |
| super(distiller_1, self).__init__() | |
| self.conv = nn.Conv2d(128, int(128*ALAPHA_2), kernel_size=1, stride=1) | |
| def forward(self, x): | |
| # encode | |
| y = self.conv(x) | |
| return y | |
| class distiller_2(nn.Module): | |
| def __init__(self): | |
| super(distiller_2, self).__init__() | |
| self.conv = nn.Conv2d(128, int(128*ALAPHA_2), kernel_size=1, stride=1) | |
| def forward(self, x): | |
| # encode | |
| y = self.conv(x) | |
| return y | |