SHARE
TWEET

Untitled

a guest Jun 19th, 2019 54 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. class Rocket_E_NN(nn.Module):
  2.     def __init__(self):
  3.         super().__init__()
  4.  
  5.         self.encoder = nn.Sequential(
  6.             nn.Conv2d(3, 32, 4, 2, 1),          # B,  32, 32, 32
  7.             nn.ReLU(True),
  8.             nn.Conv2d(32, 32, 4, 2, 1),          # B,  32, 16, 16
  9.             nn.ReLU(True),
  10.             nn.Conv2d(32, 64, 4, 2, 1),          # B,  64,  8,  8
  11.             nn.ReLU(True),
  12.             nn.Conv2d(64, 64, 4, 2, 1),          # B,  64,  4,  4
  13.             nn.ReLU(True),
  14.             nn.Conv2d(64, 256, 4, 1),            # B, 256,  1,  1
  15.             nn.ReLU(True),
  16.             View((-1, 256*1*1)),                 # B, 256
  17.             nn.Linear(256, 2),             # B, 1
  18.         )
  19.  
  20.     def forward(self, x):
  21.         z = self.encoder(x)
  22.         return z
  23.  
  24. class Rocket_D_NN(nn.Module):
  25.     def __init__(self):
  26.         super().__init__()
  27.         self.decoder = nn.Sequential(
  28.             nn.Linear(2, 256),               # B, 256
  29.             View((-1, 256, 1, 1)),               # B, 256,  1,  1
  30.             nn.ReLU(True),
  31.             nn.ConvTranspose2d(256, 64, 4),      # B,  64,  4,  4
  32.             nn.ReLU(True),
  33.             nn.ConvTranspose2d(64, 64, 4, 2, 1), # B,  64,  8,  8
  34.             nn.ReLU(True),
  35.             nn.ConvTranspose2d(64, 32, 4, 2, 1), # B,  32, 16, 16
  36.             nn.ReLU(True),
  37.             nn.ConvTranspose2d(32, 32, 4, 2, 1), # B,  32, 32, 32
  38.             nn.ReLU(True),
  39.             nn.ConvTranspose2d(32, 3, 4, 2, 1),  # B, 3, 64, 64
  40.         )
  41.  
  42.     def forward(self, z):
  43.         x = self.decoder(z)
  44.         return x
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top