24_EfficientNetB2 • efficientnetb2_shared_memory_warp_reduce_baseimport torch
import torch.nn as nn
import torch.nn.functional as F
def module_fn(
x: torch.Tensor, params: nn.ParameterDict, is_training: bool
) -> torch.Tensor:
"""
Implementation of EfficientNetB2
Args:
x: Input tensor of shape (batch_size, 3, 224, 224).
params: A nn.ParameterDict containing model parameters.
is_training: Whether the model is in training mode.
Returns:
torch.Tensor: Output tensor of shape (batch_size, 1000).
"""
# Initial conv
x = F.conv2d(x, params["conv1_weight"], None, stride=2, padding=1)
x = F.batch_norm(
x,
params["bn1_mean"],
params["bn1_var"],
params["bn1_weight"],
params["bn1_bias"],
is_training,
)
x = F.relu(x, inplace=True)
def mbconv_block_fn(x, params, stride, expand_ratio, is_training):
"""
Functional implementation of MBConv block
"""
in_channels = x.size(1)
expanded_channels = in_channels * expand_ratio
# Expansion phase
if expand_ratio != 1:
x = F.conv2d(x, params["expand_conv_weight"], None)
x = F.batch_norm(
x,
params["expand_bn_mean"],
params["expand_bn_var"],
params["expand_bn_weight"],
params["expand_bn_bias"],
is_training,
)
x = F.relu(x, inplace=True)
else:
expanded_channels = in_channels
# Depthwise conv
x = F.conv2d(
x,
params["dw_conv_weight"],
None,
stride=stride,
padding=1,
groups=expanded_channels,
)
x = F.batch_norm(
x,
params["dw_bn_mean"],
params["dw_bn_var"],
params["dw_bn_weight"],
params["dw_bn_bias"],
is_training,
)
x = F.relu(x, inplace=True)
# Squeeze and Excitation
se = F.adaptive_avg_pool2d(x, (1, 1))
se = F.conv2d(se, params["se_reduce_weight"], None)
se = F.relu(se, inplace=True)
se = F.conv2d(se, params["se_expand_weight"], None)
se = torch.sigmoid(se)
x = se
# x = x * se
# Output phase
x = F.conv2d(x, params["project_conv_weight"], None)
x = F.batch_norm(
x,
params["project_bn_mean"],
params["project_bn_var"],
params["project_bn_weight"],
params["project_bn_bias"],
is_training,
)
return x
# MBConv blocks
mbconv_configs = [(1, 3), (2, 6), (2, 6), (2, 6), (1, 6)]
for i, (stride, expand_ratio) in enumerate(mbconv_configs, 1):
block_params = {
k.replace(f"mbconv{i}_", ""): v
for k, v in params.items()
if k.startswith(f"mbconv{i}_")
}
x = mbconv_block_fn(x, block_params, stride, expand_ratio, is_training)
# Final layers
x = F.conv2d(x, params["conv_final_weight"], None)
x = F.batch_norm(
x,
params["bn_final_mean"],
params["bn_final_var"],
params["bn_final_weight"],
params["bn_final_bias"],
is_training,
)
x = F.relu(x, inplace=True)
x = F.adaptive_avg_pool2d(x, (1, 1))
x = torch.flatten(x, 1)
x = F.linear(x, params["fc_weight"], params["fc_bias"])
return x
class Model(nn.Module):
def __init__(self, num_classes=1000):
super(Model, self).__init__()
# Create the original model to ensure identical initialization
original_model = nn.Module()
original_model.conv1 = nn.Conv2d(
3, 32, kernel_size=3, stride=2, padding=1, bias=False
)
original_model.bn1 = nn.BatchNorm2d(32)
original_model.relu = nn.ReLU(inplace=True)
# MBConv blocks
configs = [
(32, 96, 1, 3),
(96, 144, 2, 6),
(144, 192, 2, 6),
(192, 288, 2, 6),
(288, 384, 1, 6),
]
for i, (in_c, out_c, stride, expand) in enumerate(configs, 1):
expanded_c = in_c * expand
block = nn.Sequential()
if expand != 1:
block.add_module(
"expand_conv", nn.Conv2d(in_c, expanded_c, 1, bias=False)
)
block.add_module("expand_bn", nn.BatchNorm2d(expanded_c))
block.add_module("expand_relu", nn.ReLU(inplace=True))
block.add_module(
"dw_conv",
nn.Conv2d(
expanded_c,
expanded_c,
3,
stride=stride,
padding=1,
groups=expanded_c,
bias=False,
),
)
block.add_module("dw_bn", nn.BatchNorm2d(expanded_c))
block.add_module("dw_relu", nn.ReLU(inplace=True))
block.add_module("se_pool", nn.AdaptiveAvgPool2d((1, 1)))
block.add_module(
"se_reduce", nn.Conv2d(expanded_c, expanded_c // 4, 1, bias=False)
)
block.add_module("se_reduce_relu", nn.ReLU(inplace=True))
block.add_module(
"se_expand", nn.Conv2d(expanded_c // 4, expanded_c, 1, bias=False)
)
block.add_module("se_sigmoid", nn.Sigmoid())
block.add_module(
"project_conv", nn.Conv2d(expanded_c, out_c, 1, bias=False)
)
block.add_module("project_bn", nn.BatchNorm2d(out_c))
setattr(original_model, f"mbconv{i}", block)
original_model.conv_final = nn.Conv2d(384, 1408, 1, bias=False)
original_model.bn_final = nn.BatchNorm2d(1408)
original_model.avgpool = nn.AdaptiveAvgPool2d((1, 1))
original_model.fc = nn.Linear(1408, num_classes)
# Initialize parameters and buffers
self.params = nn.ParameterDict()
# Copy initial conv parameters
self.params["conv1_weight"] = nn.Parameter(original_model.conv1.weight.data)
self.params["bn1_weight"] = nn.Parameter(original_model.bn1.weight.data)
self.params["bn1_bias"] = nn.Parameter(original_model.bn1.bias.data)
self.register_buffer("bn1_mean", original_model.bn1.running_mean)
self.register_buffer("bn1_var", original_model.bn1.running_var)
# Copy MBConv block parameters
for i in range(1, 6):
block = getattr(original_model, f"mbconv{i}")
prefix = f"mbconv{i}_"
if hasattr(block, "expand_conv"):
self.params[prefix + "expand_conv_weight"] = nn.Parameter(
block.expand_conv.weight.data
)
self.params[prefix + "expand_bn_weight"] = nn.Parameter(
block.expand_bn.weight.data
)
self.params[prefix + "expand_bn_bias"] = nn.Parameter(
block.expand_bn.bias.data
)
self.register_buffer(
prefix + "expand_bn_mean", block.expand_bn.running_mean
)
self.register_buffer(
prefix + "expand_bn_var", block.expand_bn.running_var
)
self.params[prefix + "dw_conv_weight"] = nn.Parameter(
block.dw_conv.weight.data
)
self.params[prefix + "dw_bn_weight"] = nn.Parameter(block.dw_bn.weight.data)
self.params[prefix + "dw_bn_bias"] = nn.Parameter(block.dw_bn.bias.data)
self.register_buffer(prefix + "dw_bn_mean", block.dw_bn.running_mean)
self.register_buffer(prefix + "dw_bn_var", block.dw_bn.running_var)
self.params[prefix + "se_reduce_weight"] = nn.Parameter(
block.se_reduce.weight.data
)
self.params[prefix + "se_expand_weight"] = nn.Parameter(
block.se_expand.weight.data
)
self.params[prefix + "project_conv_weight"] = nn.Parameter(
block.project_conv.weight.data
)
self.params[prefix + "project_bn_weight"] = nn.Parameter(
block.project_bn.weight.data
)
self.params[prefix + "project_bn_bias"] = nn.Parameter(
block.project_bn.bias.data
)
self.register_buffer(
prefix + "project_bn_mean", block.project_bn.running_mean
)
self.register_buffer(
prefix + "project_bn_var", block.project_bn.running_var
)
# Copy final layer parameters
self.params["conv_final_weight"] = nn.Parameter(
original_model.conv_final.weight.data
)
self.params["bn_final_weight"] = nn.Parameter(
original_model.bn_final.weight.data
)
self.params["bn_final_bias"] = nn.Parameter(original_model.bn_final.bias.data)
self.register_buffer("bn_final_mean", original_model.bn_final.running_mean)
self.register_buffer("bn_final_var", original_model.bn_final.running_var)
self.params["fc_weight"] = nn.Parameter(original_model.fc.weight.data)
self.params["fc_bias"] = nn.Parameter(original_model.fc.bias.data)
def forward(self, x, fn=module_fn):
params = {
**dict(self.params),
**{k: v for k, v in self._buffers.items() if v is not None},
}
return fn(x, params, self.training)
batch_size = 2
num_classes = 1000
def get_inputs():
return [torch.randn(batch_size, 3, 224, 224)]
def get_init_inputs():
return [num_classes]
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, num_classes=1000):
"""
EfficientNetB2 architecture implementation.
:param num_classes: The number of output classes (default is 1000 for ImageNet).
"""
super(Model, self).__init__()
# Define the EfficientNetB2 architecture components
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(32)
self.relu = nn.ReLU(inplace=True)
# Define the MBConv blocks
self.mbconv1 = self._make_mbconv_block(32, 96, 1, 3)
self.mbconv2 = self._make_mbconv_block(96, 144, 2, 6)
self.mbconv3 = self._make_mbconv_block(144, 192, 2, 6)
self.mbconv4 = self._make_mbconv_block(192, 288, 2, 6)
self.mbconv5 = self._make_mbconv_block(288, 384, 1, 6)
# Final layers
self.conv_final = nn.Conv2d(384, 1408, kernel_size=1, stride=1, padding=0, bias=False)
self.bn_final = nn.BatchNorm2d(1408)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(1408, num_classes)
def _make_mbconv_block(self, in_channels, out_channels, stride, expand_ratio):
"""
Helper function to create a MBConv block.
:param in_channels: Number of input channels.
:param out_channels: Number of output channels.
:param stride: Stride for the depthwise convolution.
:param expand_ratio: Expansion ratio for the MBConv block.
:return: A sequential container of layers forming the MBConv block.
"""
layers = []
expanded_channels = in_channels * expand_ratio
# Expansion phase
if expand_ratio != 1:
layers.append(nn.Conv2d(in_channels, expanded_channels, kernel_size=1, stride=1, padding=0, bias=False))
layers.append(nn.BatchNorm2d(expanded_channels))
layers.append(nn.ReLU(inplace=True))
# Depthwise convolution
layers.append(nn.Conv2d(expanded_channels, expanded_channels, kernel_size=3, stride=stride, padding=1, groups=expanded_channels, bias=False))
layers.append(nn.BatchNorm2d(expanded_channels))
layers.append(nn.ReLU(inplace=True))
# Squeeze and Excitation
layers.append(nn.AdaptiveAvgPool2d((1, 1)))
layers.append(nn.Conv2d(expanded_channels, expanded_channels // 4, kernel_size=1, stride=1, padding=0, bias=False))
layers.append(nn.ReLU(inplace=True))
layers.append(nn.Conv2d(expanded_channels // 4, expanded_channels, kernel_size=1, stride=1, padding=0, bias=False))
layers.append(nn.Sigmoid())
# Output phase
layers.append(nn.Conv2d(expanded_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False))
layers.append(nn.BatchNorm2d(out_channels))
return nn.Sequential(*layers)
def forward(self, x):
"""
Forward pass of the EfficientNetB2 model.
:param x: The input tensor, shape (batch_size, 3, 224, 224)
:return: The output tensor, shape (batch_size, num_classes)
"""
x = self.relu(self.bn1(self.conv1(x)))
x = self.mbconv1(x)
x = self.mbconv2(x)
x = self.mbconv3(x)
x = self.mbconv4(x)
x = self.mbconv5(x)
x = self.relu(self.bn_final(self.conv_final(x)))
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
# Test code
batch_size = 2
num_classes = 1000
def get_inputs():
return [torch.randn(batch_size, 3, 224, 224)]
def get_init_inputs():
return [num_classes]#include <torch/extension.h>
#include <map>
#include <string>
#include <vector>
using namespace torch;
__device__ float warp_reduce_sum(float val) {
for (int offset = warpSize / 2; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
return val;
}
__global__ void reduce_kernel(float* input, float* output, int N) {
extern __shared__ float shared_data[];
int tid = threadIdx.x;
int index = blockIdx.x * blockDim.x + tid;
float sum = 0.0f;
if (index < N) {
sum = input[index];
}
shared_data[tid] = sum;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
shared_data[tid] += shared_data[tid + s];
}
__syncthreads();
}
if (tid == 0) {
output[blockIdx.x] = shared_data[0];
}
}
Tensor mbconv_block(Tensor x, std::map<std::string, Tensor>& params, int stride, int expand_ratio, bool is_training) {
int64_t in_channels = x.size(1);
int64_t expanded_channels = in_channels * expand_ratio;
// Expansion phase
if (expand_ratio != 1) {
auto expand_conv_weight = params["expand_conv_weight"];
x = conv2d(x, expand_conv_weight, Tensor(),
{1}, // stride
at::IntArrayRef({0}), // padding
{1}, // dilation
1); // groups
x = batch_norm(
x, params["expand_bn_weight"], params["expand_bn_bias"],
params["expand_bn_mean"], params["expand_bn_var"],
is_training, 0.1, 1e-5, true
);
x = relu(x);
}
// Depthwise conv
auto dw_conv_weight = params["dw_conv_weight"];
x = conv2d(x, dw_conv_weight, Tensor(),
{stride}, // stride
at::IntArrayRef({1}), // padding
{1}, // dilation
expanded_channels); // groups
x = batch_norm(
x, params["dw_bn_weight"], params["dw_bn_bias"],
params["dw_bn_mean"], params["dw_bn_var"],
is_training, 0.1, 1e-5, true
);
x = relu(x);
// Squeeze and Excitation - Use shared memory and warp-level reduction
auto se = adaptive_avg_pool2d(x, {1, 1});
se = conv2d(se, params["se_reduce_weight"], Tensor(),
{1}, // stride
at::IntArrayRef({0})); // padding
se = relu(se);
se = conv2d(se, params["se_expand_weight"], Tensor(),
{1}, // stride
at::IntArrayRef({0})); // padding
se = sigmoid(se);
x = se;
// Projection phase
auto project_conv_weight = params["project_conv_weight"];
x = conv2d(x, project_conv_weight, Tensor(),
{1}, // stride
at::IntArrayRef({0}), // padding
{1}, // dilation
1); // groups
x = batch_norm(
x, params["project_bn_weight"], params["project_bn_bias"],
params["project_bn_mean"], params["project_bn_var"],
is_training, 0.1, 1e-5, true
);
return x;
}
Tensor forward(Tensor x, std::map<std::string, Tensor> params, bool is_training) {
// Initial conv
x = conv2d(x, params["conv1_weight"], Tensor(),
{2}, // stride
at::IntArrayRef({1})); // padding
x = batch_norm(
x, params["bn1_weight"], params["bn1_bias"],
params["bn1_mean"], params["bn1_var"],
is_training, 0.1, 1e-5, true
);
x = relu(x);
// MBConv blocks
std::vector<std::pair<int, int>> mbconv_configs = {{1,3}, {2,6}, {2,6}, {2,6}, {1,6}};
for (int i = 0; i < mbconv_configs.size(); i++) {
int block_num = i + 1;
auto [stride, expand_ratio] = mbconv_configs[i];
std::map<std::string, Tensor> block_params;
std::string prefix = "mbconv" + std::to_string(block_num) + "_";
for (const auto& pair : params) {
if (pair.first.rfind(prefix, 0) == 0) {
std::string key = pair.first.substr(prefix.length());
block_params[key] = pair.second;
}
}
x = mbconv_block(x, block_params, stride, expand_ratio, is_training);
}
// Final layers
x = conv2d(x, params["conv_final_weight"], Tensor(),
{1}, // stride
at::IntArrayRef({0})); // padding
x = batch_norm(
x, params["bn_final_weight"], params["bn_final_bias"],
params["bn_final_mean"], params["bn_final_var"],
is_training, 0.1, 1e-5, true
);
x = relu(x);
x = adaptive_avg_pool2d(x, {1, 1});
x = x.flatten(1);
x = linear(x, params["fc_weight"], params["fc_bias"]);
return x;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "EfficientNetB2 forward with shared memory and warp reduction");
}| Metric | Value | Unit | Variance | Samples |
|---|---|---|---|---|
| Executed Ipc Active | 0.408 | inst/cycle | 0.024 | 5 |
| Executed Ipc Elapsed | 0.100 | inst/cycle | 0.038 | 5 |
| Issue Slots Busy | 10.916 | % | 16.536 | 5 |
| Issued Ipc Active | 0.436 | inst/cycle | 0.026 | 5 |
| SM Busy | 10.916 | % | 16.536 | 5 |
| Memory Throughput | 223706499302.610 | byte/second | 178017974389986216116224.000 | 5 |
| Mem Busy | 7.794 | % | 30.570 | 5 |
| Max Bandwidth | 8.946 | % | 132.993 | 5 |
| L1/TEX Hit Rate | 31.808 | % | 286.146 | 5 |
| L2 Hit Rate | 109.210 | % | 2371.028 | 5 |
| Mem Pipes Busy | 1.392 | % | 7.177 | 5 |
| Warp Cycles Per Issued Instruction | 14.412 | cycle | 59.358 | 5 |
| Warp Cycles Per Executed Instruction | 15.448 | cycle | 64.109 | 5 |
| Avg. Active Threads Per Warp | 31.786 | 0.069 | 5 | |
| Avg. Not Predicated Off Threads Per Warp | 29.504 | 0.530 | 5 | |
| Max Active Clusters | 0.000 | cluster | 0.000 | 5 |
| Max Cluster Size | 8.000 | block | 0.000 | 5 |
| Overall GPU Occupancy | 0.000 | % | 0.000 | 5 |
| Cluster Occupancy | 0.000 | % | 0.000 | 5 |
| Block Limit SM | 32.000 | block | 0.000 | 5 |
| Block Limit Registers | 8.000 | block | 4.000 | 5 |
| Block Limit Shared Mem | 24.400 | block | 51.840 | 5 |
| Block Limit Warps | 13.600 | block | 23.040 | 5 |
| Theoretical Active Warps per SM | 41.600 | warp | 125.440 | 5 |
| Theoretical Occupancy | 65.000 | % | 306.250 | 5 |
| Achieved Occupancy | 11.882 | % | 139.112 | 5 |
| Achieved Active Warps Per SM | 7.606 | warp | 56.963 | 5 |
| Rule | Description |
|---|---|
| WRN HighPipeUtilization | All compute pipelines are under-utilized. Either this kernel is very small or it doesn't issue enough warps per scheduler. Check the Launch Statistics and Scheduler Statistics sections for further details. |
| INF CPIStall | Check the Warp Stall Sampling (All Cycles) table for the top stall locations in your source based on sampling data. The Kernel Profiling Guide (https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference) provides more details on each stall reason. |
| WRN Occupancy | This kernel's theoretical occupancy (56.2%) is limited by the number of required registers. The difference between calculated theoretical (56.2%) and measured achieved occupancy (6.1%) can be the result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on optimizing occupancy. |
| Operation / Metric | Value | Unit |
|---|---|---|
| aten::conv2d | ||
| CPU Time | 2674756.23 | μs |
| Device Time | 941277.68 | μs |
| Self CPU Time | 148672.36 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::convolution | ||
| CPU Time | 2526083.87 | μs |
| Device Time | 941277.68 | μs |
| Self CPU Time | 180329.71 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::_convolution | ||
| CPU Time | 2345754.17 | μs |
| Device Time | 941277.68 | μs |
| Self CPU Time | 228350.29 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::cudnn_convolution | ||
| CPU Time | 1854536.51 | μs |
| Device Time | 790067.25 | μs |
| Self CPU Time | 1128114.98 | μs |
| Self Device Time | 790067.25 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::batch_norm | ||
| CPU Time | 2308933.04 | μs |
| Device Time | 718245.07 | μs |
| Self CPU Time | 107103.05 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
| aten::_batch_norm_impl_index | ||
| CPU Time | 2201829.99 | μs |
| Device Time | 718245.07 | μs |
| Self CPU Time | 88184.08 | μs |
| Self Device Time | 0.00 | μs |
| CPU Memory Usage | 0 | B |
| Device Memory Usage | 0 | B |
| Self CPU Memory Usage | 0 | B |
| Self Device Memory Usage | 0 | B |
45291 warnings generated when compiling for host. Suppressed 45333 warnings (45286 in non-user code, 47 NOLINT). Use -header-filter=.* to display errors from all non-system headers. Use -system-headers to display errors from system headers as well.