File size: 14,951 Bytes
0a6452f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 |
"""
GPU预加载数据加载器
GPU Preloaded Data Loader - 优化小数据集训练速度
通过一次性将所有数据加载到GPU，消除每个batch的CPU-GPU传输开销。
适用于可以完全放入GPU显存的小数据集。
"""
import torch
import numpy as np
import pandas as pd
from typing import Union, Tuple, Optional, Dict, Any
from pathlib import Path
from loguru import logger
class GPUPreloadDataLoader:
"""
GPU预加载数据加载器
将所有数据一次性加载到GPU显存中，在GPU上进行切片操作，
避免每个batch的CPU-GPU传输开销。
优点:
- 消除数据传输瓶颈，训练速度提升1-5%（取决于数据类型）
- GPU上的tensor切片操作非常快
- 简化了训练循环
缺点:
- 占用更多GPU显存
- 不支持数据增强
- 不适合大数据集
适用场景:
- 小数据集（能完全放入GPU显存）
- 表格数据（CSV等结构化数据）
- 不需要复杂数据预处理的场景
"""
def __init__(
self,
data: Union[str, Path, np.ndarray, pd.DataFrame],
batch_size: int = 4096,
shuffle: bool = True,
device: Optional[torch.device] = None,
normalize_features: bool = True,
normalize_labels: bool = False,
input_dim: Optional[int] = None,
output_dim: Optional[int] = None,
feature_cols: Optional[Union[slice, list]] = None,
label_cols: Optional[Union[slice, list]] = None,
feature_names: Optional[list] = None,
label_names: Optional[list] = None
):
"""
初始化GPU预加载数据加载器
Args:
data: 数据路径或数组
batch_size: 批次大小（可以设置更大，如4096/8192）
shuffle: 是否在每个epoch开始时打乱数据
device: 目标设备（默认使用cuda如果可用）
normalize_features: 是否标准化特征
normalize_labels: 是否标准化标签
input_dim: 输入特征维度（如果提供，会自动确定特征列范围）
output_dim: 输出标签维度（如果提供，会自动确定标签列范围）
feature_cols: 特征列的切片范围或列名列表
label_cols: 标签列的切片范围或列名列表（推荐使用列名列表）
feature_names: 特征列名列表（用于从CSV中选择列）
label_names: 标签列名列表（用于从CSV中选择列）
"""
self.batch_size = batch_size
self.shuffle = shuffle
self.normalize_features = normalize_features
self.normalize_labels = normalize_labels
# 确定设备和列范围
if device is None:
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
else:
self.device = device
# 优先使用列名（更安全）
if feature_names is not None and label_names is not None:
self.feature_cols = feature_names
self.label_cols = label_names
self.use_column_names = True
elif feature_cols is not None and label_cols is not None:
self.feature_cols = feature_cols
self.label_cols = label_cols
self.use_column_names = isinstance(feature_cols, list) and isinstance(label_cols, list)
elif input_dim is not None and output_dim is not None:
# 根据input_dim和output_dim自动确定列范围
# 假设格式：前input_dim列是特征，最后output_dim列是标签
# 注意：这种假设可能不安全，推荐使用列名
self.feature_cols = slice(0, input_dim)
self.label_cols = slice(-output_dim, None)
self.use_column_names = False
else:
# 默认：最后一列是标签，其余是特征
self.feature_cols = slice(0, -1)
self.label_cols = slice(-1, None)
self.use_column_names = False
# 加载和预处理数据
features, labels = self._load_and_preprocess_data(data)
# 转换为GPU上的tensor
self.features = torch.FloatTensor(features).to(self.device)
self.labels = torch.FloatTensor(labels).to(self.device)
self.num_samples = self.features.size(0)
self.num_batches = (self.num_samples + self.batch_size - 1) // self.batch_size
logger.info(f"GPU预加载数据加载器初始化完成:")
logger.info(f" 样本数: {self.num_samples}")
logger.info(f" 特征维度: {self.features.size(1)}")
logger.info(f" 标签维度: {self.labels.size(1)}")
logger.info(f" 批次大小: {self.batch_size}")
logger.info(f" 批次数: {self.num_batches}")
logger.info(f" 设备: {self.device}")
logger.info(f" 显存占用: {self.features.element_size() * self.features.nelement() / 1024**2:.2f} MB (特征) + "
f"{self.labels.element_size() * self.labels.nelement() / 1024**2:.2f} MB (标签)")
def _load_and_preprocess_data(
self,
data: Union[str, Path, np.ndarray, pd.DataFrame]
) -> Tuple[np.ndarray, np.ndarray]:
"""
加载和预处理数据
Args:
data: 数据路径或数组
Returns:
特征数组和标签数组
"""
# 加载数据
if isinstance(data, (str, Path)):
# 从文件加载
df = pd.read_csv(data)
elif isinstance(data, pd.DataFrame):
df = data
elif isinstance(data, np.ndarray):
# numpy数组直接使用切片
data_array = data
features = data_array[:, self.feature_cols]
labels = data_array[:, self.label_cols]
# 确保标签是2D数组
if labels.ndim == 1:
labels = labels.reshape(-1, 1)
logger.info(f"数据分割: 特征列 {self.feature_cols}, 标签列 {self.label_cols}")
logger.info(f"特征形状: {features.shape}, 标签形状: {labels.shape}")
return features, labels
else:
raise ValueError(f"不支持的数据类型: {type(data)}")
# 如果是DataFrame，根据是否使用列名来选择列
if self.use_column_names:
# 使用列名选择（更安全）
features = df[self.feature_cols].values
labels = df[self.label_cols].values
logger.info(f"使用列名选择: 特征列 {self.feature_cols}, 标签列 {self.label_cols}")
else:
# 使用列索引切片
data_array = df.values
features = data_array[:, self.feature_cols]
labels = data_array[:, self.label_cols]
logger.info(f"使用索引切片: 特征列 {self.feature_cols}, 标签列 {self.label_cols}")
# 确保标签是2D数组
if labels.ndim == 1:
labels = labels.reshape(-1, 1)
logger.info(f"特征形状: {features.shape}, 标签形状: {labels.shape}")
# 标准化
if self.normalize_features:
features = self._normalize(features, fit=True)
if self.normalize_labels:
labels = self._normalize(labels, fit=True)
return features, labels
def _normalize(
self,
data: np.ndarray,
fit: bool = True
) -> np.ndarray:
"""
标准化数据
Args:
data: 数据数组
fit: 是否拟合标准化参数
Returns:
标准化后的数据
"""
if fit:
# 计算均值和标准差
self.mean = np.mean(data, axis=0)
self.std = np.std(data, axis=0)
# 避免除零
self.std[self.std < 1e-8] = 1.0
return (data - self.mean) / self.std
def __iter__(self):
"""
创建迭代器
Returns:
迭代器对象
"""
self.current_batch = 0
# 生成索引
if self.shuffle:
# 在GPU上生成随机索引
self.indices = torch.randperm(self.num_samples, device=self.device)
else:
self.indices = torch.arange(self.num_samples, device=self.device)
return self
def __next__(self) -> Tuple[torch.Tensor, torch.Tensor]:
"""
获取下一个batch
Returns:
(特征, 标签) 元组
Raises:
StopIteration: 当迭代完成时
"""
if self.current_batch >= self.num_batches:
raise StopIteration
# 计算当前batch的索引范围
start_idx = self.current_batch * self.batch_size
end_idx = min(start_idx + self.batch_size, self.num_samples)
# 获取当前batch的索引
batch_indices = self.indices[start_idx:end_idx]
# 在GPU上进行切片操作
batch_features = self.features[batch_indices]
batch_labels = self.labels[batch_indices]
self.current_batch += 1
return batch_features, batch_labels
def __len__(self) -> int:
"""
返回批次数
Returns:
批次数
"""
return self.num_batches
def to(self, device: torch.device) -> 'GPUPreloadDataLoader':
"""
将数据移动到指定设备
Args:
device: 目标设备
Returns:
self
"""
self.device = device
self.features = self.features.to(device)
self.labels = self.labels.to(device)
logger.info(f"数据已移动到设备: {device}")
return self
class GPUPreloadDataLoaderFactory:
"""
GPU预加载数据加载器工厂类
用于创建训练、验证和测试的GPU预加载数据加载器
"""
def __init__(self, config: Optional[Dict[str, Any]] = None):
"""
初始化工厂
Args:
config: 配置字典
"""
self.config = config or {}
def create_train_loader(
self,
data_path: Union[str, Path],
input_dim: Optional[int] = None,
output_dim: Optional[int] = None,
**kwargs
) -> GPUPreloadDataLoader:
"""
创建训练数据加载器
Args:
data_path: 数据路径
input_dim: 输入特征维度
output_dim: 输出标签维度
**kwargs: 额外参数
Returns:
训练数据加载器
"""
config = {**self.config, **kwargs}
config['shuffle'] = True # 训练时打乱数据
# 明确指定列名（优先级高于 input_dim/output_dim）
default_feature_names = [
'user_pad_p', 'user_pad_a', 'user_pad_d',
'vitality',
'ai_current_pad_p', 'ai_current_pad_a', 'ai_current_pad_d'
]
default_label_names = ['ai_delta_p', 'ai_delta_a', 'ai_delta_d']
# 使用列名而不是索引切片（更安全）
config['feature_names'] = config.get('feature_names', default_feature_names)
config['label_names'] = config.get('label_names', default_label_names)
# 移除 input_dim 和 output_dim（不再使用）
config.pop('input_dim', None)
config.pop('output_dim', None)
return GPUPreloadDataLoader(
data=data_path,
**config
)
def create_val_loader(
self,
data_path: Union[str, Path],
input_dim: Optional[int] = None,
output_dim: Optional[int] = None,
**kwargs
) -> GPUPreloadDataLoader:
"""
创建验证数据加载器
Args:
data_path: 数据路径
input_dim: 输入特征维度
output_dim: 输出标签维度
**kwargs: 额外参数
Returns:
验证数据加载器
"""
config = {**self.config, **kwargs}
config['shuffle'] = False # 验证时不打乱数据
# 明确指定列名（优先级高于 input_dim/output_dim）
default_feature_names = [
'user_pad_p', 'user_pad_a', 'user_pad_d',
'vitality',
'ai_current_pad_p', 'ai_current_pad_a', 'ai_current_pad_d'
]
default_label_names = ['ai_delta_p', 'ai_delta_a', 'ai_delta_d']
# 使用列名而不是索引切片（更安全）
config['feature_names'] = config.get('feature_names', default_feature_names)
config['label_names'] = config.get('label_names', default_label_names)
# 移除 input_dim 和 output_dim（不再使用）
config.pop('input_dim', None)
config.pop('output_dim', None)
return GPUPreloadDataLoader(
data=data_path,
**config
)
def create_test_loader(
self,
data_path: Union[str, Path],
input_dim: Optional[int] = None,
output_dim: Optional[int] = None,
**kwargs
) -> GPUPreloadDataLoader:
"""
创建测试数据加载器
Args:
data_path: 数据路径
input_dim: 输入特征维度
output_dim: 输出标签维度
**kwargs: 额外参数
Returns:
测试数据加载器
"""
config = {**self.config, **kwargs}
config['shuffle'] = False # 测试时不打乱数据
# 明确指定列名（优先级高于 input_dim/output_dim）
default_feature_names = [
'user_pad_p', 'user_pad_a', 'user_pad_d',
'vitality',
'ai_current_pad_p', 'ai_current_pad_a', 'ai_current_pad_d'
]
default_label_names = ['ai_delta_p', 'ai_delta_a', 'ai_delta_d']
# 使用列名而不是索引切片（更安全）
config['feature_names'] = config.get('feature_names', default_feature_names)
config['label_names'] = config.get('label_names', default_label_names)
# 移除 input_dim 和 output_dim（不再使用）
config.pop('input_dim', None)
config.pop('output_dim', None)
return GPUPreloadDataLoader(
data=data_path,
**config
)
def create_gpu_preload_loader(
data_path: Union[str, Path],
batch_size: int = 4096,
shuffle: bool = True,
device: Optional[torch.device] = None,
**kwargs
) -> GPUPreloadDataLoader:
"""
创建GPU预加载数据加载器的便捷函数
Args:
data_path: 数据路径
batch_size: 批次大小
shuffle: 是否打乱数据
device: 目标设备
**kwargs: 其他参数
Returns:
GPU预加载数据加载器实例
"""
return GPUPreloadDataLoader(
data=data_path,
batch_size=batch_size,
shuffle=shuffle,
device=device,
**kwargs
)
|