Merge pull request #2 from ritheshkumar95/tristan/gitignore

Add gitignore & Pytorch 0.4 syntax

Author: tristandeleu

.gitignore

@@ -0,0 +1,112 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# Temporary
+tmp/
+run.sh
+
+# Logs & Saves
+logs/
+saves/
+
+# Slurm
+*.out

main.py

@@ -19,6 +19,7 @@
 LAMDA = 1
 LR = 3e-4
 
+DEVICE = torch.device('cuda') # torch.device('cpu')
 
 preproc_transform = transforms.Compose([
     transforms.ToTensor(),
@@ -39,15 +40,15 @@
     num_workers=NUM_WORKERS, pin_memory=True
 )
 
-model = AutoEncoder(INPUT_DIM, DIM, K).cuda()
+model = AutoEncoder(INPUT_DIM, DIM, K).to(DEVICE)
 opt = torch.optim.Adam(model.parameters(), lr=LR)
 
 
 def train():
     train_loss = []
     for batch_idx, (x, _) in enumerate(train_loader):
         start_time = time.time()
-        x = x.cuda()
+        x = x.to(DEVICE)
 
         opt.zero_grad()
 
@@ -85,7 +86,7 @@ def test():
     start_time = time.time()
     val_loss = []
     for batch_idx, (x, _) in enumerate(test_loader):
-        x = x.cuda()
+        x = x.to(DEVICE)
         x_tilde, z_e_x, z_q_x = model(x)
         loss_recons = F.mse_loss(x_tilde, x)
         loss_vq = F.mse_loss(z_q_x, z_e_x.detach())
@@ -100,7 +101,7 @@ def test():
 
 def generate_samples():
     x, _ = test_loader.__iter__().next()
-    x = x[:32].cuda()
+    x = x[:32].to(DEVICE)
     x_tilde, _, _ = model(x)
 
     x_cat = torch.cat([x, x_tilde], 0)

modules.py

@@ -190,7 +190,9 @@ def forward(self, x, label):
         return self.output_conv(x_h)
 
     def generate(self, label, shape=(8, 8), batch_size=64):
-        x = torch.zeros(batch_size, *shape).long().cuda()
+        param = next(self.parameters())
+        x = torch.zeros((batch_size, *shape),
+            dtype=torch.int64, device=param.device)
 
         for i in range(shape[0]):
             for j in range(shape[1]):

pixelcnn.py

@@ -22,6 +22,7 @@
 K = 512
 LR = 3e-4
 
+DEVICE = torch.device('cuda') # torch.device('cpu')
 
 preproc_transform = transforms.Compose([
     transforms.ToTensor(),
@@ -42,23 +43,23 @@
     num_workers=NUM_WORKERS, pin_memory=True
 )
 
-autoencoder = AutoEncoder(INPUT_DIM, VAE_DIM, K).cuda()
+autoencoder = AutoEncoder(INPUT_DIM, VAE_DIM, K).to(DEVICE)
 autoencoder.load_state_dict(
     torch.load('models/{}_autoencoder.pt'.format(DATASET))
 )
 autoencoder.eval()
 
-model = GatedPixelCNN(K, DIM, N_LAYERS).cuda()
-criterion = nn.CrossEntropyLoss().cuda()
+model = GatedPixelCNN(K, DIM, N_LAYERS).to(DEVICE)
+criterion = nn.CrossEntropyLoss().to(DEVICE)
 opt = torch.optim.Adam(model.parameters(), lr=LR)
 
 
 def train():
     train_loss = []
     for batch_idx, (x, label) in enumerate(train_loader):
         start_time = time.time()
-        x = x.cuda()
-        label = label.cuda()
+        x = x.to(DEVICE)
+        label = label.to(DEVICE)
 
         # Get the latent codes for image x
         latents, _ = autoencoder.encode(x)
@@ -93,8 +94,8 @@ def test():
     val_loss = []
     with torch.no_grad():
         for batch_idx, (x, label) in enumerate(test_loader):
-            x = x.cuda()
-            label = label.cuda()
+            x = x.to(DEVICE)
+            label = label.to(DEVICE)
 
             latents, _ = autoencoder.encode(x)
             logits = model(latents.detach(), label)
@@ -114,7 +115,7 @@ def test():
 
 def generate_samples():
     label = torch.arange(10).expand(10, 10).contiguous().view(-1)
-    label = label.long().cuda()
+    label = label.to(device=DEVICE, dtype=torch.int64)
 
     latents = model.generate(label, shape=LATENT_SHAPE, batch_size=100)
     x_tilde, _ = autoencoder.decode(latents)
@@ -129,7 +130,7 @@ def generate_samples():
 
 def generate_reconstructions():
     x, _ = test_loader.__iter__().next()
-    x = x[:32].cuda()
+    x = x[:32].to(DEVICE)
 
     latents, _ = autoencoder.encode(x)
     x_tilde, _ = autoencoder.decode(latents)