drbh

Migrated from kernels-community/deformable-detr

249fa33 unverified about 1 month ago

7.84 kB

	import torch
	import torch.nn.functional as F

	from kernels.benchmark import Benchmark


	def ms_deform_attn_reference(
	value: torch.Tensor,
	spatial_shapes: torch.Tensor,
	level_start_index: torch.Tensor,
	sampling_locations: torch.Tensor,
	attention_weights: torch.Tensor,
	) -> torch.Tensor:
	batch, _, num_heads, channels = value.shape
	_, num_query, _, num_levels, num_points, _ = sampling_locations.shape

	# Split value by levels
	value_list = []
	for level_id in range(num_levels):
	H, W = spatial_shapes[level_id]
	start_idx = level_start_index[level_id]
	end_idx = (
	level_start_index[level_id + 1]
	if level_id < num_levels - 1
	else value.shape[1]
	)
	# (batch, H*W, num_heads, channels) -> (batch, num_heads, channels, H, W)
	value_level = value[:, start_idx:end_idx, :, :].view(
	batch, H, W, num_heads, channels
	)
	value_level = value_level.permute(0, 3, 4, 1, 2).contiguous()
	value_list.append(value_level)

	# Sample from each level
	output = torch.zeros(
	batch, num_query, num_heads, channels, device=value.device, dtype=value.dtype
	)

	for level_id in range(num_levels):
	H, W = spatial_shapes[level_id]
	value_level = value_list[level_id] # (batch, num_heads, channels, H, W)

	# Get sampling locations for this level: (batch, num_query, num_heads, num_points, 2)
	sampling_loc_level = sampling_locations[:, :, :, level_id, :, :]

	# Convert from [0, 1] to [-1, 1] for grid_sample
	grid = (
	2.0 * sampling_loc_level - 1.0
	) # (batch, num_query, num_heads, num_points, 2)

	# Reshape for grid_sample: need (batch * num_heads, channels, H, W) and (batch * num_heads, num_query, num_points, 2)
	value_level = value_level.view(batch * num_heads, channels, H.item(), W.item())
	grid = grid.permute(
	0, 2, 1, 3, 4
	).contiguous() # (batch, num_heads, num_query, num_points, 2)
	grid = grid.view(batch * num_heads, num_query, num_points, 2)

	# Sample: output is (batch * num_heads, channels, num_query, num_points)
	sampled = F.grid_sample(
	value_level,
	grid,
	mode="bilinear",
	padding_mode="zeros",
	align_corners=False,
	)

	# Reshape back: (batch, num_heads, channels, num_query, num_points)
	sampled = sampled.view(batch, num_heads, channels, num_query, num_points)
	# -> (batch, num_query, num_heads, num_points, channels)
	sampled = sampled.permute(0, 3, 1, 4, 2).contiguous()

	# Get attention weights for this level: (batch, num_query, num_heads, num_points)
	attn_level = attention_weights[:, :, :, level_id, :]

	# Weighted sum over points: (batch, num_query, num_heads, channels)
	output += (sampled * attn_level.unsqueeze(-1)).sum(dim=3)

	# Reshape to (batch, num_query, num_heads * channels)
	output = output.view(batch, num_query, num_heads * channels)
	return output


	class MSDeformAttnBenchmark(Benchmark):
	seed: int = 42

	def setup(self):
	batch = 2
	num_heads = 8
	channels = 32 # embed_dim = num_heads * channels = 256
	num_levels = 4
	num_query = 300
	num_points = 4
	im2col_step = 64

	# Spatial shapes for 4 levels: 64x64, 32x32, 16x16, 8x8
	spatial_shapes = torch.tensor(
	[[64, 64], [32, 32], [16, 16], [8, 8]],
	dtype=torch.int64,
	device=self.device,
	)
	# Calculate spatial_size = sum of H*W for all levels
	spatial_size = (64 * 64) + (32 * 32) + (16 * 16) + (8 * 8) # 5440

	# Level start indices
	level_start_index = torch.tensor(
	[0, 64 * 64, 64 * 64 + 32 * 32, 64 * 64 + 32 * 32 + 16 * 16],
	dtype=torch.int64,
	device=self.device,
	)

	self.value = torch.randn(
	batch,
	spatial_size,
	num_heads,
	channels,
	device=self.device,
	dtype=torch.float32,
	)
	self.spatial_shapes = spatial_shapes
	self.level_start_index = level_start_index
	self.sampling_loc = torch.rand(
	batch,
	num_query,
	num_heads,
	num_levels,
	num_points,
	2,
	device=self.device,
	dtype=torch.float32,
	)
	self.attn_weight = torch.rand(
	batch,
	num_query,
	num_heads,
	num_levels,
	num_points,
	device=self.device,
	dtype=torch.float32,
	)
	# Normalize attention weights
	self.attn_weight = self.attn_weight / self.attn_weight.sum(-1, keepdim=True)
	self.im2col_step = im2col_step

	self.out = torch.empty(
	batch,
	num_query,
	num_heads * channels,
	device=self.device,
	dtype=torch.float32,
	)

	def benchmark_forward(self):
	self.out = self.kernel.ms_deform_attn_forward(
	self.value,
	self.spatial_shapes,
	self.level_start_index,
	self.sampling_loc,
	self.attn_weight,
	self.im2col_step,
	)

	def verify_forward(self) -> torch.Tensor:
	return ms_deform_attn_reference(
	self.value,
	self.spatial_shapes,
	self.level_start_index,
	self.sampling_loc,
	self.attn_weight,
	)

	def setup_large(self):
	batch = 8
	num_heads = 8
	channels = 32
	num_levels = 4
	num_query = 900
	num_points = 4
	im2col_step = 64

	spatial_shapes = torch.tensor(
	[[64, 64], [32, 32], [16, 16], [8, 8]],
	dtype=torch.int64,
	device=self.device,
	)
	spatial_size = (64 * 64) + (32 * 32) + (16 * 16) + (8 * 8)

	level_start_index = torch.tensor(
	[0, 64 * 64, 64 * 64 + 32 * 32, 64 * 64 + 32 * 32 + 16 * 16],
	dtype=torch.int64,
	device=self.device,
	)

	self.value = torch.randn(
	batch,
	spatial_size,
	num_heads,
	channels,
	device=self.device,
	dtype=torch.float32,
	)
	self.spatial_shapes = spatial_shapes
	self.level_start_index = level_start_index
	self.sampling_loc = torch.rand(
	batch,
	num_query,
	num_heads,
	num_levels,
	num_points,
	2,
	device=self.device,
	dtype=torch.float32,
	)
	self.attn_weight = torch.rand(
	batch,
	num_query,
	num_heads,
	num_levels,
	num_points,
	device=self.device,
	dtype=torch.float32,
	)
	self.attn_weight = self.attn_weight / self.attn_weight.sum(-1, keepdim=True)
	self.im2col_step = im2col_step

	self.out = torch.empty(
	batch,
	num_query,
	num_heads * channels,
	device=self.device,
	dtype=torch.float32,
	)

	def benchmark_large(self):
	self.out = self.kernel.ms_deform_attn_forward(
	self.value,
	self.spatial_shapes,
	self.level_start_index,
	self.sampling_loc,
	self.attn_weight,
	self.im2col_step,
	)

	def verify_large(self) -> torch.Tensor:
	return ms_deform_attn_reference(
	self.value,
	self.spatial_shapes,
	self.level_start_index,
	self.sampling_loc,
	self.attn_weight,
	)