谷歌云VM训练深度学习比本地PC慢，求优化方案

问题背景 #

使用PyTorch训练LSTM神经网络，因模型复杂度、数据集规模大导致本地训练缓慢，迁移至谷歌云Deep Learning VM（配置：PyTorch 1.13+CUDA 11.3、1块Nvidia V100 GPU、4vCPU、26GB内存）后，每个Epoch训练时间反而达本地两倍，通过Jupyter Notebook运行脚本，需针对性提速方案。

训练脚本 #

import time
import pandas as pd
import torch
import torch.nn as nn 
import numpy as np
from torch.utils.data import Dataset
from torch.utils.data import DataLoader 


test_indicator = '_test'

indi = '5d'

l = list(range(0,43,1))
l.remove(1)
l.remove(3)
l.remove(7)
l.remove(6)

if indi == '5d':
    l.remove(42) 
else:
    l.remove(41)  

print(l)


epochs = 200
lr = 0.01
batch_size = 131072

look_back = 21

lstm_input_dim = 36
Linear_output_dim = 3
lstm_hidden_dim = 72
Linear_hidden_dim1 = 24 
Linear_hidden_dim2 = 12 
Linear_hidden_dim3 = 6 
lstm_num_layers = 3

path = r'ServerFolder/Data/crsp_train'+test_indicator+'.csv'

# Define a PyTorch dataset for the stock data
class StockDataset(Dataset):
    def __init__(self, path, look_back):
        self.look_back =look_back
        self.df = pd.read_csv(path, usecols=l)
        self.stocks = np.unique(self.df["PERMNO"])
        self.stock_data = {}

        # Split the data by stock and store it in a dictionary
        for stock in self.stocks:
            stock_df = self.df[self.df["PERMNO"] == stock]
            stock_data = stock_df.values
            self.stock_data[stock] = stock_data

    def __len__(self):
        # Return the total number of sequences across all stocks
        return sum(len(self.stock_data[stock]) - self.look_back for stock in self.stocks)

    def __getitem__(self, idx):
        # Determine which stock and which sequence within the stock to use
        stock_idx = 0
        while idx >= len(self.stock_data[self.stocks[stock_idx]]) - self.look_back:
            idx -= len(self.stock_data[self.stocks[stock_idx]]) - self.look_back
            stock_idx += 1

        stock = self.stocks[stock_idx]
        start_idx = idx
        end_idx = idx + self.look_back

        # Get the input and target sequences for the current stock and sequence
        inputs = self.stock_data[stock][start_idx:end_idx, 1:-1]
        target = self.stock_data[stock][end_idx, -1]

        # Convert the numpy arrays to PyTorch tensors
        x = torch.tensor(inputs, dtype=torch.float32)
        y = torch.tensor(target, dtype=torch.long)

        return x,y


# Create a dataset for the entire dataset
dataset = StockDataset(path, look_back)

# Create a data loader for the dataset
loader_train = DataLoader(dataset, batch_size=batch_size)


device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'using {device} device')


from ServerFolder.Code.NeuralNetworks.Models import LSTMModel
model = LSTMModel(lstm_input_dim, lstm_hidden_dim, lstm_num_layers, Linear_output_dim).to(device)


print(model)


loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

test_stats = {
    'loss': [],
    "acc": []
}
train_stats = {
    'loss': [],
    "acc": []
}

def train(dataloader, model, loss_fn, optimizer, multi_acc):
  model.train()
  train_loss = 0
  train_acc = 0

  for i, (x, y) in enumerate(dataloader):
    x, y = x.to(device), y.to(device)

    y_hat = model(x)
    loss = loss_fn(y_hat, y)
    train_loss += loss.item()
    acc = multi_acc(y_hat, y)
    train_acc += acc.item()

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

  num_batches = len(dataloader)
  train_loss = train_loss / num_batches
  train_acc = train_acc / num_batches

  train_stats['loss'].append(train_loss)
  train_stats['acc'].append(train_acc)

  print(
    f'Epoch {epoch + 1:03}: | Train Loss: {train_loss:.5f} | Train Acc: {train_acc:.3f}| ')

def multi_acc(y_hat, y):
    y_pred_softmax = torch.log_softmax(y_hat, dim=1)
    _, y_pred_tags = torch.max(y_pred_softmax, dim=1)

    correct_pred = (y_pred_tags == y).float()
    acc = correct_pred.sum() / len(correct_pred)

    acc = acc * 100

    return acc


for epoch in range(epochs):
    start_time = time.time()

    train(loader_train, model, loss_fn, optimizer, multi_acc)

    train_train_df = pd.DataFrame.from_dict(train_stats).rename(
        columns={"index": "epochs"})

    train_train_df.to_csv(f'ServerFolder/Results/Loss/train_data_{indi}_lstm.csv')

    torch.save(model.state_dict(),
               f'ServerFolder/Results/Model/{indi}'
               f'indicator'+f'{epoch+1}'+'_lstm.pth')


    print("--- %s seconds ---" % (time.time() - start_time))

核心优化建议 #

1. 脱离Jupyter Notebook运行脚本 #

Jupyter的交互特性会引入额外开销，尤其在大批次数据处理场景下：

• 将脚本保存为train_lstm.py
• 终端直接执行：python train_lstm.py
• 需后台运行时，使用nohup python train_lstm.py > train.log 2>&1 &，避免终端断开中断训练

2. 优化数据加载Pipeline（瓶颈大概率在此） #

自定义StockDataset存在多处效率短板，需针对性调整：

• 预转换并缓存张量：在__init__阶段就将numpy数组转为PyTorch张量存储，避免__getitem__重复转换
• 重构索引查找逻辑：当前__getitem__的循环索引查找效率极低，预计算所有序列对应的stock和起始位置，存入列表后直接索引
• 开启多进程数据加载：给DataLoader添加参数，利用多CPU进程加速数据读取与预处理：

  loader_train = DataLoader(dataset, batch_size=batch_size, num_workers=4, pin_memory=True)
  

• 改用高效数据格式：将CSV转为Parquet/Feather格式，Pandas读取速度可提升数倍，替换pd.read_csv为pd.read_parquet

3. 最大化GPU利用率 #

• 验证GPU状态：运行nvidia-smi查看GPU使用率，若使用率低则说明数据加载是瓶颈（对应上述Pipeline优化）
• 调整Batch Size：当前batch_size=131072过大，可能导致内存占用过高或数据加载延迟，尝试减半或调整至GPU内存刚好容纳的大小（可通过torch.cuda.memory_allocated()实时查看内存）
• 启用混合精度训练：PyTorch内置自动混合精度，仅需少量代码即可提升训练速度并降低内存占用：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  
  def train(dataloader, model, loss_fn, optimizer, multi_acc):
      model.train()
      train_loss = 0
      train_acc = 0
  
      for i, (x, y) in enumerate(dataloader):
          x, y = x.to(device), y.to(device)
  
          with autocast():
              y_hat = model(x)
              loss = loss_fn(y_hat, y)
          
          train_loss += loss.item()
          acc = multi_acc(y_hat, y)
          train_acc += acc.item()
  
          optimizer.zero_grad()
          scaler.scale(loss).backward()
          scaler.step(optimizer)
          scaler.update()
  

4. 减少训练过程IO开销 #

当前每个Epoch都保存CSV和模型权重，频繁磁盘IO会拖慢训练：

• 降低保存频率：改为每10个Epoch保存一次模型与统计数据，而非每个Epoch都保存
• 使用高速存储：将ServerFolder迁移至谷歌云VM挂载的SSD本地磁盘，IO速度远快于默认持久磁盘

5. 模型与训练细节优化 #

• 简化准确率计算：multi_acc中torch.log_softmax是冗余操作，softmax不改变最大值位置，直接使用torch.max(y_hat, dim=1)即可：

  def multi_acc(y_hat, y):
      _, y_pred_tags = torch.max(y_hat, dim=1)
      correct_pred = (y_pred_tags == y).float()
      acc = correct_pred.sum() / len(correct_pred)
      return acc * 100
  

• 添加学习率调度：固定lr=0.01可能过大，使用torch.optim.lr_scheduler.StepLR实现学习率衰减，既提升模型收敛性，也能间接减少训练时间

内容的提问来源于stack exchange，提问作者Sebastian