使用Colab训练数据时，CUDA总是会内存不足

Question

所以我有2DCNN模型对图像进行分类，只有2个类，每个类有300个图像。 这是我的nn模块类

class Flatten(Module):

    def forward(self, input):

        return input.view(input.size(0), -1)

class ActionNet(Module):

    def __init__(self, num_class=4):

        super(ActionNet, self).__init__()



        self.cnn_layer = Sequential(

            #conv1

            Conv2d(in_channels=1, out_channels=32, kernel_size=1, bias=False),

            BatchNorm2d(32),

            PReLU(num_parameters=32),

            MaxPool2d(kernel_size=3),

            #conv2

            Conv2d(in_channels=32, out_channels=64, kernel_size=1, bias=False),

            BatchNorm2d(64),

            PReLU(num_parameters=64),

            MaxPool2d(kernel_size=3),

            #flatten

            Flatten(),

            Linear(576, 128),

            BatchNorm1d(128),

            ReLU(inplace=True),

            Dropout(0.5),

            Linear(128, num_class)

        )



    def forward(self, x):

        x = self.cnn_layer(x)

        return x

我将训练和测试分为80％：20％，有关更多详细信息，这是我的训练代码：https://pastebin.com/uzBFTrDc，当我尝试使用colab进行训练时出现错误：

RuntimeError：CUDA内存不足。尝试分配20.00 MiB（GPU 0； 15.90 GiB总容量；已分配15.20 GiB；免费1.88 MiB； PyTorch总共保留15.20 GiB）

我使用google colab是因为我没有强大的GPU并实现了我不知道它是否正确的批处理，训练数据只是大小为32x32的400张图像，我认为colabs足以训练我的数据。