RuntimeError: CUDA out of memory怎么办?求助,救急

import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
#data
train_data = dataset.MNIST(root="mnist",train=True,transform=transforms.ToTensor(),download=True)
test_data = dataset.MNIST(root="mnist",train=False,transform=transforms.ToTensor(),download=False)

#batchsize
train_loader = data_utils.DataLoader(dataset=train_data,batch_size=64,shuffle=True)
test_loader = data_utils.DataLoader(dataset=test_data,batch_size=64,shuffle=True)

#Net
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, kernel_size=5, padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.fc = torch.nn.Linear(14 * 14 * 32,10)
    def forward(self,x):
        out=self.conv(x)
        out=out.view(out.size()[0],-1)
        out=self.fc(out)
        return out
cnn=CNN()
cnn = cnn.cuda()
#loss

loss_func = torch.nn.CrossEntropyLoss()

#optimizer

optimizer = torch.optim.Adam(cnn.parameters(), lr=0.01)

#training
for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        images = images.cuda()
        labels = labels.cuda()

        outputs = cnn(images)
        loss = loss_func(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print("epoch is {}, ite is {}/{}, loss is {}".format(epoch+1, i, len(train_data)//64, loss.item()))


#eval/test
    loss_test = 0
    accuracy = 0
    for i, (images, labels) in enumerate(test_loader):
        images = images.cuda()
        labels = labels.cuda()
        outputs = cnn(images)
        loss_test += loss_func(outputs, labels)
        _, pred = outputs.max(1)
        accuracy = (pred == labels).sum().item()
    accuracy = accuracy / len(test_data)
    loss_test = loss_test / (len(test_data) // 64)
    print("epoch is {}, accuracy is {}, loss is {}".format(epoch + 1, accuracy, loss_test.item()))

#save
torch.save(cnn, "model/mnist_model.pkl")

上面就是我的代码,然后它的运行结果报错:

具体报错如下:

求哪位大佬告知一下该怎么办,救急救救孩子!! 

评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值