Using RPC only decreases “backward time”.
Hello I am using distributed.rpc changing the example in examples/main.py at 01539f9eada34aef67ae7b3d674f30634c35f468 · pytorch/examples · GitHub 2
I have N agents and 1 simulation. I would like to select actions in parallel, if possible in different GPU’s.
However, when I execute my code the only thing that seems to reduce is the update part. Below there is the code for selecting the actions which does not seem to be reduced at all, any ideas?
from torch.distributed.rpc import RRef, rpc_sync, rpc_async, remote

def _call_method(method, rref, *args, **kwargs):
    r"""
    a helper function to call a method on the given RRef
    """
    return method(rref.local_value(), *args, **kwargs)


def _remote_method(method, rref, *args, **kwargs):
    r"""
    a helper function to run method on the owner of rref and fetch back the
    result using RPC
    """
    args = [method, rref] + list(args)
    return rpc_sync(rref.owner(), _call_method, args=args, kwargs=kwargs)

 def select_actions_all_agents(self, state, step):

        self.current_actions = np.zeros(self.current_actions.shape, dtype=np.int32) * -1000
        futs = []
        start_time = time.time()
        for ag_rreff in self.ag_rrefs:
            # make async RPC to kick off an episode on all observers
            futs.append(
                rpc_async(
                    ag_rreff.owner(),
                    _call_method,
                    args=(Agent.select_action, ag_rreff, state, step)
                )
            )
            # if step % 50 == 0:
            #    print("Done", ag_rreff, time.time()-start_time)
        # wait until all obervers have finished this episode
        for fut in futs:
           fut.wait()

class Agent:
    def __init__(self):
        self.id = rpc.get_worker_info().id
        device = (self.id - 1) % args.num_gpus
        print("Initialasing Agent ID: {}, Device {}".format(self.id, device))
        self.device = ("cuda:" + str(device) if torch.cuda.is_available() else "cpu")
        self.rewards = []
        self.saved_log_probs = []
        # torch.manual_seed(args.seed+self.id)
        self.policy = Policy()
        self.policy.to(self.device)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2)
        self.eps = np.finfo(np.float32).eps.item()

    def select_action(self, state, step):
        r"""
        This function is mostly borrowed from the Reinforcement Learning example.
        See https://github.com/pytorch/examples/tree/master/reinforcement_learning
        The main difference is that instead of keeping all probs in one list,
        the agent keeps probs in a dictionary, one key per observer.

        NB: no need to enforce thread-safety here as GIL will serialize
        executions.
        """
        start_time = time.time()
        state = torch.from_numpy(state).float().unsqueeze(0).to(self.device)
        start_time_policy = time.time()
        probs = self.policy(state)
        m = Categorical(probs)
        action = m.sample()
        self.saved_log_probs.append(m.log_prob(action))
        _remote_method(Master.report_action, master_rref, self.id,  action.item())