Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,317 Bytes
02ba63a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import torch
from torch import nn
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
super(ConvBlock, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
self.batchnorm = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU()
def forward(self, x):
return self.relu(self.batchnorm(self.conv(x)))
class DeconvBlock(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding):
super(DeconvBlock, self).__init__()
self.deconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
self.batchnorm = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU()
def forward(self, x):
return self.relu(self.batchnorm(self.deconv(x)))
class Autoencoder(nn.Module):
def __init__(self, feature_dim=32):
super(Autoencoder, self).__init__()
self.feature_dim = feature_dim
# エンコーダ
self.enc1 = ConvBlock(3, 16, 10, 1, 0)
self.enc2 = ConvBlock(16, 32, 10, 1, 0)
self.enc3 = ConvBlock(32, 64, 2, 2, 0)
self.enc4 = ConvBlock(64, 128, 2, 2, 0)
self.enc5 = ConvBlock(128, 256, 2, 2, 0)
# デコーダ
self.dec1 = DeconvBlock(256, 128, 2, 2, 0, 1)
self.dec2 = DeconvBlock(256, 64, 2, 2, 0, 1) # 128 + 128
self.dec3 = DeconvBlock(128, 32, 2, 2, 0, 0) # 64 + 64
self.dec4 = DeconvBlock(64, 16, 10, 1, 0, 0) # 32 + 32
self.dec5 = DeconvBlock(32, self.feature_dim, 10, 1, 0, 0)
self.dec6 = nn.Conv2d(self.feature_dim, 32, 1, 1, 0)
self.dec7 = nn.Conv2d(32, 3, 1, 1, 0)
def forward(self, x):
# エンコーダ
enc1 = self.enc1(x)
enc2 = self.enc2(enc1)
enc3 = self.enc3(enc2)
enc4 = self.enc4(enc3)
enc5 = self.enc5(enc4)
# デコーダ
dec1 = self.dec1(enc5)
dec2 = self.dec2(torch.cat((dec1, enc4), 1))
dec3 = self.dec3(torch.cat((dec2, enc3), 1))
dec4 = self.dec4(torch.cat((dec3, enc2), 1))
dec5 = self.dec5(torch.cat((dec4, enc1), 1))
dec6 = self.dec6(dec5)
dec7 = self.dec7(dec6)
return dec5, dec7
|