我正在实现全卷积网络（UNet）的尖峰版本。我测试了没有 LIF 神经元的正常模型，虽然它非常占用内存，但效果很好，所以我必须使用批量大小 < 4. After converting the model to use LIF neurons, it runs for 5 iterations and then runs out of memory. Every subsequent iteration is slower than the last as well. My inputs are 4 channel 512x512 so they take a lot of memory to handle even without the added temporal dimension, but I still get the CUDA out of memory error even with a batch size of 1 and num_steps = 1 (essentially no temporal dimension at all). So the problem isn't that I don't have the juice to run the model, it's that something is clogging up the memory and not being removed when it needs to (at least I think). I've been fiddling around with it for a while but I can't seem to find the problem. Below I've attached my code plus the outputs of a profiler running the loop with num_steps = 2, batch size = 1 and for 5 iterations.
链接到我的探查器输出和跟踪： https://drive.google.com/drive/folders/1Lg_92gmSAzVlVb3AnsbxJokGVxfDRGd_?usp=drive_link
# UNet parts
class DoubleConv(nn.Module):
"""(convolution => [BNTT] => Spikes) * 2 + Dropout"""
def __init__(self, in_channels, out_channels, beta, SG_func, mid_channels=None):
super().__init__()
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
#snn.bntt.BatchNormTT2d(mid_channels, time_steps=num_steps),
snn.Leaky(beta=beta, spike_grad=SG_func, init_hidden=True),
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
#snn.bntt.BatchNormTT2d(out_channels, time_steps=num_steps),
snn.Leaky(beta=beta, spike_grad=SG_func, init_hidden=True),
nn.Dropout2d()
)
def forward(self, x):
UT.reset(self)
out_spks= self.double_conv(x)
return out_spks
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels, beta, SG_func):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels, beta=beta, SG_func=SG_func)
)
def forward(self, x):
return self.maxpool_conv(x)
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, beta, SG_func, bilinear=True):
super().__init__()
# transposed conv option omiited, only bilinear used instead
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.conv = DoubleConv(in_channels, out_channels, beta=beta, SG_func=SG_func, mid_channels=in_channels // 2)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is BxCxHxW
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class OutConv(nn.Module):
"""Output conv with 1x1 kernel"""
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)
# full Unet model
class UNet(nn.Module):
def __init__(self, n_channels, n_classes, bilinear=True, beta=0.5, SG_func=surrogate.ATan.apply):
super(UNet, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.beta = beta
self.SG_func = SG_func
self.inc = DoubleConv(n_channels, 64, beta, SG_func)
self.down1 = Down(64, 128, beta, SG_func)
self.down2 = Down(128, 256, beta, SG_func)
self.down3 = Down(256, 512, beta, SG_func)
self.down4 = Down(512, 1024 // 2, beta, SG_func)
self.up1 = Up(1024, 512 // 2, beta, SG_func)
self.up2 = Up(512, 256 // 2, beta, SG_func)
self.up3 = Up(256, 128 // 2,beta, SG_func)
self.up4 = Up(128, 64, beta, SG_func)
self.outc = OutConv(64, n_classes)
def forward(self, x):
spikes = []
torch.cuda.empty_cache()
UT.reset(self)
for step in range(num_steps):
x1 = self.inc(x[step])
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.up1(x5, x4)
x7 = self.up2(x6, x3)
x8 = self.up3(x7, x2)
x9 = self.up4(x8, x1)
spikes.append(x9)
x1 = x1.detach()
x2 = x2.detach()
x3 = x3.detach()
x4 = x4.detach()
x5 = x5.detach()
x6 = x6.detach()
x7 = x7.detach()
x8 = x8.detach()
x9 = x9.detach()
accumulated_spks=torch.sum(torch.stack(spikes), dim=0)
accumulated_spks=accumulated_spks.detach()
output = self.outc(accumulated_spks)
return output
# very simple training loop just to mess around with
def training(loss_f, model, Optimizer, dataloader):
training_acc = 0
batch_no = 0
model.train() # set network to training mode
for batch in dataloader:
torch.cuda.empty_cache()
print(f"iteration: {batch_no}")
start = time.time()
X = spikegen.rate(batch["X"], num_steps=num_steps).to(device) # ensure network and data are running on GPU if available
Y = batch["Y"].to(device)
pred = model(X) # compute network output
loss = loss_f(pred, Y) # compute loss and backpropagate on it
Optimizer.zero_grad()
loss.backward()
Optimizer.step()
stop = time.time()
batch_no += 1
if batch_no == 5:
break;
if batch_no % 50 == 0:
print(f"Loss at batch no: {batch_no}: {loss}")
PATH = os.path.join(root_dir, "model"+".pth")
model = UNet(n_channels = 4, n_classes = 1).to(device)
torch.save(model.state_dict(), PATH)
loss_f = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
training(loss_f, model, optimizer, train_dataloader)
我尝试添加 torch.cuda.empty_cache() 并使用 snnTorch utils.reset(model) 清除任何不必要的变量或状态堵塞内存，但无济于事。我是 PyTorch 的新手，所以我不太清楚跟踪的情况，但我认为它看起来可能还不错。 idk 我真的很感谢大家的帮助，提前谢谢你们。
我发现问题了。 snnTorch 的 util.reset() 没有为 double_conv 正确实现，我传递了 self，但我应该传递 self.double_conv。因此，尖峰神经元的隐藏状态使记忆膨胀并且无法清除。