INNER CODE UNIT · Python
ResidualNet
digantamisra98/Mish · PyTorch Benchmarks/resnet.py:169
def ResidualNet(network_type, depth, num_classes):
assert network_type in [
"ImageNet",
"CIFAR10",
"CIFAR100",
], "network type should be ImageNet or CIFAR10 / CIFAR100"
assert depth in [18, 34, 50, 101], "network depth should be 18, 34, 50 or 101"
if depth == 18:
model = ResNet(BasicBlock, [2, 2, 2, 2], network_type, num_classes)
if depth == 50:
model = ResNet(Bottleneck, [3, 4, 6, 3], network_type, num_classes)
return model
# %%