patrickloeber/snake-ai-pytorch

cpu to gpu

Opened this issue · 11 comments

It has been implemented with cpu, I want to implement this project with gpu, how should I do it, please advise

All you have to do is in the agent.py place the following code right after the model init:

model.to('cuda:0')

Alright I figured out the solution.

Changes to agent.py

Add the below right after line "self.model..."
self.model.to('cuda:0')

Modify def get_action, in the else (add .to('cuda:0')
state0 = torch.tensor(state, dtype=torch.float).to('cuda:0')

Changes to model.py...

"import numpy as np"

add the below just after line "self.model = model"
self.model = model.to('cuda:0')

def train_setup needs to have these changes.

def train_step(self, states, actions, rewards, next_states, dones):
    states = torch.tensor(np.array(states), dtype=torch.float).to('cuda:0')
    actions = torch.tensor(np.array(actions), dtype=torch.long).to('cuda:0')
    rewards = torch.tensor(np.array(rewards), dtype=torch.float).to('cuda:0')
    next_states = torch.tensor(np.array(next_states), dtype=torch.float).to('cuda:0')

    if len(states.shape) == 1:
        states = torch.unsqueeze(states, 0)
        actions = torch.unsqueeze(actions, 0)
        rewards = torch.unsqueeze(rewards, 0)
        next_states = torch.unsqueeze(next_states, 0)
        dones = (dones,)

Or you can do for agent.py:

Action = self.model(state0.to('cuda:0'))

Hi guys, thx for the help i did it and it works, BUT on my PC with AMD 7950X3D and Nvidia 4090 cpu device performance is actually faster than cuda, why is that? In both cases are devices utilized 100%, thx

Sometimes GPUs are meant for graphics strictly, which may hinder them in terms of multiprocessing and calculating each action the neurons will take, it is a fairly uncommon issue.

Yea, but that's THE best PC GPU right now although CPU is top too. Any other cases to process AI is like 100-1000x faster with GPU, for example LLM inferences. Im still think that there is some issue in the code.

If you could show the code, I could probably help optimize it.

Additionally, it would probably be better to train the AI in a headless environment.