训练期间增加GPU内存使用量

时间:2020-10-15 10:15:46

标签: python pytorch gpu

我正在使用常规MNIST数据集训练网络,遇到了下一个问题:
当我开始将valid_metrics添加到loss_list和precision_list时,正在使用的GPU内存量每1或2个周期开始增加。 这是train_loop的代码:

def train_model(model: torch.nn.Module, 
                train_dataset: torch.utils.data.Dataset,
                valid_dataset: torch.utils.data.Dataset,
                loss_function: torch.nn.Module = torch.nn.CrossEntropyLoss(),
                optimizer_class: Type[torch.optim.Optimizer] = torch.optim,
                optimizer_params: Dict = {},
                initial_lr = 0.01,
                lr_scheduler_class: Any = torch.optim.lr_scheduler.ReduceLROnPlateau,
                lr_scheduler_params: Dict = {},
                batch_size = 64,
                max_epochs = 1000,
                early_stopping_patience = 20):
    
    
    optimizer = torch.optim.Adam(model.parameters(), lr=initial_lr, **optimizer_params)
    lr_scheduler = lr_scheduler_class(optimizer, **lr_scheduler_params)
    
    train_loader = torch.utils.data.DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
    valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size)

    best_valid_loss = None
    best_epoch = None
    loss_list = list()
    accuracy_list = list()
    
    for epoch in range(max_epochs):
        
        print(f'Epoch {epoch}')
        
        start = timer()
        
        train_single_epoch(model, optimizer, loss_function, train_loader)
        
        valid_metrics = validate_single_epoch(model, loss_function, valid_loader)
        loss_list.append(valid_metrics['loss'])
        accuracy_list.append(valid_metrics['accuracy'])
        
        print('time:', timer() - start)
        print(f'Validation metrics: \n{valid_metrics}')

        lr_scheduler.step(valid_metrics['loss'])
        
        if best_valid_loss is None or best_valid_loss > valid_metrics['loss']:
            print(f'Best model yet, saving')
            best_valid_loss = valid_metrics['loss']
            best_epoch = epoch
            torch.save(model, './best_model.pth')
            
        if epoch - best_epoch > early_stopping_patience:
            print('Early stopping triggered')
            return loss_list, accuracy_list

和代码validate_single_epoch

def validate_single_epoch(model: torch.nn.Module,
                          loss_function: torch.nn.Module, 
                          data_loader: torch.utils.data.DataLoader):
    
    loss_total = 0
    accuracy_total = 0
    
    for data in data_loader:

        X, y = data
        X, y = X.view(-1, 784), y.to(device)
        X = X.to(device)
        
        output = model(X)
        
        loss = loss_function(output, y)
        loss_total += loss
        

        y_pred = output.argmax(dim = 1, keepdim=True).to(device)
        accuracy_total += y_pred.eq(y.view_as(y_pred)).sum().item()
        
    loss_avg = loss_total / len(data_loader.dataset)
    accuracy_avg = 100.0 * accuracy_total / len(data_loader.dataset)
    
    return {'loss' : loss_avg, 'accuracy' : accuracy_avg}

我使用GeForce MX250作为GPU

1 个答案:

答案 0 :(得分:0)

该问题可能是因为正在计算渐变并将其存储在验证循环中。为了解决这个问题,也许最简单的方法是将验证调用包装在ActiveX control '8856f961-340a-11d0-a96b-00c04fd705a2' cannot be instantiated because the current thread is not in a single-threaded apartment. 上下文中:

no_grad

如果愿意,还可以用with torch.no_grad(): valid_metrics = validate_single_epoch(model, loss_function, valid_loader) 装饰validate_single_epoch(...)

@torch.no_grad()

与您的问题无关,但请注意,在验证期间您正在训练模式下使用模型,这可能不是您想要的。验证功能中可能缺少对@torch.no_grad() def validate_single_epoch(...): # ... 的调用。