Hello everybody!
I try to write and train a deep multiple instance learning model for 3D image classification, I always meet the question of “out of memory”. And cannot find out what leads to this problem.
Since the 3D images are too large, I segment each 3D image into overlapped patches. So my input is an array of 3D patches, and my output is a label.
In the training process, each patch is as the input of a 3D CNN with the Softmax layer, and the 3D CNN outputs probability of each patch. After obtaining the probabilities of all the patches of a 3D image, the deep multiple instance leanring model selects the patches with the max probability to proceed back-propagation. During the process of back-propagation, I always get the problem of “cudaError: out of memory”. My codes are as bellow:
    self.model.train(True)

    running_loss, running_corrects = 0,0
    patch_unit_size = 200

    for pat_batch in self.dcm_datasets['train']:
        inputs, labels, data_dir = pat_batch   
       # the shape of inputs are [batch_numr, patch_num, channel_num, path_height, patch_width, patch_length], where batch_num=1, patch_num varies with differente 3D images, channel_num=3, patch_height=patch_width=patch_length=24

        for input_each_batch in inputs:
            patches_size = input_each_batch.shape[0]
            num_patches = math.ceil(1.0 * patches_size / patch_unit_size)

            patch_out_max, patch_out_prob_max = None, None
            for i in range(num_patches):  # Find out the patches with the maximum probability
                inputs_tmp = input_each_batch[i * self.patch_size: (i + 1) *
                          self.patch_size] if i < num_patches - 1 else input_each_batch[i * self.patch_size:]
                
                with torch.cuda.device(self.cuda_ids[0]):
                    inputs_new = Variable(inputs_tmp.cuda()).to(self.cuda_ids[0])
                    labels_new = Variable(labels.cuda()).to((self.cuda_ids[0]))

                outputs = self.model(inputs_new)  # 1*2
                out_probs = torch.nn.functional.softmax(outputs, dim=1).data
                patch_out_prob_max = out_probs if patch_out_prob_max is None else torch.cat(
                    (out_probs, patch_out_prob_max), dim=0)

                '''find all the indices of the maximum'''
                patch_prob_max = patch_out_prob_max.cpu().numpy()
                inds_x, inds_y = np.where(patch_prob_max == np.max(patch_prob_max))
                patch_out_prob_max = patch_out_prob_max[inds_x]

                inds_x_out = inds_x[inds_x < outputs.shape[0]]
                inds_x_rest = inds_x[inds_x >= outputs.shape[0]]
                if inds_x_rest.shape[0] != 0:
                    inds_x_rest = inds_x_rest - outputs.shape[0]
                    patch_out_max = patch_out_max[inds_x_rest]
                if inds_x_out.shape[0] != 0:
                    patch_out_max_tmp = outputs[inds_x_out]
                    if inds_x_rest.shape[0] != 0:
                        patch_out_max = torch.cat((patch_out_max, patch_out_max_tmp), dim=0)
                    else:
                        patch_out_max = patch_out_max_tmp

                outputs, out_probs, inputs_new = 0, 0, 0

            labels_new_1 = None
            for i in range(patch_out_max.shape[0]):
                labels_new_1 = labels_new if labels_new_1 is None else torch.cat((labels_new_1, labels_new),dim=0)

            loss = self.criterion(patch_out_max, labels_new_1)
            self.optimizer.zero_grad()  # zero the parameter gradients
            loss.backward()
            self.optimizer.step()

            preds = torch.argmax(patch_out_max[0].data)  # preds is still a tensor
            running_loss += loss.item()  # running_loss is a Python data
            running_corrects += np.sum(preds.item() == labels_new.item())

            loss, labels_new_1 = 0, None
            torch.cuda.empty_cache()
            torch.cuda.memory_allocated()

    data_len = len(self.dcm_datasets['train'])
    epoch_loss = running_loss / data_len
    epoch_acc = running_corrects / data_len

    return epoch_acc, epoch_loss

Since the patch_num may be higher to 1500, I am wondering whether the “CudaError: out of memory” is caused by the large computation graph? But I am not sure.
So my question is how large is my computation graph? Providing the computation graph of a self.patch_size is O, my computation graph is O or num_patches*O?
If my computation graph is just O, the cuda usage is not too much. In this case, what leads to “cuda out of memory”?