def create_net(num_classes, dnn='resnet20', **kwargs):
ext = None
if dnn in ['resnet20', 'resnet56', 'resnet110']:
net = models.__dict__[dnn](num_classes=num_classes)
elif dnn == 'resnet50':
net = torchvision.models.resnet50(num_classes=num_classes)
elif dnn == 'resnet101':
net = torchvision.models.resnet101(num_classes=num_classes)
elif dnn == 'resnet152':
net = torchvision.models.resnet152(num_classes=num_classes)
elif dnn == 'densenet121':
net = torchvision.models.densenet121(num_classes=num_classes)
elif dnn == 'densenet161':
net = torchvision.models.densenet161(num_classes=num_classes)
elif dnn == 'densenet201':
net = torchvision.models.densenet201(num_classes=num_classes)
elif dnn == 'inceptionv4':
net = models.inceptionv4(num_classes=num_classes)
elif dnn == 'inceptionv3':
net = torchvision.models.inception_v3(num_classes=num_classes)
elif dnn == 'vgg16i': # vgg16 for imagenet
net = torchvision.models.vgg16(num_classes=num_classes)
elif dnn == 'googlenet':
net = models.googlenet()
elif dnn == 'mnistnet':
net = MnistNet()
elif dnn == 'fcn5net':
net = models.FCN5Net()
elif dnn == 'lenet':
net = models.LeNet()
elif dnn == 'lr':
net = models.LinearRegression()
elif dnn == 'vgg16':
net = models.VGG(dnn.upper())
elif dnn == 'alexnet':
#net = models.AlexNet()
net = torchvision.models.alexnet()
elif dnn == 'lstman4':
net, ext = models.LSTMAN4(datapath=kwargs['datapath'])
elif dnn == 'lstm':
# model = lstm(embedding_dim=args.hidden_size, num_steps=args.num_steps, batch_size=args.batch_size,
# vocab_size=vocab_size, num_layers=args.num_layers, dp_keep_prob=args.dp_keep_prob)
net = lstmpy.lstm(vocab_size=kwargs['vocab_size'],
batch_size=kwargs['batch_size'])
else:
errstr = 'Unsupport neural network %s' % dnn
logger.error(errstr)
raise errstr
return net, ext
def create_net(num_classes, dnn='resnet20', **kwargs):
ext = None
if dnn in ['resnet20', 'resnet56', 'resnet110']:
net = models.__dict__[dnn](num_classes=num_classes)
elif dnn == 'resnet50':
#net = models.__dict__['resnet50'](num_classes=num_classes)
net = torchvision.models.resnet50(num_classes=num_classes)
elif dnn == 'inceptionv4':
net = models.inceptionv4(num_classes=num_classes)
elif dnn == 'inceptionv3':
net = torchvision.models.inception_v3(num_classes=num_classes)
elif dnn == 'vgg16i': # vgg16 for imagenet
net = torchvision.models.vgg16(num_classes=num_classes)
elif dnn == 'googlenet':
net = models.googlenet()
elif dnn == 'mnistnet':
net = MnistNet()
elif dnn == 'fcn5net':
net = models.FCN5Net()
elif dnn == 'lenet':
net = models.LeNet()
elif dnn == 'lr':
net = models.LinearRegression()
elif dnn == 'vgg16':
net = models.VGG(dnn.upper())
elif dnn == 'alexnet':
net = torchvision.models.alexnet()
elif dnn == 'lstman4':
net, ext = models.LSTMAN4(datapath=kwargs['datapath'])
elif dnn == 'lstm':
net = lstmpy.lstm(vocab_size=kwargs['vocab_size'],
batch_size=kwargs['batch_size'])
else:
errstr = 'Unsupport neural network %s' % dnn
logger.error(errstr)
raise errstr
return net, ext