CNN that only looks at local area

1

I’m trying to look at local orientation in images. I have a set of images which have the size 150x150x150 and for each voxel I would like to produce a unit vector i.e. an output of size 3x150x150x150 for each image. Each of these unit vectors should reflect the dominant orientation in its neighborhood, in my case the size of this neighborhood is 17x17x17 (kernel_size = 17), which means my model shouldn’t care about other values which are outside of this neighborhood.

I tried naively setting up a series of 3D convolutions and trained on a small test-set. But I’m pretty sure that I’m already reaching beyond my neighborhood in my second convolution.

class CNN3(torch.nn.Module):
    def __init__(self,kernel_size):
        super(CNN3, self).__init__()
        self.k = kernel_size
        self.pad = kernel_size//2
        self.conv1 = torch.nn.Conv3d(1, 3, kernel_size=self.k, stride=1, padding=self.pad)
        self.conv2 = torch.nn.Conv3d(3, 3, kernel_size=self.k, stride=1, padding=self.pad)
        self.conv3 = torch.nn.Conv3d(3, 3, kernel_size=self.k, stride=1, padding=self.pad)
        self.conv4 = torch.nn.Conv3d(3, 3, kernel_size=self.k, stride=1, padding=self.pad)
    def forward(self,x):
        out = F.relu(self.conv1(x))
        out = F.relu(self.conv2(out))
        out = F.relu(self.conv3(out))
        out = self.conv4(out)
        return out

Is there a way to set up a CNN to only look at a particular neighborhood or is it simply not the right tool for the job?