import torch
import torch.nn as nn
# --- ResNet 基本ブロック定義 ---
class ResBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
super().__init__()
self.expansion = 1
# 第一層
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3,
stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=True)
# 第二層
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channels)
# ショートカット経路
self.shortcut = nn.Sequential()
if stride != 1 or in_channels != self.expansion * out_channels:
self.shortcut = nn.Sequential(
nn.Conv2d(in_channels, self.expansion * out_channels,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(self.expansion * out_channels)
)
def forward(self, x):
out = self.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
out = self.relu(out)
return out
torch.manual_seed(42)
# 入力テンソル（バッチサイズ=2, チャンネル数=64, サイズ=32x32）
x = torch.randn(2, 64, 32, 32)
# ブロック作成（同じチャンネル数）
block1 = ResBlock(64, 64)
y1 = block1(x)
print("出力 shape（stride=1）:", y1.shape)
# ブロック作成（チャンネル変化あり・stride=2）
block2 = ResBlock(64, 128, stride=2)
y2 = block2(x)
print("出力 shape（stride=2）:", y2.shape)
