下采样方法

AntiAliasInterpolation2d代码解读

注记

最近在看一些视频驱动的代码时，常见一种特殊的下采样方法，故在这里记录一下。

Class AntiAliasInterpolation2d(nn.Module):
	## 初始化
	def __init__(self, channels, scale):
		## channels
		## scale: 下采样比例 取 < 1， 这里假设取0.5
		super(AntiAliasInterpolation2d, self).__init__()
		sigma = (1 / scale - 1) / 2  ## 0.5
		kernel_size = 2* round(sigma * 4) + 1 ## 5
		self.ka = kernel_size // 2 ## 2
		self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka ## 2
		kernel_size = [kernel_size, kernel_size] ##[5, 5]
		sigma = [sigma, sigma]   ## [0.5, 0.5] 
		kernel = 1
		meshgrids = torch.mershgrid(   ## size: (2, 5, 5)
			[	
				torch.arange(size, dtype=torc.float32)
					for size in kernel_size	
				]
		)
		for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
			mean = (size - 1)  / 2   ## 2
			## 高斯处理（标准化，使得数据在处理之后，符合正态分布，符合模型推理要求）
			kernel *= torch.exp(-(mgrid - mean)**2 / (2 * std **2)) ## shape: [5, 5]
		## 保证卷积核中元素和为1
		kernel = kernel / torch.sum(kernel) ## shape: [5, 5]
		kernel = kernel.view(1, 1, *kernel.size()) ##shape[1, 1, 5, 5]
		## 数组广播
		kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) ## shape:[channels, 1, 5, 5],这里第二个维度由(in_channel/group)计算得到
		## 向模块添加持久缓冲区（将kennel加入到“weight”中，weight不参与训练）
		self.register_buffer("weight", kernel)
		self.groups = channels ## group=channel,为deepwise
		self.scale = scale
		inv_scale = 1 / scale
		self.int_inv_scale = int(inv_scale)
	
	## 推理
	def forward(self, input):
		if self.scale == 1.0:
			return input
		## 填充（conv 操作不改变size）
		out = F.pad(input, (self.ka, self.kb, self.ka, self,kb))
		## 卷积处理
		out = F.conv2d(out, weight=self.weight, groups=self.groups)
		## 下采样，int_inv_scale为scale 倒数的整数
		out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale]
		return out