Autoencoder2.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. from torch import nn
  2. class Autoencoder(nn.Module):
  3. def __init__(self, dropout=0.1, latent_channels=32):
  4. super(Autoencoder, self).__init__()
  5. self.encoder = nn.Sequential(
  6. nn.Dropout(dropout),
  7. nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3),
  8. nn.ReLU(True),
  9. nn.Dropout(dropout),
  10. nn.Conv2d(32, 64, kernel_size=5, stride=2, padding=2),
  11. nn.ReLU(True),
  12. nn.Dropout(dropout),
  13. nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
  14. nn.ReLU(True),
  15. nn.Dropout(dropout),
  16. nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
  17. nn.ReLU(True),
  18. nn.Dropout(dropout),
  19. nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
  20. nn.ReLU(True),
  21. nn.Dropout(dropout),
  22. nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),
  23. nn.ReLU(True),
  24. nn.Dropout(dropout),
  25. nn.Conv2d(128, latent_channels, kernel_size=3, padding="same"),
  26. nn.ReLU(True),
  27. )
  28. self.decoder = nn.Sequential(
  29. nn.Dropout(dropout),
  30. nn.Conv2d(latent_channels, 128, kernel_size=3, padding="same"),
  31. nn.ReLU(True),
  32. nn.Dropout(dropout),
  33. nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1),
  34. nn.ReLU(True),
  35. nn.Dropout(dropout),
  36. nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
  37. nn.ReLU(True),
  38. nn.Dropout(dropout),
  39. nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1),
  40. nn.ReLU(True),
  41. nn.Dropout(dropout),
  42. nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1),
  43. nn.ReLU(True),
  44. nn.Dropout(dropout),
  45. nn.ConvTranspose2d(64, 32, kernel_size=6, stride=2, padding=2),
  46. nn.ReLU(True),
  47. nn.Dropout(dropout),
  48. nn.ConvTranspose2d(32, 16, kernel_size=8, stride=2, padding=3),
  49. nn.ReLU(True),
  50. nn.Dropout(dropout),
  51. nn.Conv2d(16, 3, kernel_size=3, stride=1, padding="same"),
  52. nn.Tanh(),
  53. )
  54. def forward(self, x):
  55. x = self.encoder(x)
  56. x = self.decoder(x)
  57. return x