Autoencoder3.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. from torch import nn
  2. class Autoencoder(nn.Module):
  3. def __init__(self, dropout=0.1, latent_features=512):
  4. super(Autoencoder, self).__init__()
  5. self.encoder = nn.Sequential(
  6. nn.Dropout(dropout),
  7. nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
  8. nn.ReLU(True),
  9. nn.Dropout(dropout),
  10. nn.Conv2d(64, 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, 64, kernel_size=3, stride=2, padding=1),
  20. nn.ReLU(True),
  21. nn.Dropout(dropout),
  22. nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
  23. nn.ReLU(True),
  24. nn.Dropout(dropout),
  25. nn.Flatten(),
  26. nn.Linear(1024, latent_features),
  27. nn.ReLU(True),
  28. )
  29. self.decoder = nn.Sequential(
  30. nn.Linear(512, 1024),
  31. nn.ReLU(True),
  32. nn.Unflatten(1, (64, 4, 4)),
  33. nn.Dropout(dropout),
  34. nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1),
  35. nn.ReLU(True),
  36. nn.Dropout(dropout),
  37. nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1),
  38. nn.ReLU(True),
  39. nn.Dropout(dropout),
  40. nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1),
  41. nn.ReLU(True),
  42. nn.Dropout(dropout),
  43. nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1),
  44. nn.ReLU(True),
  45. nn.Dropout(dropout),
  46. nn.ConvTranspose2d(64, 64, kernel_size=6, stride=2, padding=2),
  47. nn.ReLU(True),
  48. nn.Dropout(dropout),
  49. nn.ConvTranspose2d(64, 64, kernel_size=8, stride=2, padding=3),
  50. nn.ReLU(True),
  51. nn.Dropout(dropout),
  52. nn.Conv2d(64, 3, kernel_size=3, stride=1, padding="same"),
  53. nn.Tanh(),
  54. )
  55. def forward(self, x):
  56. x = self.encoder(x)
  57. x = self.decoder(x)
  58. return x