Unable to produce full resolution output from trained model due to lack of CUDA memory

Ma_John: <p>Hello to all.<br>
I have a trained model that is only able to accept inputs of size <strong>1024 * 1024</strong>(which is resized from the original image) and produce outputs of the same size.<br>
However, I need it to produce outputs of full resolutions, which is no bigger than <strong>5000 * 5000.</strong> I wonder if <strong>my implementation is wrong</strong> or <strong>my model is simply too big</strong> to produce such output? (According to my knowledge that should not be the case, cause I’ve seen others produce <strong>1024 * 1024</strong> outputs with a model that’s composed of <strong>100+</strong> conv layers.)<br>
Any advice would be appreciated!</p>
<p>Here are my model, and my testing procedure.</p>
<pre><code class="lang-auto">class model(nn.Module):
    def __init__(self, pretrain=False, IENetPath="", InvertIENetPath=""):
        super(model, self).__init__()
        self.IENetwork = UNet(3, 3) # UNet is composed of about 20 conv layers
        self.InvertedIENetwork = UNet(3, 3)

        if pretrain:
            self.IENetwork.load_state_dict(torch.load(IENetPath))
            self.InvertedIENetwork.load_state_dict(torch.load(InvertIENetPath))

        self.conv = conv3x3_ReLU(3, 64)
        self.attention = AttentionModule(128, 64) # AttentionModule is compsed of 2 conv layers
        self.conv2 = conv3x3_ReLU(192, 64)
        self.DRDB1 = RDB(64, 6, 32) # RDB is composed of 6 conv layers
        self.DRDB2 = RDB(64, 6, 32)
        self.DRDB3 = RDB(64, 6, 32)
        self.conv3 = conv3x3_ReLU(192, 64)
        self.conv4 = conv3x3_ReLU(64, 64)
        self.conv5 = conv3x3_ReLU(64, 3)
</code></pre>
<pre><code class="lang-auto">with torch.no_grad():
            dataset = mydataset(ROOT, TEST_DIR, LABEL_DIR, TEST_IMG_NUM, TEST_TRANSFORM, training=True)
            model = model().cuda()
            testloader = data.DataLoader(dataset, batch_size=1)

            model.load_state_dict(torch.load(MODEL_PATH))
            model.eval()

            if not os.path.exists(RES_DIR):
                os.makedirs(RES_DIR)

            for i, data in enumerate(testloader):
                Input, Inverted_Input, Label = data['train_img'], data['inverted_img'], data['label_img']

                savefolder = RES_DIR + "/test" + str(i)

                # Images saving
                SaveImg(Input, savefolder + "/1input.jpg")
                SaveImg(Inverted_Input, savefolder + "/2inverted.jpg")


                Output, IEmap, IEmap_Inverted = model(Input, Inverted_Input, Label, path=savefolder, training=TRAIN)
                SaveImg(Output, savefolder + "/7finalRes.jpg")
                SaveImg(Label, savefolder + "/8label.jpg")

            print("TestEnds")
</code></pre>

Ma_John: <p>It is worth to note that the testing procedure has no problem reading in data of full resolution, but when inferring the input (i.e. <strong>model.forward()</strong>) cuda out of memory error would occur.</p>

ptrblck: <p>Input tensors of <code>[5000, 5000]</code> are quite huge. How much memory was used for a single sample of your [<code>1024, 1024]</code> input? Note that the <code>[5000, 5000]</code> has approx. 24x more elements.</p>
<p>Also, have a look at <a href="https://pytorch.org/docs/stable/checkpoint.html" rel="nofollow noopener">torch.utils.checkpoint</a> to trade compute for memory.</p>

Ma_John: <p>A single sample of [ <code>1024, 1024]</code>  input takes up almost my entire GTX1080TI GPU, that’s roughly 10G+ memory.<br>
Also, I wonder if there is any implementation error in my code? Cause I’ve seen a bigger model that can handle  [ <code>1024, 1024]</code>  input with less memory.</p>

ptrblck: <p>Based on your code snippet I cannot see anything obviously wrong, but since you are using custom sumbodules, I’m not sure if they contain any bugs.</p>
<p>The memory usage depends on a lot of factors besides the number of layers.<br>
E.g. if you pool your activations quickly, you could use a lot of more layers.</p>