Generalize device declaration in mnist_hogwild and reinforcement_learning

mnist_hogwild/main.py

@@ -27,10 +27,6 @@
                     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('--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 +54,8 @@ 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")
-    else:
-        device = torch.device("cpu")
+    # Set the device to run on
+    device = torch.accelerator.current_accelerator() or torch.device('cpu')
 
     transform=transforms.Compose([
         transforms.ToTensor(),
@@ -77,7 +67,7 @@ def forward(self, x):
                        transform=transform)
     kwargs = {'batch_size': args.batch_size,
               'shuffle': True}
-    if use_cuda:
+    if torch.accelerator.is_available():
         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

@@ -29,6 +29,9 @@
 env.reset(seed=args.seed)
 torch.manual_seed(args.seed)
 
+# Set the device to run on
+device = torch.accelerator.current_accelerator() or torch.device('cpu')
+
 
 SavedAction = namedtuple('SavedAction', ['log_prob', 'value'])
 
@@ -70,13 +73,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).to(dtype=torch.float, device=device)
     probs, state_value = model(state)
 
     # create a categorical distribution over the list of probabilities of actions
@@ -118,7 +121,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

@@ -27,6 +27,9 @@
 env.reset(seed=args.seed)
 torch.manual_seed(args.seed)
 
+# Set the device to run on
+device = torch.accelerator.current_accelerator() or torch.device('cpu')
+
 
 class Policy(nn.Module):
     def __init__(self):
@@ -46,13 +49,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).to(dtype=torch.float, device=device).unsqueeze(0)
     probs = policy(state)
     m = Categorical(probs)
     action = m.sample()
@@ -67,7 +70,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)

run_python_examples.sh

@@ -98,7 +98,7 @@ function mnist_forward_forward() {
 
 }
 function mnist_hogwild() {
-  uv run main.py --epochs 1 --dry-run $CUDA_FLAG || error "mnist hogwild failed"
+  uv run main.py --epochs 1 --dry-run || error "mnist hogwild failed"
 }
 
 function mnist_rnn() {