尽管批次大小不同,但是每次循环迭代都会导致内存快速消耗。仅在CPU而非GPU上运行时会观察到此情况。
这是重现该问题的摘要,
import torch
import random
from torchvision.models.resnet import resnext50_32x4d
model = resnext50_32x4d(pretrained=False)#Using GPU here makes the problem go away!!
_ = model.eval()
my_predictions = []
for i in range(100):
x = torch.zeros((random.randint(1, 10), 3, 224, 224)).float()
with torch.no_grad():
pred = model(x)
my_predictions.append(pred.mean().item())
有人可以帮助解决问题吗?谢谢!