Exemplo n.º 1
0
def CreatNet(opt):
    name = opt.model_name
    label_num = opt.label
    if name == 'lstm':
        net = lstm.lstm(opt.input_size,
                        opt.time_step,
                        input_nc=opt.input_nc,
                        num_classes=label_num)
    elif name == 'cnn_1d':
        net = cnn_1d.cnn(opt.input_nc, num_classes=label_num)
    elif name == 'resnet18_1d':
        net = resnet_1d.resnet18()
        net.conv1 = nn.Conv1d(opt.input_nc, 64, 7, 2, 3, bias=False)
        net.fc = nn.Linear(512, label_num)
    elif name == 'resnet34_1d':
        net = resnet_1d.resnet34()
        net.conv1 = nn.Conv1d(opt.input_nc, 64, 7, 2, 3, bias=False)
        net.fc = nn.Linear(512, label_num)
    elif name == 'multi_scale_resnet_1d':
        net = multi_scale_resnet_1d.Multi_Scale_ResNet(inchannel=opt.input_nc,
                                                       num_classes=label_num)
    elif name == 'micro_multi_scale_resnet_1d':
        net = micro_multi_scale_resnet_1d.Multi_Scale_ResNet(
            inchannel=opt.input_nc, num_classes=label_num)
    elif name == 'multi_scale_resnet':
        net = multi_scale_resnet.Multi_Scale_ResNet(inchannel=opt.input_nc,
                                                    num_classes=label_num)
    elif name == 'dfcnn':
        net = dfcnn.dfcnn(num_classes=label_num)
    elif name in ['resnet101', 'resnet50', 'resnet18']:
        if name == 'resnet101':
            net = resnet.resnet101(pretrained=False)
            net.fc = nn.Linear(2048, label_num)
        elif name == 'resnet50':
            net = resnet.resnet50(pretrained=False)
            net.fc = nn.Linear(2048, label_num)
        elif name == 'resnet18':
            net = resnet.resnet18(pretrained=False)
            net.fc = nn.Linear(512, label_num)
        net.conv1 = nn.Conv2d(opt.input_nc, 64, 7, 2, 3, bias=False)

    elif 'densenet' in name:
        if name == 'densenet121':
            net = densenet.densenet121(pretrained=False, num_classes=label_num)
        elif name == 'densenet201':
            net = densenet.densenet201(pretrained=False, num_classes=label_num)
    elif name == 'squeezenet':
        net = squeezenet.squeezenet1_1(pretrained=False,
                                       num_classes=label_num,
                                       inchannel=1)

    return net
Exemplo n.º 2
0
def CreatNet(name):
    if name == 'lstm':
        net = lstm.lstm(100, 27, num_classes=5)
    elif name == 'cnn_1d':
        net = cnn_1d.cnn(1, num_classes=5)
    elif name == 'resnet18_1d':
        net = resnet_1d.resnet18()
        net.conv1 = nn.Conv1d(1, 64, 7, 2, 3, bias=False)
        net.fc = nn.Linear(512, 5)
    elif name == 'multi_scale_resnet_1d':
        net = multi_scale_resnet_1d.Multi_Scale_ResNet(inchannel=1,
                                                       num_classes=5)
    elif name == 'multi_scale_resnet':
        net = multi_scale_resnet.Multi_Scale_ResNet(inchannel=1, num_classes=5)
    elif name == 'dfcnn':
        net = dfcnn.dfcnn(num_classes=5)
    elif name in ['resnet101', 'resnet50', 'resnet18']:
        if name == 'resnet101':
            net = resnet.resnet101(pretrained=False)
            net.fc = nn.Linear(2048, 5)
        elif name == 'resnet50':
            net = resnet.resnet50(pretrained=False)
            net.fc = nn.Linear(2048, 5)
        elif name == 'resnet18':
            net = resnet.resnet18(pretrained=False)
            net.fc = nn.Linear(512, 5)
        net.conv1 = nn.Conv2d(1, 64, 7, 2, 3, bias=False)

    elif 'densenet' in name:
        if name == 'densenet121':
            net = densenet.densenet121(pretrained=False, num_classes=5)
        elif name == 'densenet201':
            net = densenet.densenet201(pretrained=False, num_classes=5)
    elif name == 'squeezenet':
        net = squeezenet.squeezenet1_1(pretrained=False,
                                       num_classes=5,
                                       inchannel=1)

    return net