Merge pull request #247 from SpikeAI/master

Add MPS as a possible computing backend
jeshraghian · Oct 10, 2023 · 5f2d292 · 5f2d292
2 parents 71836f8 + 6d1c48e
commit 5f2d292
Show file tree

Hide file tree

Showing 31 changed files with 2,523 additions and 2,709 deletions.
diff --git a/docs/quickstart.rst b/docs/quickstart.rst
@@ -48,8 +48,8 @@ Define variables for dataloading.
 ::
 
     batch_size = 128
-    data_path='/data/mnist'
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    data_path='/tmp/data/mnist'
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
 
 Load MNIST dataset.
 

diff --git a/docs/tutorials/legacy/tutorial_1_old.rst b/docs/tutorials/legacy/tutorial_1_old.rst
@@ -51,7 +51,7 @@ Let's define a few variables:
 
   # Training Parameters
   batch_size=128
-  data_path='/data/mnist'
+  data_path='/tmp/data/mnist'
   num_classes = 10
 
   # Temporal Dynamics

diff --git a/docs/tutorials/legacy/tutorial_3_old.rst b/docs/tutorials/legacy/tutorial_3_old.rst
@@ -161,7 +161,7 @@ Much of the following code has already been explained in the first two tutorials
 
   # Training Parameters
   batch_size=128
-  data_path='/data/mnist'
+  data_path='/tmp/data/mnist'
 
   # Temporal Dynamics
   num_steps = 25

diff --git a/docs/tutorials/tutorial_1.rst b/docs/tutorials/tutorial_1.rst
@@ -95,7 +95,7 @@ Install the latest PyPi distribution of snnTorch:
 
     # Training Parameters
     batch_size=128
-    data_path='/data/mnist'
+    data_path='/tmp/data/mnist'
     num_classes = 10  # MNIST has 10 output classes
     
     # Torch Variables

diff --git a/docs/tutorials/tutorial_5.rst b/docs/tutorials/tutorial_5.rst
@@ -388,10 +388,10 @@ training a fully-connected spiking neural net.
 
     # dataloader arguments
     batch_size = 128
-    data_path='/data/mnist'
+    data_path='/tmp/data/mnist'
     
     dtype = torch.float
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
 
 ::
 

diff --git a/docs/tutorials/tutorial_6.rst b/docs/tutorials/tutorial_6.rst
@@ -191,10 +191,10 @@ here. <https://snntorch.readthedocs.io/en/latest/snntorch.surrogate.html>`__
 
     # dataloader arguments
     batch_size = 128
-    data_path='/data/mnist'
+    data_path='/tmp/data/mnist'
     
     dtype = torch.float
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
 
 ::
 

diff --git a/docs/tutorials/tutorial_7.rst b/docs/tutorials/tutorial_7.rst
@@ -235,7 +235,7 @@ previous tutorial. The convolutional network architecture to be used is:
 
 ::
 
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
     
     # neuron and simulation parameters
     spike_grad = surrogate.atan()

diff --git a/docs/tutorials/tutorial_ipu_1.rst b/docs/tutorials/tutorial_ipu_1.rst
@@ -62,7 +62,7 @@ Load in the MNIST dataset.
     from torchvision import datasets, transforms
 
     batch_size = 128
-    data_path='/data/mnist'
+    data_path='/tmp/data/mnist'
     
     # Define a transform
     transform = transforms.Compose([

diff --git a/docs/tutorials/tutorial_pop.rst b/docs/tutorials/tutorial_pop.rst
@@ -62,8 +62,8 @@ Define variables for dataloading.
 ::
 
     batch_size = 128
-    data_path='/data/fmnist'
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    data_path='/tmp/data/fmnist'
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
 
 Load FashionMNIST dataset.
 

diff --git a/docs/tutorials/tutorial_regression_1.rst b/docs/tutorials/tutorial_regression_1.rst
@@ -318,7 +318,7 @@ Instantiate the network below:
 ::
 
     hidden = 128
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
     model = Net(timesteps=num_steps, hidden=hidden).to(device)
 
 

diff --git a/docs/tutorials/tutorial_regression_2.rst b/docs/tutorials/tutorial_regression_2.rst
@@ -305,7 +305,7 @@ Instantiate the network below:
 ::
 
     hidden = 128
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
     model = Net(timesteps=num_steps, hidden=hidden, beta=0.9).to(device)
 
 4. Construct Training Loop
@@ -332,7 +332,7 @@ temporal data is an exercise left to the reader/coder.
 ::
 
     batch_size = 128
-    data_path='/data/mnist'
+    data_path='/tmp/data/mnist'
     
     # Define a transform
     transform = transforms.Compose([

diff --git a/docs/tutorials/tutorial_sae.rst b/docs/tutorials/tutorial_sae.rst
@@ -177,10 +177,10 @@
 
       # dataloader arguments
       batch_size = 250
-      data_path='/data/mnist'
+      data_path='/tmp/data/mnist'
 
       dtype = torch.float
-      device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+      device = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
 
 .. container:: cell code
 

diff --git a/docs/tutorials/zh-cn/tutorial_1_cn.rst b/docs/tutorials/zh-cn/tutorial_1_cn.rst
@@ -85,7 +85,7 @@ snnTorch 教程系列基于以下论文。如果您发现这些资源或代码
 
     # Training Parameters
     batch_size=128
-    data_path='/data/mnist'
+    data_path='/tmp/data/mnist'
     num_classes = 10  # MNIST has 10 output classes
     
     # Torch Variables