reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
import torch.nn.functional as F
class InceptionModule(nn.Module):
def __init__(self, in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj):
"""
:param in_channels: Number of input channels
:param out_1x1: Number of output channels for the 1x1 convolution
:param reduce_3x3: Number of output channels for the 1x1 reduction before 3x3 convolution
:param out_3x3: Number of output channels for the 3x3 convolution
:param reduce_5x5: Number of output channels for the 1x1 reduction before 5x5 convolution
:param out_5x5: Number of output channels for the 5x5 convolution
:param pool_proj: Number of output channels for the pooling projection
"""
super(InceptionModule, self).__init__()
# 1x1 convolution branch
self.branch1x1 = nn.Conv2d(in_channels, out_1x1, kernel_size=1)
# 3x3 convolution branch
self.branch3x3 = nn.Sequential(
nn.Conv2d(in_channels, reduce_3x3, kernel_size=1),
nn.Conv2d(reduce_3x3, out_3x3, kernel_size=3, padding=1)
)
# 5x5 convolution branch
self.branch5x5 = nn.Sequential(
nn.Conv2d(in_channels, reduce_5x5, kernel_size=1),
nn.Conv2d(reduce_5x5, out_5x5, kernel_size=5, padding=2)
)
# Max pooling branch
self.branch_pool = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
nn.Conv2d(in_channels, pool_proj, kernel_size=1)
)
def forward(self, x):
"""
:param x: Input tensor, shape (batch_size, in_channels, height, width)
:return: Output tensor, shape (batch_size, out_channels, height, width)
"""
branch1x1 = self.branch1x1(x)
branch3x3 = self.branch3x3(x)
branch5x5 = self.branch5x5(x)
branch_pool = self.branch_pool(x)
outputs = [branch1x1, branch3x3, branch5x5, branch_pool]
return torch.cat(outputs, 1)
class Model(nn.Module):
def __init__(self, num_classes=1000):
"""
:param num_classes: Number of output classes
"""
super(Model, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
self.maxpool1 = nn.MaxPool2d(3, stride=2, padding=1)
self.conv2 = nn.Conv2d(64, 64, kernel_size=1)
self.conv3 = nn.Conv2d(64, 192, kernel_size=3, padding=1)
self.maxpool2 = nn.MaxPool2d(3, stride=2, padding=1)
self.inception3a = InceptionModule(192, 64, 96, 128, 16, 32, 32)
self.inception3b = InceptionModule(256, 128, 128, 192, 32, 96, 64)
self.maxpool3 = nn.MaxPool2d(3, stride=2, padding=1)
self.inception4a = InceptionModule(480, 192, 96, 208, 16, 48, 64)
self.inception4b = InceptionModule(512, 160, 112, 224, 24, 64, 64)
self.inception4c = InceptionModule(512, 128, 128, 256, 24, 64, 64)
self.inception4d = InceptionModule(512, 112, 144, 288, 32, 64, 64)
self.inception4e = InceptionModule(528, 256, 160, 320, 32, 128, 128)
self.maxpool4 = nn.MaxPool2d(3, stride=2, padding=1)
self.inception5a = InceptionModule(832, 256, 160, 320, 32, 128, 128)
self.inception5b = InceptionModule(832, 384, 192, 384, 48, 128, 128)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.dropout = nn.Dropout(0.0)
self.fc = nn.Linear(1024, num_classes)
def forward(self, x):
"""
:param x: Input tensor, shape (batch_size, 3, height, width)
:return: Output tensor, shape (batch_size, num_classes)
"""
x = self.maxpool1(F.relu(self.conv1(x)))
x = F.relu(self.conv2(x))
x = self.maxpool2(F.relu(self.conv3(x)))
x = self.inception3a(x)
x = self.inception3b(x)
x = self.maxpool3(x)
x = self.inception4a(x)
x = self.inception4b(x)
x = self.inception4c(x)
x = self.inception4d(x)
x = self.inception4e(x)
x = self.maxpool4(x)
x = self.inception5a(x)
x = self.inception5b(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.dropout(x)
x = self.fc(x)
return x
# Test code
batch_size = 10
input_channels = 3
height = 224
width = 224
num_classes = 1000
def get_inputs():
return [torch.rand(batch_size, input_channels, height, width)]
def get_init_inputs():
return [num_classes]final submitted kernel (255 lines, 9,440 chars) — full source, untruncated
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor relu_maxpool2d_cuda(torch::Tensor x);
torch::Tensor maxpool2d_cuda(torch::Tensor x);
torch::Tensor spatial_avg_cuda(torch::Tensor x);
torch::Tensor relu_maxpool2d_forward(torch::Tensor x) {
return relu_maxpool2d_cuda(x);
}
torch::Tensor maxpool2d_forward(torch::Tensor x) {
return maxpool2d_cuda(x);
}
torch::Tensor spatial_avg_forward(torch::Tensor x) {
return spatial_avg_cuda(x);
}
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>
__global__ void maxpool3s2p1_kernel(const float* __restrict__ x, float* __restrict__ y,
int N, int C, int H, int W, int OH, int OW,
int64_t xs0, int64_t xs1, int64_t xs2, int64_t xs3,
int64_t ys0, int64_t ys1, int64_t ys2, int64_t ys3,
bool do_relu) {
// Iterate in NHWC logical order. This gives coalesced accesses/stores for
// channels-last tensors while still working for arbitrary strided inputs.
int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
int64_t total = (int64_t)N * OH * OW * C;
for (; idx < total; idx += (int64_t)blockDim.x * gridDim.x) {
int t = (int)idx;
int c = t % C; t /= C;
int ow = t % OW; t /= OW;
int oh = t % OH; t /= OH;
int n = t;
int hs = oh * 2 - 1;
int ws = ow * 2 - 1;
int he = hs + 3;
int we = ws + 3;
if (hs < 0) hs = 0;
if (ws < 0) ws = 0;
if (he > H) he = H;
if (we > W) we = W;
float m = -FLT_MAX;
for (int h = hs; h < he; ++h) {
for (int w = ws; w < we; ++w) {
float v = x[(int64_t)n*xs0 + (int64_t)c*xs1 + (int64_t)h*xs2 + (int64_t)w*xs3];
if (do_relu && v < 0.0f) v = 0.0f;
m = fmaxf(m, v);
}
}
y[(int64_t)n*ys0 + (int64_t)c*ys1 + (int64_t)oh*ys2 + (int64_t)ow*ys3] = m;
}
}
__global__ void spatial_avg_kernel(const float* __restrict__ x, float* __restrict__ y,
int N, int C, int H, int W,
int64_t xs0, int64_t xs1, int64_t xs2, int64_t xs3) {
int bc = blockIdx.x;
int c = bc % C;
int n = bc / C;
float sum = 0.0f;
int HW = H * W;
for (int i = threadIdx.x; i < HW; i += blockDim.x) {
int h = i / W;
int w = i - h * W;
sum += x[(int64_t)n*xs0 + (int64_t)c*xs1 + (int64_t)h*xs2 + (int64_t)w*xs3];
}
__shared__ float smem[64];
smem[threadIdx.x] = sum;
__syncthreads();
for (int s = 32; s > 0; s >>= 1) {
if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];
__syncthreads();
}
if (threadIdx.x == 0) {
y[(int64_t)n*C + c] = smem[0] / (float)HW;
}
}
static torch::Tensor pool_common(torch::Tensor x, bool relu) {
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
c10::cuda::CUDAGuard device_guard(x.device());
int N = (int)x.size(0);
int C = (int)x.size(1);
int H = (int)x.size(2);
int W = (int)x.size(3);
int OH = (H + 1) / 2;
int OW = (W + 1) / 2;
auto y = at::empty({N, C, OH, OW}, x.options().memory_format(c10::MemoryFormat::ChannelsLast));
int64_t total = (int64_t)N * OH * OW * C;
int threads = 256;
int blocks = (int)((total + threads - 1) / threads);
if (blocks > 65535) blocks = 65535;
maxpool3s2p1_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, OH, OW,
x.stride(0), x.stride(1), x.stride(2), x.stride(3),
y.stride(0), y.stride(1), y.stride(2), y.stride(3), relu);
C10_CUDA_KERNEL_LAUNCH_CHECK();
return y;
}
torch::Tensor relu_maxpool2d_cuda(torch::Tensor x) {
return pool_common(x, true);
}
torch::Tensor maxpool2d_cuda(torch::Tensor x) {
return pool_common(x, false);
}
torch::Tensor spatial_avg_cuda(torch::Tensor x) {
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
c10::cuda::CUDAGuard device_guard(x.device());
int N = (int)x.size(0);
int C = (int)x.size(1);
int H = (int)x.size(2);
int W = (int)x.size(3);
auto y = at::empty({N, C}, x.options());
spatial_avg_kernel<<<N * C, 64, 0, at::cuda::getCurrentCUDAStream()>>>(
x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W,
x.stride(0), x.stride(1), x.stride(2), x.stride(3));
C10_CUDA_KERNEL_LAUNCH_CHECK();
return y;
}
'''
_ext = load_inline(
name="googlenet_custom_pool_ops_v1",
cpp_sources=_cpp_src,
cuda_sources=_cuda_src,
functions=["relu_maxpool2d_forward", "maxpool2d_forward", "spatial_avg_forward"],
extra_cflags=["-O3"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False,
)
class InceptionModule(nn.Module):
def __init__(self, in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj):
super(InceptionModule, self).__init__()
self.branch1x1 = nn.Conv2d(in_channels, out_1x1, kernel_size=1)
self.branch3x3 = nn.Sequential(
nn.Conv2d(in_channels, reduce_3x3, kernel_size=1),
nn.Conv2d(reduce_3x3, out_3x3, kernel_size=3, padding=1),
)
self.branch5x5 = nn.Sequential(
nn.Conv2d(in_channels, reduce_5x5, kernel_size=1),
nn.Conv2d(reduce_5x5, out_5x5, kernel_size=5, padding=2),
)
self.branch_pool = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
nn.Conv2d(in_channels, pool_proj, kernel_size=1),
)
def forward(self, x):
b1 = self.branch1x1(x)
b3 = self.branch3x3(x)
b5 = self.branch5x5(x)
# The reference pool branch is maxpool stride=1, not stride=2; keep it as
# PyTorch/cudnn for exact shape and optimized behavior.
bp = self.branch_pool(x)
return torch.cat((b1, b3, b5, bp), 1)
class ModelNew(nn.Module):
def __init__(self, num_classes=1000):
super(ModelNew, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
self.maxpool1 = nn.MaxPool2d(3, stride=2, padding=1)
self.conv2 = nn.Conv2d(64, 64, kernel_size=1)
self.conv3 = nn.Conv2d(64, 192, kernel_size=3, padding=1)
self.maxpool2 = nn.MaxPool2d(3, stride=2, padding=1)
self.inception3a = InceptionModule(192, 64, 96, 128, 16, 32, 32)
self.inception3b = InceptionModule(256, 128, 128, 192, 32, 96, 64)
self.maxpool3 = nn.MaxPool2d(3, stride=2, padding=1)
self.inception4a = InceptionModule(480, 192, 96, 208, 16, 48, 64)
self.inception4b = InceptionModule(512, 160, 112, 224, 24, 64, 64)
self.inception4c = InceptionModule(512, 128, 128, 256, 24, 64, 64)
self.inception4d = InceptionModule(512, 112, 144, 288, 32, 64, 64)
self.inception4e = InceptionModule(528, 256, 160, 320, 32, 128, 128)
self.maxpool4 = nn.MaxPool2d(3, stride=2, padding=1)
self.inception5a = InceptionModule(832, 256, 160, 320, 32, 128, 128)
self.inception5b = InceptionModule(832, 384, 192, 384, 48, 128, 128)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.dropout = nn.Dropout(0.0)
self.fc = nn.Linear(1024, num_classes)
# NHWC improves cuDNN convolution throughput for this network. State-dict
# loading copies values into these parameters without changing names.
self.to(memory_format=torch.channels_last)
def forward(self, x):
if x.is_cuda and x.dtype is torch.float32:
x = x.contiguous(memory_format=torch.channels_last)
x = _ext.relu_maxpool2d_forward(self.conv1(x))
x = F.relu(self.conv2(x))
x = _ext.relu_maxpool2d_forward(self.conv3(x))
x = self.inception3a(x)
x = self.inception3b(x)
x = _ext.maxpool2d_forward(x)
x = self.inception4a(x)
x = self.inception4b(x)
x = self.inception4c(x)
x = self.inception4d(x)
x = self.inception4e(x)
x = _ext.maxpool2d_forward(x)
x = self.inception5a(x)
x = self.inception5b(x)
x = _ext.spatial_avg_forward(x)
return self.fc(x)
else:
x = self.maxpool1(F.relu(self.conv1(x)))
x = F.relu(self.conv2(x))
x = self.maxpool2(F.relu(self.conv3(x)))
x = self.inception3a(x)
x = self.inception3b(x)
x = self.maxpool3(x)
x = self.inception4a(x)
x = self.inception4b(x)
x = self.inception4c(x)
x = self.inception4d(x)
x = self.inception4e(x)
x = self.maxpool4(x)
x = self.inception5a(x)
x = self.inception5b(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.dropout(x)
x = self.fc(x)
return x