自己实现的Conv1d一维卷积,运行速度很慢,效果较差,有没有人帮忙改进一下
class Conv1d(nn.Module):
def __init__(self, in_channels:int, out_channels:int, kernel_size:int=3,
stride:int=1, padding:int=0, bias:bool=True) -> None:
super(Conv1d, self).__init__()
self.stride = stride
self.padding = padding
self.ker_size = kernel_size
self.out_channels = out_channels
self.kernel = nn.Parameter(pt.randn(out_channels, in_channels, kernel_size))
self.bias = None
if bias:
self.bias = nn.Parameter(pt.randn(out_channels))
def forward(self, x):
# x:(batch_size, in_channel, length)
out_len = (x.shape[2] + 2*self.padding - self.ker_size)//self.stride + 1
out = pt.empty(x.shape[0], self.out_channels, out_len, device='cuda:0')
x = nn.functional.pad(x, (self.padding, self.padding), "constant", 0)
for l in range(out_len):
# bik:(batch_size,in_ch,ker_size) oik:(out_ch,in_ch,ker_size)
out[:,:,l] = pt.einsum('bik, oik -> bo',
x[:,:,l*self.stride:l*self.stride+self.ker_size], self.kernel)
if self.bias is not None:
b = pt.unsqueeze(self.bias, 1)
b = b.repeat(x.shape[0], 1, out_len) # (out_chan,1)->(batch_size,out_chan,out_len)
out += b
return out