CUDAError: недостаточно памяти для среды RL с использованием Gymnasium. ⇐ Python

[Anonymous]

Код: Выделить всё

from TTset_main import load_TTset

# File paths for training and validation data
filted_tset_path = "path.csv"
filted_vset_path = "path.csv"

# Load training and validation sets
TTset_training, TTset_vali, training_data_length = load_TTset(filted_tset_path, filted_vset_path)

# Extract frame tensors and labels from the training set
frame_tensors = []
labels = []

for inputs, targets in TTset_training:
frame_tensors.append(inputs)
labels.append(targets)

# Print dimensions of the frame tensors and labels
print("Inputs")
print("Dimension 1: ", len(frame_tensors)) #1927
print("Dimension 2: ", len(frame_tensors[0])) #1
print("Dimension 3: ", len(frame_tensors[0][0])) #3
print("Dimension 4: ", len(frame_tensors[0][0][0])) #60
print("Dimension 5: ", len(frame_tensors[0][0][0][0])) #230
print("Dimension 6: ", len(frame_tensors[0][0][0][0][0])) #230
print()
print("Training data length:", training_data_length)
print("Targets")
print("Dimension 1: ", len(labels)) #1927
print("Dimension 2: ", len(labels[0])) #1

import numpy as np
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback
from tqdm import tqdm

# Custom gym environment for table tennis
class TableTennisEnv(gym.Env):
def __init__(self, frame_tensors, labels, frame_size=(3, 60, 230, 230)):
super(TableTennisEnv, self).__init__()
self.frame_tensors = frame_tensors
self.labels = labels
self.current_step = 0
self.current_substep = 0
self.frame_size = frame_size

self.n_actions = 20  # Number of unique actions
self.observation_space = spaces.Box(low=0, high=255, shape=(3, 60, 230, 230), dtype=np.float32)
self.action_space = spaces.Discrete(self.n_actions)
self.normalize_images = False

def reset(self, seed=None):
self.current_step = 0
self.current_substep = 0
return self.frame_tensors[self.current_step][self.current_substep], {}

def step(self, action):
reward = 1 if action == self.labels[self.current_step][self.current_substep] else -1
self.current_substep += 1

if self.current_substep >= len(self.frame_tensors[self.current_step]):
self.current_substep = 0
self.current_step += 1

done = self.current_step >= len(self.frame_tensors)

obs = self.frame_tensors[self.current_step][self.current_substep] if not done else np.zeros_like(self.frame_tensors[0][0])

return obs, reward, done, {}

def render(self, mode='human'):
pass

# Reduce memory usage by processing in smaller batches
env = DummyVecEnv([lambda: TableTennisEnv(frame_tensors, labels, frame_size=(3, 60, 230, 230))])

# Callback for progress bar during training
class ProgressBarCallback(BaseCallback):
def __init__(self, total_timesteps, verbose=0):
super(ProgressBarCallback, self).__init__(verbose)
self.total_timesteps = total_timesteps
self.pbar = None

def _on_training_start(self):
self.pbar = tqdm(total=self.total_timesteps)

def _on_step(self):
self.pbar.update(self.model.n_steps)
return True

def _on_training_end(self):
self.pbar.close()

# Set total timesteps for training
total_timesteps = 2048  # Adjust timesteps as needed

# Initialize PPO model with a smaller batch size
model = PPO("MlpPolicy", env, verbose=1, n_steps=64, tensorboard_log="./ppo_tt_tensorboard/")  # Reduce batch size
progress_bar_callback = ProgressBarCallback(total_timesteps=total_timesteps)

# Train the model
model.learn(total_timesteps=total_timesteps, callback=progress_bar_callback)

# Save the trained model
model.save("ppo_table_tennis")

Я использую приведенный выше код для обучения модели RL для классификации видео. Видео вводятся в виде тензоров. Frame_tensors имеет размеры 1927x1x3x60x230x230, а labels — 1927x1. Одно видео представлено в размерах 3x60x230x230. Я получаю следующую ошибку:

Код: Выделить всё

Traceback (most recent call last):
File "d:/PJ/TTARnet_v2_to_lab_兆仁code/prototype.py", line 107, in 
model.learn(total_timesteps=total_timesteps, callback=progress_bar_callback)
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\stable_baselines3\ppo\ppo.py", line 315, in learn
return super().learn(
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\stable_baselines3\common\on_policy_algorithm.py", line 313, in learn
self.train()
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\stable_baselines3\ppo\ppo.py", line 282, in train
self.policy.optimizer.step()
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\torch\optim\optimizer.py", line 391, in wrapper
out = func(*args, **kwargs)
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\torch\optim\optimizer.py", line 76, in _use_grad
ret = func(self, *args, **kwargs)
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\torch\optim\adam.py", line 168, in step
adam(
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\torch\optim\adam.py", line 318, in adam
func(params,
File "C:\Users\User1\anaconda3\envs\ML\lib\site-packages\torch\optim\adam.py", line 581, in _multi_tensor_adam
exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.27 GiB. GPU

Изменить:
Устранена предыдущая ошибка. Новая ошибка связана с памятью графического процессора. Я запускаю этот код на 4070Ti, у него 12 ГБ памяти, но по какой-то причине код требует больше, примерно на 2,27 ГБ больше.
Я уменьшил размер пакета до 4 и не могу уменьшите размер кадра, так как я потеряю много информации.

Подробнее здесь: https://stackoverflow.com/questions/786 ... -gymnasium