import torch.nn as nn
class Masking_CNN(nn.Module):
def __init__(self,input_channels):
"""
input_channels: int, number of channels in the input
"""
super().__init__()
#describe the network
self.conv_layer= nn.Conv2d(input_channels,4,kernel_size=3,padding=1)
def forward(self,x,mask):
"""
x: pytorch tensor of size(batch_size,input_channels,input_size,input_size)
mask: pytorch tensor of size(1,4,input_size,input_size) containing either 0 or 1
"""
x= self.conv_layer(x)
x= torch.mul(x,mask)
return x
1条答案
按热度按时间qni6mghb1#
您是否检查了torch.mul函数(https://pytorch.org/docs/stable/generated/torch.mul.html)?它将对两个大小相等的输入执行哈达玛乘积(对于大小不相等的输入,将使用广播)。按照惯例,在要求_grad=True的Tensor上执行此函数,梯度将被正确计算并存储在Tensor中。
由于你没有提供足够的信息来重建/完全理解你想要解决的问题,我只是举了一个例子。假设你在pytorch中有一个简单的CNN,但是(不管出于什么原因)你想在某些区域屏蔽掉层的输出。这可以通过以下方式实现:
因此,这将是一个如何将函数集成到网络中的例子。您所需要的只是将torch.穆尔函数添加到前向传递中,并将您希望按元素乘的Tensor(哈达玛乘积)作为参数传递。
如果您需要更多关于您的问题的缩小范围的帮助,请指定您的问题并提供关于您的问题的更多信息。