def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
groups=1,
base_width=64,
dilation=1,
norm_layer=None,
*,
reduction=16):
super(BasicBlock, self).__init__()
if (norm_layer is None):
norm_layer = nn.BatchNorm
if ((groups != 1) or (base_width != 64)):
raise ValueError(
'BasicBlock only supports groups=1 and base_width=64')
if (dilation > 1):
raise NotImplementedError(
'Dilation > 1 not supported in BasicBlock')
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = norm_layer(planes)
self.relu = nn.GELU()
self.conv2 = conv3x3(planes, planes)
self.bn2 = norm_layer(planes)
#self.se = SELayer(planes, reduction)
self.downsample = downsample
self.stride = stride
def __init__(self,
block,
layers,
num_classes=1000,
zero_init_residual=False,
groups=1,
width_per_group=64,
replace_stride_with_dilation=None,
norm_layer=None):
super(ResNet, self).__init__()
if (norm_layer is None):
norm_layer = nn.BatchNorm
self._norm_layer = norm_layer
self.inplanes = 64
self.dilation = 1
if (replace_stride_with_dilation is None):
replace_stride_with_dilation = [False, False, False]
if (len(replace_stride_with_dilation) != 3):
raise ValueError(
'replace_stride_with_dilation should be None or a 3-element tuple, got {}'
.format(replace_stride_with_dilation))
self.groups = groups
self.base_width = width_per_group
self.conv1 = nn.Conv(3,
self.inplanes,
kernel_size=7,
stride=2,
padding=3,
bias=False)
jt.init.relu_invariant_gauss_(self.conv1.weight, mode="fan_out")
self.bn1 = norm_layer(self.inplanes)
self.relu = nn.GELU()
self.maxpool = nn.Pool(kernel_size=3,
stride=2,
padding=1,
op='maximum')
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block,
128,
layers[1],
stride=2,
dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(block,
256,
layers[2],
stride=2,
dilate=replace_stride_with_dilation[1])
self.layer4 = self._make_layer(block,
512,
layers[3],
stride=2,
dilate=replace_stride_with_dilation[2])
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
# self.conv2 = conv1x1((512 * block.expansion), 1024)
# self.at = Attention(1024, num_heads=1, kdim=1024,
# vdim=1024, self_attention=True)
# self.conv3 = conv1x1(1024, (512 * block.expansion))
self.fc = nn.Linear((512 * block.expansion), num_classes)
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
groups=1,
base_width=64,
dilation=1,
norm_layer=None,
*,
reduction=16):
super(Bottleneck, self).__init__()
if (norm_layer is None):
norm_layer = nn.BatchNorm
width = (int((planes * (base_width / 64.0))) * groups)
self.conv1 = conv1x1(inplanes, width)
self.bn1 = norm_layer(width)
self.conv2 = conv3x3(width, width, stride, groups, dilation)
self.bn2 = norm_layer(width)
self.conv3 = conv1x1(width, (planes * self.expansion))
self.bn3 = norm_layer((planes * self.expansion))
self.relu = nn.GELU()
#self.se = SELayer((planes * self.expansion), reduction)
self.downsample = downsample
self.stride = stride
def test_relu(self):
# ***************************************************************
# Test ReLU Layer
# ***************************************************************
arr = np.random.randn(16, 10, 224, 224)
check_equal(arr, jnn.ReLU(), tnn.ReLU())
# ***************************************************************
# Test PReLU Layer
# ***************************************************************
arr = np.random.randn(16, 10, 224, 224)
check_equal(arr, jnn.PReLU(), tnn.PReLU())
check_equal(arr, jnn.PReLU(10, 99.9), tnn.PReLU(10, 99.9))
check_equal(arr, jnn.PReLU(10, 2), tnn.PReLU(10, 2))
check_equal(arr, jnn.PReLU(10, -0.2), tnn.PReLU(10, -0.2))
# ***************************************************************
# Test ReLU6 Layer
# ***************************************************************
arr = np.random.randn(16, 10, 224, 224)
check_equal(arr, jnn.ReLU6(), tnn.ReLU6())
# ***************************************************************
# Test LeakyReLU Layer
# ***************************************************************
arr = np.random.randn(16, 10, 224, 224)
check_equal(arr, jnn.LeakyReLU(), tnn.LeakyReLU())
check_equal(arr, jnn.LeakyReLU(2), tnn.LeakyReLU(2))
check_equal(arr, jnn.LeakyReLU(99.9), tnn.LeakyReLU(99.9))
# ***************************************************************
# Test ELU Layer
# ***************************************************************
arr = np.random.randn(16, 10, 224, 224)
check_equal(arr, jnn.ELU(), tnn.ELU())
check_equal(arr, jnn.ELU(0.3), tnn.ELU(0.3))
check_equal(arr, jnn.ELU(2), tnn.ELU(2))
check_equal(arr, jnn.ELU(99.9), tnn.ELU(99.9))
# ***************************************************************
# Test GELU Layer
# ***************************************************************
if hasattr(tnn, "GELU"):
arr = np.random.randn(16, 10, 224, 224)
check_equal(arr, jnn.GELU(), tnn.GELU())
# ***************************************************************
# Test Softplus Layer
# ***************************************************************
arr = np.random.randn(16, 10, 224, 224)
check_equal(arr, jnn.Softplus(), tnn.Softplus())
check_equal(arr, jnn.Softplus(2), tnn.Softplus(2))
check_equal(arr, jnn.Softplus(2, 99.9), tnn.Softplus(2, 99.9))
def __init__(self,
*,
image_size,
patch_size,
num_classes,
dim,
depth,
heads,
mlp_dim,
transformer=None,
channels=3,
dropout=0.,
emb_dropout=0.):
super(ViT, self).__init__()
assert image_size % patch_size == 0, 'image sizes must be divisible by the patch size'
num_patches = (image_size // patch_size)**2
patch_dim = channels * patch_size**2
assert num_patches > MIN_NUM_PATCHES, f'your num_patches is too small'
self.patch_size = patch_size
self.pos_embedding = jt.random((1, num_patches + 1, dim),
dtype='float32')
self.patch_to_embedding = nn.Linear(patch_dim, dim)
self.cls_token = jt.random((1, 1, dim), dtype='float32')
self.dropout = nn.Dropout(emb_dropout)
if transformer is None:
self.transformer = Transformer(dim, depth, heads, mlp_dim, dropout)
else:
self.transformer = transformer
self.to_cls_token = nn.Identity()
self.mlp_head = nn.Sequential(nn.LayerNorm(dim),
nn.Linear(dim, mlp_dim), nn.GELU(),
nn.Dropout(dropout),
nn.Linear(mlp_dim, num_classes))
def __init__(self, dim, hidden_dim, dropout=0.):
super(FeedForward, self).__init__()
self.net = nn.Sequential(nn.Linear(dim, hidden_dim), nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim,
dim), nn.Dropout(dropout))