Generalize device declaration in mnist_hogwild and reinforcement_learning

mnist_hogwild/main.py

@@ -27,10 +27,8 @@
                     help='how many batches to wait before logging training status')
 parser.add_argument('--num-processes', type=int, default=2, metavar='N',
                     help='how many training processes to use (default: 2)')
-parser.add_argument('--cuda', action='store_true', default=False,
-                    help='enables CUDA training')
-parser.add_argument('--mps', action='store_true', default=False,
-                    help='enables macOS GPU training')
+parser.add_argument('--no-accel', action='store_true', default=False,
+                    help='disables accelerator')
 parser.add_argument('--save_model', action='store_true', default=False,
                     help='save the trained model to state_dict')
 parser.add_argument('--dry-run', action='store_true', default=False,
@@ -58,14 +56,13 @@ def forward(self, x):
 if __name__ == '__main__':
     args = parser.parse_args()
 
-    use_cuda = args.cuda and torch.cuda.is_available()
-    use_mps = args.mps and torch.backends.mps.is_available()
-    if use_cuda:
-        device = torch.device("cuda")
-    elif use_mps:
-        device = torch.device("mps")
+    use_accel = not args.no_accel and torch.accelerator.is_available()
+
+    # Set the device to run on
+    if use_accel:
+        device = torch.accelerator.current_accelerator()
     else:
-        device = torch.device("cpu")
+        device = torch.device('cpu')
 
     transform=transforms.Compose([
         transforms.ToTensor(),
@@ -77,7 +74,7 @@ def forward(self, x):
                        transform=transform)
     kwargs = {'batch_size': args.batch_size,
               'shuffle': True}
-    if use_cuda:
+    if use_accel:
         kwargs.update({'num_workers': 1,
                        'pin_memory': True,
                       })

mnist_hogwild/requirements.txt

@@ -1,2 +1,2 @@
 torch
-torchvision==0.20.0
+torchvision

reinforcement_learning/actor_critic.py

@@ -21,6 +21,8 @@
                     help='render the environment')
 parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                     help='interval between training status logs (default: 10)')
+parser.add_argument('--no-accel', action='store_true', default=False,
+                    help='disables accelerator')
 args = parser.parse_args()
 
 
@@ -29,6 +31,14 @@
 env.reset(seed=args.seed)
 torch.manual_seed(args.seed)
 
+use_accel = not args.no_accel and torch.accelerator.is_available()
+
+# Set the device to run on
+if use_accel:
+    device = torch.accelerator.current_accelerator()
+else:
+    device = torch.device('cpu')
+
 
 SavedAction = namedtuple('SavedAction', ['log_prob', 'value'])
 
@@ -70,13 +80,13 @@ def forward(self, x):
         return action_prob, state_values
 
 
-model = Policy()
+model = Policy().to(device)
 optimizer = optim.Adam(model.parameters(), lr=3e-2)
 eps = np.finfo(np.float32).eps.item()
 
 
 def select_action(state):
-    state = torch.from_numpy(state).float()
+    state = torch.from_numpy(state).float().to(device)
     probs, state_value = model(state)
 
     # create a categorical distribution over the list of probabilities of actions
@@ -118,7 +128,7 @@ def finish_episode():
         policy_losses.append(-log_prob * advantage)
 
         # calculate critic (value) loss using L1 smooth loss
-        value_losses.append(F.smooth_l1_loss(value, torch.tensor([R])))
+        value_losses.append(F.smooth_l1_loss(value, torch.tensor([R], device=device)))
 
     # reset gradients
     optimizer.zero_grad()

reinforcement_learning/reinforce.py

@@ -19,6 +19,8 @@
                     help='render the environment')
 parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                     help='interval between training status logs (default: 10)')
+parser.add_argument('--no-accel', action='store_true', default=False,
+                    help='disables accelerator')
 args = parser.parse_args()
 
 
@@ -27,6 +29,14 @@
 env.reset(seed=args.seed)
 torch.manual_seed(args.seed)
 
+use_accel = not args.no_accel and torch.accelerator.is_available()
+
+# Set the device to run on
+if use_accel:
+    device = torch.accelerator.current_accelerator()
+else:
+    device = torch.device('cpu')
+
 
 class Policy(nn.Module):
     def __init__(self):
@@ -46,13 +56,13 @@ def forward(self, x):
         return F.softmax(action_scores, dim=1)
 
 
-policy = Policy()
+policy = Policy().to(device)
 optimizer = optim.Adam(policy.parameters(), lr=1e-2)
 eps = np.finfo(np.float32).eps.item()
 
 
 def select_action(state):
-    state = torch.from_numpy(state).float().unsqueeze(0)
+    state = torch.from_numpy(state).float().unsqueeze(0).to(device)
     probs = policy(state)
     m = Categorical(probs)
     action = m.sample()
@@ -67,7 +77,7 @@ def finish_episode():
     for r in policy.rewards[::-1]:
         R = r + args.gamma * R
         returns.appendleft(R)
-    returns = torch.tensor(returns)
+    returns = torch.tensor(returns, device=device)
     returns = (returns - returns.mean()) / (returns.std() + eps)
     for log_prob, R in zip(policy.saved_log_probs, returns):
         policy_loss.append(-log_prob * R)