如何解决PyTorch-lightning 模型在第一个纪元后内存不足
我在 PyTorch 上看到了一个 Kaggle 内核,并使用相同的 img_size、batch_size 等运行它,并创建了另一个具有完全相同值的 PyTorch-lightning 内核,但我的闪电模型在大约 1.5 个 epoch 后内存不足(每个 epoch 包含 8750步骤)在第一次折叠,而原生 PyTorch 模型运行整个 5 次折叠。有没有办法改进代码或释放内存?我本可以尝试删除模型或进行一些垃圾收集,但如果它甚至没有完成第一次折叠,我就无法删除模型和东西。
def run_fold(fold):
df_train = train[train['fold'] != fold]
df_valid = train[train['fold'] == fold]
train_dataset = G2NetDataset(df_train,get_train_aug())
valid_dataset = G2NetDataset(df_valid,get_test_aug())
train_dl = DataLoader(train_dataset,batch_size = config.batch_size,num_workers = config.num_workers,shuffle = True,drop_last = True,pin_memory = True)
valid_dl = DataLoader(valid_dataset,shuffle = False,drop_last = False,pin_memory = True)
model = Classifier()
logger = pl.loggers.WandbLogger(project='G2Net',name=f'fold: {fold}')
trainer = pl.Trainer(gpus = 1,max_epochs = config.epochs,fast_dev_run = config.debug,logger = logger,log_every_n_steps=10)
trainer.fit(model,train_dl,valid_dl)
result = trainer.test(test_dataloaders = valid_dl)
wandb.run.finish()
return result
def main():
if config.train:
results = []
for fold in range(config.n_fold):
result = run_fold(fold)
results.append(result)
return results
results = main()
解决方法
如果不查看您的模型类,我不能说太多,但我遇到的几个可能的问题是日志记录的度量和损失评估。 例如,像
这样的东西pl.metrics.Accuracy(compute_on_step=False)
需要并显式调用 .compute()
def training_epoch_end(self,outputs):
loss = sum([out['loss'] for out in outputs])/len(outputs)
self.log_dict({'train_loss' : loss.detach(),'train_accuracy' : self.train_metric.compute()})
在纪元结束时。
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。