torch.nn.functional.conv2d для нескольких каналов / пакетов

Question

У меня есть изображение, которое я хочу заполнить (чтобы сохранить ту же форму), а затем выполнить свертку с данным ядром. Это работает нормально, если у меня есть только один канал и одно изображение в пакете. Но как правильно переписать его для нескольких партий и каналов? Я полагаю, что для пакетов я могу просто дублировать ядро ​​по измерению 0. Но как насчет каналов? Как правильно сделать это в факеле? См. Пример с игрушкой ниже.

import torch.nn.functional as f

x = torch.zeros((1,1,16,16))
x[...,6:10,6:10] = 1.
ker = torch.ones(3,3)
ker[1,1] = -4
padding = (ker.shape[1] // 2, ker.shape[1] // 2, 
           ker.shape[0] // 2, ker.shape[0] // 2)
x = f.pad(x, padding, mode='replicate')    
kernel = ker.reshape((1, 1, ker.shape[0], -1))
result = f.conv2d(x, ker, groups=1)

Как переписать этот фрагмент кода для работы с несколькими каналами и пакетами? Я читаю документацию, но, честно говоря, она мне не показалась очень подробной.

Answer 1

Документация на https://pytorch.org/docs/stable/nn.functional.html#torch .nn.functional.conv2d , кажется, отвечает на ваш вопрос:

 input – input tensor of shape (minibatch,in_channels,iH,iW)
 weight – filters of shape (out_channels,in_channels/groups,kH,kW)

, поэтому ваш x должен быть размером (batch_size, in_channels, 16, 16 )

и ваше ядро ​​(batch_size, out_channels / groups, kH, kW)