def define_model(is_resnet, is_densenet, is_senet):
use18 = True # True
if is_resnet:
if not use18:
original_model = resnet.resnet18(pretrained = True)
Encoder = modules.E_resnet(original_model)
model = net.model(Encoder, num_features=512, block_channel = [64, 128, 256, 512])
else:
stereoModel = Resnet18Encoder(3)
model_dict = stereoModel.state_dict()
encoder_dict = torch.load('./models/monodepth_resnet18_001.pth',map_location='cpu' )
new_dict = {}
for key in encoder_dict:
if key in model_dict:
new_dict[key] = encoder_dict[key]
stereoModel.load_state_dict(new_dict )
Encoder = stereoModel
model = net.model(Encoder, num_features=512, block_channel = [64, 128, 256, 512])
if is_densenet:
original_model = densenet.densenet161(pretrained=True)
Encoder = modules.E_densenet(original_model)
model = net.model(Encoder, num_features=2208, block_channel = [192, 384, 1056, 2208])
if is_senet:
original_model = senet.senet154(pretrained='imagenet')
Encoder = modules.E_senet(original_model)
model = net.model(Encoder, num_features=2048, block_channel = [256, 512, 1024, 2048])
return model
def define_model(encoder='resnet'):
if encoder is 'resnet':
original_model = resnet.resnet50(pretrained=True)
Encoder = modules.E_resnet(original_model)
model = net.model(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048])
if encoder is 'densenet':
original_model = densenet.densenet161(pretrained=True)
Encoder = modules.E_densenet(original_model)
model = net.model(Encoder,
num_features=2208,
block_channel=[192, 384, 1056, 2208])
if encoder is 'senet':
original_model = senet.senet154(pretrained='imagenet')
Encoder = modules.E_senet(original_model)
model = net.model(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048])
if encoder is 'resnet4':
original_model = resnet4.resnet50(pretrained=True)
Encoder = modules.E_resnet(original_model)
model = net.model(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048])
return model
def get_model_dic(device):
models = {}
#densenet_121 = densenet121(num_classes=110)
#load_model(densenet_121,"./pre_weights/ep_38_densenet121_val_acc_0.6527.pth",device)
densenet_161 = densenet161(num_classes=110)
load_model(densenet_161,
"./pre_weights/ep_30_densenet161_val_acc_0.6990.pth", device)
resnet_50 = resnet50(num_classes=110)
load_model(resnet_50, "./pre_weights/ep_41_resnet50_val_acc_0.6900.pth",
device)
incept_v3 = inception_v3(num_classes=110)
load_model(incept_v3,
"./pre_weights/ep_36_inception_v3_val_acc_0.6668.pth", device)
#incept_v1 = googlenet(num_classes=110)
#load_model(incept_v1,"./pre_weights/ep_33_googlenet_val_acc_0.7091.pth",device)
#vgg16 = vgg16_bn(num_classes=110)
#load_model(vgg16, "./pre_weights/ep_30_vgg16_bn_val_acc_0.7282.pth",device)
incept_resnet_v2_adv = InceptionResNetV2(num_classes=110)
load_model(incept_resnet_v2_adv,
"./pre_weights/ep_22_InceptionResNetV2_val_acc_0.8214.pth",
device)
incept_v4_adv = InceptionV4(num_classes=110)
load_model(incept_v4_adv,
"./pre_weights/ep_37_InceptionV4_val_acc_0.7119.pth", device)
MainModel = imp.load_source('MainModel',
"./models_old/tf_to_pytorch_resnet_v1_50.py")
resnet_model = torch.load(
'./models_old/tf_to_pytorch_resnet_v1_50.pth').to(device)
MainModel = imp.load_source('MainModel',
"./models_old/tf_to_pytorch_vgg16.py")
vgg_model = torch.load('./models_old/tf_to_pytorch_vgg16.pth').to(device)
MainModel = imp.load_source('MainModel',
"./models_old/tf_to_pytorch_inception_v1.py")
inception_model = torch.load(
'./models_old/tf_to_pytorch_inception_v1.pth').to(device)
models = { #"densenet121":densenet_121,
"densenet161": densenet_161,
"resnet_50": resnet_50,
# "incept_v1":incept_v1,
"incept_v3": incept_v3,
"incept_resnet_v2_adv": incept_resnet_v2_adv,
"incept_v4_adv": incept_v4_adv,
#"vgg16":vgg16
"old_incept": inception_model,
"old_res": resnet_model,
"old_vgg": vgg_model
}
return models
def define_test_model():
#archs = {"Resnet", "Densenet", "SEnet", "Custom"}
is_resnet = args.arch == "Resnet" #True #False #True
is_densenet = args.arch == "Densenet" # #False #True #False # False
is_senet = args.arch == "SEnet" # True #False #True #False
is_custom = args.arch == "Custom"
if is_resnet:
#original_model = resnet.resnet18(pretrained = pretrain_logical)
#Encoder = modules.E_resnet(original_model)
#model = net.model(Encoder, num_features=2048, block_channel = [256, 512, 1024, 2048])
stereoModel = Resnet18Encoder(3)
model_dict = stereoModel.state_dict()
encoder_dict = torch.load('./models/monodepth_resnet18_001.pth',
map_location='cpu')
new_dict = {}
for key in encoder_dict:
# print(key)
if key in model_dict:
new_dict[key] = encoder_dict[key]
stereoModel.load_state_dict(new_dict)
Encoder = stereoModel
model = net.model(Encoder,
num_features=512,
block_channel=[64, 128, 256, 512])
print("Loading a model...")
print("/model_epoch_{}.pth".format(str(args.load_epoch)))
model = model.cuda().float()
#print(stereoModel)
#print(model)
model_dict = torch.load(
args.load_dir +
"/original_model_epoch_{}.pth".format(str(args.load_epoch)))
new_dict = model_dict
#new_dict = {}
#for key in model_dict:
# new_dict[key[7:]] = model_dict[key]
model.load_state_dict(new_dict)
if is_densenet:
# TODO: no dot bug
original_model = densenet.densenet161(pretrained=True)
Encoder = modules.E_densenet(original_model)
model = net.model(Encoder,
num_features=2208,
block_channel=[192, 384, 1056, 2208])
if is_senet:
original_model = senet.senet154(pretrained='imagenet')
Encoder = modules.E_senet(original_model)
model = net.model(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048])
return model
def define_model(is_resnet, is_densenet, is_senet):
if is_resnet:
original_model = resnet.resnet50(pretrained=True)
Encoder = modules.E_resnet(original_model)
model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048])
if is_densenet:
original_model = densenet.densenet161(pretrained=True)
Encoder = modules.E_densenet(original_model)
model = net.model(Encoder, num_features=2208, block_channel=[192, 384, 1056, 2208])
if is_senet:
original_model = senet.senet154(pretrained=None)
Encoder = modules.E_senet(original_model)
model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048])
return model
def _get_model(self, backbone):
if backbone == 'resnet18':
model = resnet18(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet34':
model = resnet34(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet50':
model = resnet50(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet101':
model = resnet101(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet152':
model = resnet152(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'preact_resnet18':
model = PreActResNet18(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet34':
model = PreActResNet34(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet50':
model = PreActResNet50(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet101':
model = PreActResNet101(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet152':
model = PreActResNet152(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'densenet121':
model = densenet121(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'densenet161':
model = densenet161(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'densenet169':
model = densenet169(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'densenet201':
model = densenet201(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'mlp':
model = MLPNet().to(self.args.device)
elif backbone == 'cnn_small' or backbone == "CNN_SMALL":
model = CNN_small(self.args.classnum).to(self.args.device)
elif backbone == "cnn" or backbone == "CNN":
model = CNN(n_outputs=self.args.classnum, input_channel=self.args.input_dim, linear_num=self.args.linear_num).to(self.args.device)
else:
print("No matched backbone. Using ResNet50...")
model = resnet50(pretrained=True, num_classes=self.args.classnum,
input_size=self.args.image_size).to(self.args.device)
return model
def define_model():
if args.net_arch == 'resnext_64x4d':
clz_model = resnext.resnext(groups=64, width_per_group=4)
fea_model = modules.FeatureResnext(clz_model.features)
model = modules.DABC(fea_model, args.num_classes)
elif args.net_arch == 'resnet':
pass
elif args.net_arch == 'densenet':
clz_model = densenet.densenet161(pretrained=True, )
fea_model = modules.FeatureResnext(clz_model.features)
model = modules.DABC(fea_model, args.num_classes)
elif args.net_arch == 'senet':
pass
else:
raise NotImplementedError(
"Network Architecture [%s] is not recognized." % args.net_arch)
return model
def create_model(model):
print('==> 2. Building the Model...')
if model == 'lenet':
net = lenet(dataset.num_classes, dataset.img_channels).to(device)
elif model == 'convnet':
net = convnet(dataset.num_classes, dataset.img_channels).to(device)
elif model == 'resnet32_wide':
net = resnet32_wide(dataset.num_classes, dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'densenet40':
net = densenet40(num_classes = dataset.num_classes, in_channels = dataset.img_channels).to(device)
elif model == 'densenet161':
net = densenet161(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet50':
net = resnet50(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet110':
net = resnet110(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet152':
net = resnet152(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet110_SD':
net = resnet110_SD(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet152_SD':
net = resnet152_SD(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
return net
__models_small__ = {
'ResNet18':
lambda: E_resnet(resnet18(pretrained=True)),
'ResNet34':
lambda: E_resnet(resnet34(pretrained=True)),
'ResNet50':
lambda: E_resnet(resnet50(pretrained=True)),
'ResNet101':
lambda: E_resnet(resnet101(pretrained=True)),
'ResNet152':
lambda: E_resnet(resnet152(pretrained=True)),
'DenseNet121':
lambda: E_densenet(densenet121(pretrained=True)),
'DenseNet161':
lambda: E_densenet(densenet161(pretrained=True)),
'DenseNet169':
lambda: E_densenet(densenet169(pretrained=True)),
'DenseNet201':
lambda: E_densenet(densenet201(pretrained=True)),
'SENet154':
lambda: E_senet(senet154(pretrained="imagenet")),
'SE_ResNet50':
lambda: E_senet(se_resnet50(pretrained="imagenet")),
'SE_ResNet101':
lambda: E_senet(se_resnet101(pretrained="imagenet")),
'SE_ResNet152':
lambda: E_senet(se_resnet152(pretrained="imagenet")),
'SE_ResNext50_32x4d':
lambda: E_senet(se_resnext50_32x4d(pretrained="imagenet")),
'SE_ResNext101_32x4d':
def get_network(args):
""" return given network
"""
if args.net == 'vgg16':
from models.vgg import vgg16_bn
net = vgg16_bn()
elif args.net == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn()
elif args.net == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn()
elif args.net == 'vgg19':
from models.vgg import vgg19_bn
net = vgg19_bn()
elif args.net == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif args.net == 'densenet161':
from models.densenet import densenet161
net = densenet161()
elif args.net == 'densenet169':
from models.densenet import densenet169
net = densenet169()
elif args.net == 'densenet201':
from models.densenet import densenet201
net = densenet201()
elif args.net == 'googlenet':
from models.googlenet import googlenet
net = googlenet()
elif args.net == 'inceptionv3':
from models.inceptionv3 import inceptionv3
net = inceptionv3()
elif args.net == 'inceptionv4':
from models.inceptionv4 import inceptionv4
net = inceptionv4()
elif args.net == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
net = inception_resnet_v2()
elif args.net == 'xception':
from models.xception import xception
net = xception()
elif args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34()
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50()
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101()
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152()
elif args.net == 'preactresnet18':
from models.preactresnet import preactresnet18
net = preactresnet18()
elif args.net == 'preactresnet34':
from models.preactresnet import preactresnet34
net = preactresnet34()
elif args.net == 'preactresnet50':
from models.preactresnet import preactresnet50
net = preactresnet50()
elif args.net == 'preactresnet101':
from models.preactresnet import preactresnet101
net = preactresnet101()
elif args.net == 'preactresnet152':
from models.preactresnet import preactresnet152
net = preactresnet152()
elif args.net == 'resnext50':
from models.resnext import resnext50
net = resnext50()
elif args.net == 'resnext101':
from models.resnext import resnext101
net = resnext101()
elif args.net == 'resnext152':
from models.resnext import resnext152
net = resnext152()
elif args.net == 'shufflenet':
from models.shufflenet import shufflenet
net = shufflenet()
elif args.net == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2()
elif args.net == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet()
elif args.net == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet()
elif args.net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2()
elif args.net == 'nasnet':
from models.nasnet import nasnet
net = nasnet()
elif args.net == 'attention56':
from models.attention import attention56
net = attention56()
elif args.net == 'attention92':
from models.attention import attention92
net = attention92()
elif args.net == 'seresnet18':
from models.senet import seresnet18
net = seresnet18()
elif args.net == 'seresnet34':
from models.senet import seresnet34
net = seresnet34()
elif args.net == 'seresnet50':
from models.senet import seresnet50
net = seresnet50()
elif args.net == 'seresnet101':
from models.senet import seresnet101
net = seresnet101()
elif args.net == 'seresnet152':
from models.senet import seresnet152
net = seresnet152()
elif args.net == 'wideresnet':
from models.wideresidual import wideresnet
net = wideresnet()
elif args.net == 'stochasticdepth18':
from models.stochasticdepth import stochastic_depth_resnet18
net = stochastic_depth_resnet18()
elif args.net == 'stochasticdepth34':
from models.stochasticdepth import stochastic_depth_resnet34
net = stochastic_depth_resnet34()
elif args.net == 'stochasticdepth50':
from models.stochasticdepth import stochastic_depth_resnet50
net = stochastic_depth_resnet50()
elif args.net == 'stochasticdepth101':
from models.stochasticdepth import stochastic_depth_resnet101
net = stochastic_depth_resnet101()
elif args.net == 'normal_resnet':
from models.normal_resnet import resnet18
net = resnet18()
elif args.net == 'hyper_resnet':
from models.hypernet_main import Hypernet_Main
net = Hypernet_Main(
encoder="resnet18",
hypernet_params={'vqvae_dict_size': args.dict_size})
elif args.net == 'normal_resnet_wo_bn':
from models.normal_resnet_wo_bn import resnet18
net = resnet18()
elif args.net == 'hyper_resnet_wo_bn':
from models.hypernet_main import Hypernet_Main
net = Hypernet_Main(
encoder="resnet18_wobn",
hypernet_params={'vqvae_dict_size': args.dict_size})
else:
print('the network name you have entered is not supported yet')
sys.exit()
if args.gpu: #use_gpu
net = net.cuda()
return net
def get_network(args, use_gpu=True):
""" return given network
"""
if args.net == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet(args)
elif args.net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2(args)
elif args.net == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn(args)
elif args.net == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn(args)
elif args.net == 'vgg19':
# from models.vgg import vgg19_bn
# net = vgg19_bn(args)
from torchvision.models import vgg19_bn
import torch.nn as nn
net = vgg19_bn(pretrained=True)
net.classifier[6] = nn.Linear(4096, args.nc)
elif args.net == 'densenet121':
from models.densenet import densenet121
net = densenet121(args)
elif args.net == 'densenet161':
from models.densenet import densenet161
net = densenet161(args)
elif args.net == 'densenet169':
from models.densenet import densenet169
net = densenet169(args)
elif args.net == 'densenet201':
from models.densenet import densenet201
net = densenet201(args)
elif args.net == 'googlenet':
from models.googlenet import googlenet
net = googlenet(args)
elif args.net == 'inceptionv3':
from models.inceptionv3 import inceptionv3
net = inceptionv3(args)
elif args.net == 'inceptionv4':
from models.inceptionv4 import inceptionv4
net = inceptionv4(args)
elif args.net == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
net = inception_resnet_v2(args)
elif args.net == 'xception':
from models.xception import xception
net = xception(args)
elif args.net == 'resnet18':
# from models.resnet import resnet18
# net = resnet18(args)
from torchvision.models import resnet18
import torch.nn as nn
net = resnet18(pretrained=True)
net.fc = nn.Linear(512, args.nc)
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34(args)
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50(args)
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101(args)
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152(args)
elif args.net == 'preactresnet18':
from models.preactresnet import preactresnet18
net = preactresnet18(args)
elif args.net == 'preactresnet34':
from models.preactresnet import preactresnet34
net = preactresnet34(args)
elif args.net == 'preactresnet50':
from models.preactresnet import preactresnet50
net = preactresnet50(args)
elif args.net == 'preactresnet101':
from models.preactresnet import preactresnet101
net = preactresnet101(args)
elif args.net == 'preactresnet152':
from models.preactresnet import preactresnet152
net = preactresnet152(args)
elif args.net == 'resnext50':
from models.resnext import resnext50
net = resnext50(args)
elif args.net == 'resnext101':
from models.resnext import resnext101
net = resnext101(args)
elif args.net == 'resnext152':
from models.resnext import resnext152
net = resnext152(args)
elif args.net == 'shufflenet':
from models.shufflenet import shufflenet
net = shufflenet(args)
elif args.net == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2(args)
elif args.net == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet(args)
elif args.net == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet(args)
elif args.net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2(args)
elif args.net == 'mobilenetv3':
from models.mobilenetv3 import mobileNetv3
net = mobileNetv3(args)
elif args.net == 'mobilenetv3_l':
from models.mobilenetv3 import mobileNetv3
net = mobileNetv3(args, mode='large')
elif args.net == 'mobilenetv3_s':
from models.mobilenetv3 import mobileNetv3
net = mobileNetv3(args, mode='small')
elif args.net == 'nasnet':
from models.nasnet import nasnetalarge
net = nasnetalarge(args)
elif args.net == 'attention56':
from models.attention import attention56
net = attention56(args)
elif args.net == 'attention92':
from models.attention import attention92
net = attention92(args)
elif args.net == 'seresnet18':
from models.senet import seresnet18
net = seresnet18(args)
elif args.net == 'seresnet34':
from models.senet import seresnet34
net = seresnet34(args)
elif args.net == 'seresnet50':
from models.senet import seresnet50
net = seresnet50(args)
elif args.net == 'seresnet101':
from models.senet import seresnet101
net = seresnet101(args)
elif args.net == 'seresnet152':
from models.senet import seresnet152
net = seresnet152(args)
elif args.net.lower() == 'sqnxt_23_1x':
from models.SqueezeNext import SqNxt_23_1x
net = SqNxt_23_1x(args)
elif args.net.lower() == 'sqnxt_23_1xv5':
from models.SqueezeNext import SqNxt_23_1x_v5
net = SqNxt_23_1x_v5(args)
elif args.net.lower() == 'sqnxt_23_2x':
from models.SqueezeNext import SqNxt_23_2x
net = SqNxt_23_2x(args)
elif args.net.lower() == 'sqnxt_23_2xv5':
from models.SqueezeNext import SqNxt_23_2x_v5
net = SqNxt_23_2x_v5(args)
elif args.net.lower() == 'mnasnet':
# from models.MnasNet import mnasnet
# net = mnasnet(args)
from models.nasnet_mobile import nasnet_Mobile
net = nasnet_Mobile(args)
elif args.net == 'efficientnet_b0':
from models.efficientnet import efficientnet_b0
net = efficientnet_b0(args)
elif args.net == 'efficientnet_b1':
from models.efficientnet import efficientnet_b1
net = efficientnet_b1(args)
elif args.net == 'efficientnet_b2':
from models.efficientnet import efficientnet_b2
net = efficientnet_b2(args)
elif args.net == 'efficientnet_b3':
from models.efficientnet import efficientnet_b3
net = efficientnet_b3(args)
elif args.net == 'efficientnet_b4':
from models.efficientnet import efficientnet_b4
net = efficientnet_b4(args)
elif args.net == 'efficientnet_b5':
from models.efficientnet import efficientnet_b5
net = efficientnet_b5(args)
elif args.net == 'efficientnet_b6':
from models.efficientnet import efficientnet_b6
net = efficientnet_b6(args)
elif args.net == 'efficientnet_b7':
from models.efficientnet import efficientnet_b7
net = efficientnet_b7(args)
elif args.net == 'mlp':
from models.mlp import MLPClassifier
net = MLPClassifier(args)
elif args.net == 'alexnet':
from torchvision.models import alexnet
import torch.nn as nn
net = alexnet(pretrained=True)
net.classifier[6] = nn.Linear(4096, args.nc)
elif args.net == 'lambda18':
from models._lambda import LambdaResnet18
net = LambdaResnet18(num_classes=args.nc, channels=args.cs)
elif args.net == 'lambda34':
from models._lambda import LambdaResnet34
net = LambdaResnet34(num_classes=args.nc, channels=args.cs)
elif args.net == 'lambda50':
from models._lambda import LambdaResnet50
net = LambdaResnet50(num_classes=args.nc, channels=args.cs)
elif args.net == 'lambda101':
from models._lambda import LambdaResnet101
net = LambdaResnet101(num_classes=args.nc)
elif args.net == 'lambda152':
from models._lambda import LambdaResnet152
net = LambdaResnet152(num_classes=args.nc, channels=args.cs)
else:
print('the network name you have entered is not supported yet')
sys.exit()
if use_gpu:
net = net.cuda()
return net
def define_model(is_resnet,
is_densenet,
is_senet,
model='tbdp',
parallel=False,
semff=False,
pcamff=False):
if is_resnet:
original_model = resnet.resnet50(pretrained=True)
Encoder = modules.E_resnet(original_model)
if model == 'tbdp':
model = net.TBDPNet(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048],
parallel=parallel,
pcamff=pcamff)
elif model == 'hu':
model = net.Hu(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048],
semff=semff,
pcamff=pcamff)
else:
raise NotImplementedError(
"Select model type in [\'tbdp\', \'hu\']")
if is_densenet:
original_model = densenet.densenet161(pretrained=True)
Encoder = modules.E_densenet(original_model)
if model == 'tbdp':
model = net.TBDPNet(Encoder,
num_features=2208,
block_channel=[192, 384, 1056, 2208],
parallel=parallel,
pcamff=pcamff)
elif model == 'hu':
model = net.Hu(Encoder,
num_features=2208,
block_channel=[192, 384, 1056, 2208],
semff=semff,
pcamff=pcamff)
else:
raise NotImplementedError(
"Select model type in [\'tbdp\', \'hu\']")
if is_senet:
original_model = senet.senet154(pretrained='imagenet')
Encoder = modules.E_senet(original_model)
if model == 'tbdp':
model = net.TBDPNet(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048],
parallel=parallel,
pcamff=pcamff)
elif model == 'hu':
model = net.Hu(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048],
semff=semff,
pcamff=pcamff)
else:
raise NotImplementedError(
"Select model type in [\'tbdp\', \'hu\']")
return model
def main_worker(gpu, args):
global best_acc1
global best_auc
global minimum_loss
global count
global best_accdr
args.gpu = gpu
if not os.path.exists(args.model_dir):
os.makedirs(args.model_dir)
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.arch == "vgg11":
from models.vgg import vgg11
model = vgg11(num_classes=args.num_class, crossCBAM=args.crossCBAM)
elif args.arch == "resnet50":
from models.resnet50 import resnet50
model = resnet50(num_classes=args.num_class, multitask=args.multitask, liu=args.liu,
chen=args.chen, CAN_TS=args.CAN_TS, crossCBAM=args.crossCBAM,
crosspatialCBAM = args.crosspatialCBAM, choice=args.choice)
elif args.arch == "resnet34":
from models.resnet50 import resnet34
model = resnet34(num_classes=args.num_class, multitask=args.multitask, liu=args.liu,
chen=args.chen,CAN_TS=args.CAN_TS, crossCBAM=args.crossCBAM,
crosspatialCBAM = args.crosspatialCBAM)
elif args.arch == "resnet18":
from models.resnet50 import resnet18
model = resnet18(num_classes=args.num_class, multitask=args.multitask, liu=args.liu,
chen=args.chen, flagCBAM=False, crossCBAM=args.crossCBAM)
elif args.arch == "densenet161":
from models.densenet import densenet161
model = densenet161(num_classes=args.num_class, multitask=args.multitask, cosface=False, liu=args.liu,
chen=args.chen, crossCBAM=args.crossCBAM)
elif args.arch == "wired":
from models.wirednetwork import CNN
model = CNN(args, num_classes=args.num_class)
else:
print ("no backbone model")
if args.pretrained:
print ("==> Load pretrained model")
model_dict = model.state_dict()
pretrain_path = {"resnet50": "pretrain/resnet50-19c8e357.pth",
"resnet34": "pretrain/resnet34-333f7ec4.pth",
"resnet18": "pretrain/resnet18-5c106cde.pth",
"densenet161": "pretrain/densenet161-8d451a50.pth",
"vgg11": "pretrain/vgg11-bbd30ac9.pth",
"densenet121": "pretrain/densenet121-a639ec97.pth"}[args.arch]
pretrained_dict = torch.load(pretrain_path)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
pretrained_dict.pop('classifier.weight', None)
pretrained_dict.pop('classifier.bias', None)
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0))
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
if args.adam:
optimizer = torch.optim.Adam(model.parameters(), args.base_lr, weight_decay=args.weight_decay)
else:
optimizer = torch.optim.SGD(model.parameters(), args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location={'cuda:4':'cuda:0'})
# args.start_epoch = checkpoint['epoch']
# load partial weights
if not args.evaluate:
print ("load partial weights")
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in checkpoint['state_dict'].items() if k in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
else:
print("load whole weights")
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
exit(0)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
size = 224
tra = transforms.Compose([
transforms.Resize(256),
transforms.RandomResizedCrop(size),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
# transforms.RandomRotation(90),
# transforms.ColorJitter(0.05, 0.05, 0.05, 0.05),
transforms.ToTensor(),
normalize,
])
tra_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize])
# tra = transforms.Compose([
# transforms.Resize(350),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# # transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
# transforms.RandomRotation([-180, 180]),
# transforms.RandomAffine([-180, 180], translate=[0.1, 0.1], scale=[0.7, 1.3]),
# transforms.RandomCrop(224),
# # transforms.CenterCrop(224),
# transforms.ToTensor(),
# normalize
# ])
# print (args.model_dir)
# tra = transforms.Compose([
# transforms.Resize(350),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# # transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
# transforms.RandomRotation([-180, 180]),
# transforms.RandomAffine([-180, 180], translate=[0.1, 0.1], scale=[0.7, 1.3]),
# transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
# transforms.ToTensor(),
# normalize
# ])
# tra_test = transforms.Compose([
# transforms.Resize(350),
# transforms.CenterCrop(224),
# transforms.ToTensor(),
# normalize])
if args.dataset == 'amd':
from datasets.amd_dataset import traindataset
elif args.dataset == 'pm':
from datasets.pm_dataset import traindataset
elif args.dataset == "drdme":
from datasets.drdme_dataset import traindataset
elif args.dataset == "missidor":
from datasets.missidor import traindataset
elif args.dataset == "kaggle":
from datasets.kaggle import traindataset
else:
print ("no dataset")
exit(0)
val_dataset = traindataset(root=args.data, mode = 'val',
transform=tra_test, num_class=args.num_class,
multitask=args.multitask, args=args)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
if args.evaluate:
a = time.time()
# savedir = args.resume.replace("model_converge.pth.tar","")
savedir = args.resume.replace(args.resume.split("/")[-1], "")
# savedir = "./"
if not args.multitask:
acc, auc, precision_dr, recall_dr, f1score_dr = validate(val_loader, model, args)
result_list = [acc, auc, precision_dr, recall_dr, f1score_dr]
print ("acc, auc, precision, recall, f1", acc, auc, precision_dr, recall_dr, f1score_dr)
save_result_txt(savedir, result_list)
print("time", time.time() - a)
return
else:
acc_dr, acc_dme, acc_joint, other_results, se, sp = validate(val_loader, model, args)
print ("acc_dr, acc_dme, acc_joint", acc_dr, acc_dme, acc_joint)
exit(0)
print ("auc_dr, auc_dme, precision_dr, precision_dme, recall_dr, recall_dme, f1score_dr, f1score_dme",
other_results)
print ("se, sp", se, sp)
result_list = [acc_dr, acc_dme, acc_joint]
result_list += other_results
result_list += [se, sp]
save_result_txt(savedir, result_list)
print ("time", time.time()-a)
return
train_dataset = traindataset(root=args.data, mode='train', transform=tra, num_class=args.num_class,
multitask=args.multitask, args=args)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True,worker_init_fn=worker_init_fn)
writer = SummaryWriter()
writer.add_text('Text', str(args))
#
from lr_scheduler import LRScheduler
lr_scheduler = LRScheduler(optimizer, len(train_loader), args)
for epoch in range(args.start_epoch, args.epochs):
is_best = False
is_best_auc = False
is_best_acc = False
# lr = adjust_learning_rate(optimizer, epoch, args)
# writer.add_scalar("lr", lr, epoch)
# train for one epoch
loss_train = train(train_loader, model, criterion, lr_scheduler, writer, epoch, optimizer, args)
writer.add_scalar('Train loss', loss_train, epoch)
# evaluate on validation set
if epoch % 5 == 0:
if args.dataset == "kaggle":
acc_dr, auc_dr = validate(val_loader, model, args)
writer.add_scalar("Val acc_dr", acc_dr, epoch)
writer.add_scalar("Val auc_dr", auc_dr, epoch)
is_best = acc_dr >= best_acc1
best_acc1 = max(acc_dr, best_acc1)
elif not args.multitask:
acc, auc, precision, recall, f1 = validate(val_loader, model, args)
writer.add_scalar("Val acc_dr", acc, epoch)
writer.add_scalar("Val auc_dr", auc, epoch)
is_best = auc >= best_acc1
best_acc1 = max(auc, best_acc1)
else:
acc_dr, acc_dme, joint_acc, other_results, se, sp , losses = validate(val_loader, model, args,criterion)
writer.add_scalar("Val acc_dr", acc_dr, epoch)
writer.add_scalar("Val acc_dme", acc_dme, epoch)
writer.add_scalar("Val acc_joint", joint_acc, epoch)
writer.add_scalar("Val auc_dr", other_results[0], epoch)
writer.add_scalar("Val auc_dme", other_results[1], epoch)
writer.add_scalar("val loss", losses, epoch)
is_best = joint_acc >= best_acc1
best_acc1 = max(joint_acc, best_acc1)
is_best_auc = other_results[0] >= best_auc
best_auc = max(other_results[0], best_auc)
is_best_acc = acc_dr >= best_accdr
best_accdr = max(acc_dr, best_accdr)
if not args.invalid:
if is_best:
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer' : optimizer.state_dict(),
}, is_best, filename = "model_converge.pth.tar", save_dir=args.model_dir)
if is_best_auc:
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_auc,
'optimizer' : optimizer.state_dict(),
}, False, filename = "converge_auc.pth.tar", save_dir=args.model_dir)
if is_best_acc:
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_accdr,
'optimizer': optimizer.state_dict(),
}, False, filename="converge_acc.pth.tar", save_dir=args.model_dir)
logger.info("early stop")
break
for ds in dataset:
data_path = os.path.join(args.root, ds)
cls = [
x for x in os.listdir(data_path)
if os.path.isdir(os.path.join(data_path, x))
]
num_class = len(cls)
models = {
"vgg16": vgg.vgg16_bn(num_class),
"vgg19": vgg.vgg19_bn(num_class),
"densenet121": densenet.densenet121(num_class),
"densenet161": densenet.densenet161(num_class),
"resnet34": resnet.resnet34(num_class),
"resnet50": resnet.resnet50(num_class),
"resnet101": resnet.resnet101(num_class),
"seresnet34": senet.seresnet34(num_class),
"seresnet50": senet.seresnet50(num_class),
"seresnet101": senet.seresnet101(num_class),
"resnext34": resnext.resnext34(num_class),
"resnext50": resnext.resnext50(num_class),
"resnext101": resnext.resnext101(num_class),
"shufflenet": shufflenet.shufflenet(num_class),
"xception": xception.xception(num_class)
}
for net_name in models.keys():
writer = SummaryWriter('./runs/%s_%s/' % (ds, net_name))
model = models[net_name]
def get_network(args):
""" return given network
"""
if args.model == 'vgg16':
from models.vgg import vgg16_bn
model = vgg16_bn()
elif args.model == 'vgg13':
from models.vgg import vgg13_bn
model = vgg13_bn()
elif args.model == 'vgg11':
from models.vgg import vgg11_bn
model = vgg11_bn()
elif args.model == 'vgg19':
from models.vgg import vgg19_bn
model = vgg19_bn()
elif args.model == 'densenet121':
from models.densenet import densenet121
model = densenet121()
elif args.model == 'densenet161':
from models.densenet import densenet161
model = densenet161()
elif args.model == 'densenet169':
from models.densenet import densenet169
model = densenet169()
elif args.model == 'densenet201':
from models.densenet import densenet201
model = densenet201()
elif args.model == 'googlenet':
from models.googlenet import googlenet
model = googlenet()
elif args.model == 'inceptionv3':
from models.inceptionv3 import inceptionv3
model = inceptionv3()
elif args.model == 'inceptionv4':
from models.inceptionv4 import inceptionv4
model = inceptionv4()
elif args.model == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
model = inception_resnet_v2()
elif args.model == 'xception':
from models.xception import xception
model = xception()
elif args.model == 'resnet18':
from models.resnet import resnet18
model = resnet18()
elif args.model == 'resnet34':
from models.resnet import resnet34
model = resnet34()
elif args.model == 'resnet50':
from models.resnet import resnet50
model = resnet50()
elif args.model == 'resnet101':
from models.resnet import resnet101
model = resnet101()
elif args.model == 'resnet152':
from models.resnet import resnet152
model = resnet152()
elif args.model == 'preactresnet18':
from models.preactresnet import preactresnet18
model = preactresnet18()
elif args.model == 'preactresnet34':
from models.preactresnet import preactresnet34
model = preactresnet34()
elif args.model == 'preactresnet50':
from models.preactresnet import preactresnet50
model = preactresnet50()
elif args.model == 'preactresnet101':
from models.preactresnet import preactresnet101
model = preactresnet101()
elif args.model == 'preactresnet152':
from models.preactresnet import preactresnet152
model = preactresnet152()
elif args.model == 'resnext50':
from models.resnext import resnext50
model = resnext50()
elif args.model == 'resnext101':
from models.resnext import resnext101
model = resnext101()
elif args.model == 'resnext152':
from models.resnext import resnext152
model = resnext152()
elif args.model == 'shufflenet':
from models.shufflenet import shufflenet
model = shufflenet()
elif args.model == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
model = shufflenetv2()
elif args.model == 'squeezenet':
from models.squeezenet import squeezenet
model = squeezenet()
elif args.model == 'mobilenet':
from models.mobilenet import mobilenet
model = mobilenet()
elif args.model == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
model = mobilenetv2()
elif args.model == 'nasnet':
from models.nasnet import nasnet
model = nasnet()
elif args.model == 'attention56':
from models.attention import attention56
model = attention56()
elif args.model == 'attention92':
from models.attention import attention92
model = attention92()
elif args.model == 'seresnet18':
from models.senet import seresnet18
model = seresnet18()
elif args.model == 'seresnet34':
from models.senet import seresnet34
model = seresnet34()
elif args.model == 'seresnet50':
from models.senet import seresnet50
model = seresnet50()
elif args.model == 'seresnet101':
from models.senet import seresnet101
model = seresnet101()
elif args.model == 'seresnet152':
from models.senet import seresnet152
model = seresnet152()
elif args.model == 'wideresnet':
from models.wideresidual import wideresnet
model = wideresnet()
elif args.model == 'stochasticdepth18':
from models.stochasticdepth import stochastic_depth_resnet18
model = stochastic_depth_resnet18()
elif args.model == 'stochasticdepth34':
from models.stochasticdepth import stochastic_depth_resnet34
model = stochastic_depth_resnet34()
elif args.model == 'stochasticdepth50':
from models.stochasticdepth import stochastic_depth_resnet50
model = stochastic_depth_resnet50()
elif args.model == 'stochasticdepth101':
from models.stochasticdepth import stochastic_depth_resnet101
model = stochastic_depth_resnet101()
else:
print('the network name you have entered is not supported yet')
sys.exit()
return model
log = './log/log20.txt'
with open(log, 'w') as f:
f.write('\n')
nlabels = 80
display = 20
maxepoch = 40
base_lr = 0.001
gamma = 1 # each epoch learning_rate = learning_rate * gamma
# state_dict = torch.load('./snapshots/dense_plus_sssd_epoch_39.pytorch')
state_dict = torch.load('./snapshots/densenet_truncated_dot5_step2_10.pytorch')
# state_dict = torch.load('./snapshots/ckpt_epoch_48.pytorch')
# state_dict = torch.load('./snapshots/dense_sssd_cnn0_sgd_epoch_2fc_step3_4.pytorch')
net = densenet.densenet161(state_dict)
# net = densessd.densenet161(state_dict)
net.cuda()
mean = [x / 255 for x in [125.3, 123.0, 113.9]]
std = [x / 255 for x in [63.0, 62.1, 66.7]]
crop_size = 224
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomSizedCrop(crop_size),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
test_transform = transforms.Compose([
transforms.CenterCrop(crop_size),
transforms.ToTensor(),
def get_network(args):
""" return given network
"""
if args.net == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn()
elif args.net == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn()
elif args.net == 'vgg16':
from models.vgg import vgg16_bn
net = vgg16_bn()
elif args.net == 'vgg19':
from models.vgg import vgg19_bn
net = vgg19_bn()
elif args.net == 'googlenet':
from models.googLeNet import GoogLeNet
net = GoogLeNet()
elif args.net == 'inceptionv3':
from models.inceptionv3 import Inceptionv3
net = Inceptionv3()
elif args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34()
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50()
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101()
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152()
elif args.net == 'wrn':
from models.wideresnet import wideresnet
net = wideresnet()
elif args.net == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif args.net == 'densenet161':
from models.densenet import densenet161
net = densenet161()
elif args.net == 'densenet169':
from models.densenet import densenet169
net = densenet169()
elif args.net == 'densenet201':
from models.densenet import densenet201
net = densenet201()
else:
print('the network name you have entered is not supported yet')
sys.exit()
if args.gpu:
print("use gpu")
net = net.cuda()
return net
def get_network(args):
""" return given network
"""
if args.net == 'vgg16':
from models.vgg import vgg16_bn
net = vgg16_bn()
elif args.net == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn()
elif args.net == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn()
elif args.net == 'vgg19':
from models.vgg import vgg19_bn
net = vgg19_bn()
# elif args.net == 'efficientnet':
# from models.effnetv2 import effnetv2_s
# net = effnetv2_s()
elif args.net == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif args.net == 'densenet161':
from models.densenet import densenet161
net = densenet161()
elif args.net == 'densenet169':
from models.densenet import densenet169
net = densenet169()
elif args.net == 'densenet201':
from models.densenet import densenet201
net = densenet201()
elif args.net == 'googlenet':
from models.googlenet import googlenet
net = googlenet()
elif args.net == 'inceptionv3':
from models.inceptionv3 import inceptionv3
net = inceptionv3()
elif args.net == 'inceptionv4':
from models.inceptionv4 import inceptionv4
net = inceptionv4()
elif args.net == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
net = inception_resnet_v2()
elif args.net == 'xception':
from models.xception import xception
net = xception()
elif args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34()
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50()
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101()
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152()
elif args.net == 'preactresnet18':
from models.preactresnet import preactresnet18
net = preactresnet18()
elif args.net == 'preactresnet34':
from models.preactresnet import preactresnet34
net = preactresnet34()
elif args.net == 'preactresnet50':
from models.preactresnet import preactresnet50
net = preactresnet50()
elif args.net == 'preactresnet101':
from models.preactresnet import preactresnet101
net = preactresnet101()
elif args.net == 'preactresnet152':
from models.preactresnet import preactresnet152
net = preactresnet152()
elif args.net == 'resnext50':
from models.resnext import resnext50
net = resnext50()
elif args.net == 'resnext101':
from models.resnext import resnext101
net = resnext101()
elif args.net == 'resnext152':
from models.resnext import resnext152
net = resnext152()
elif args.net == 'shufflenet':
from models.shufflenet import shufflenet
net = shufflenet()
elif args.net == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2()
elif args.net == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet()
elif args.net == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet()
elif args.net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2()
elif args.net == 'nasnet':
from models.nasnet import nasnet
net = nasnet()
elif args.net == 'attention56':
from models.attention import attention56
net = attention56()
elif args.net == 'attention92':
from models.attention import attention92
net = attention92()
elif args.net == 'seresnet18':
from models.senet import seresnet18
net = seresnet18()
elif args.net == 'seresnet34':
from models.senet import seresnet34
net = seresnet34()
elif args.net == 'seresnet50':
from models.senet import seresnet50
net = seresnet50()
elif args.net == 'seresnet101':
from models.senet import seresnet101
net = seresnet101()
elif args.net == 'seresnet152':
from models.senet import seresnet152
net = seresnet152()
elif args.net == 'wideresnet':
from models.wideresidual import wideresnet
net = wideresnet()
elif args.net == 'stochasticdepth18':
from models.stochasticdepth import stochastic_depth_resnet18
net = stochastic_depth_resnet18()
elif args.net == 'stochasticdepth34':
from models.stochasticdepth import stochastic_depth_resnet34
net = stochastic_depth_resnet34()
elif args.net == 'stochasticdepth50':
from models.stochasticdepth import stochastic_depth_resnet50
net = stochastic_depth_resnet50()
elif args.net == 'stochasticdepth101':
from models.stochasticdepth import stochastic_depth_resnet101
net = stochastic_depth_resnet101()
elif args.net == 'efficientnetb0':
from models.efficientnet import efficientnetb0
net = efficientnetb0()
elif args.net == 'efficientnetb1':
from models.efficientnet import efficientnetb1
net = efficientnetb1()
elif args.net == 'efficientnetb2':
from models.efficientnet import efficientnetb2
net = efficientnetb2()
elif args.net == 'efficientnetb3':
from models.efficientnet import efficientnetb3
net = efficientnetb3()
elif args.net == 'efficientnetb4':
from models.efficientnet import efficientnetb4
net = efficientnetb4()
elif args.net == 'efficientnetb5':
from models.efficientnet import efficientnetb5
net = efficientnetb5()
elif args.net == 'efficientnetb6':
from models.efficientnet import efficientnetb6
net = efficientnetb6()
elif args.net == 'efficientnetb7':
from models.efficientnet import efficientnetb7
net = efficientnetb7()
elif args.net == 'efficientnetl2':
from models.efficientnet import efficientnetl2
net = efficientnetl2()
elif args.net == 'eff':
from models.efficientnet_pytorch import EfficientNet
net = EfficientNet.from_pretrained('efficientnet-b7', num_classes=2)
else:
print('the network name you have entered is not supported yet')
sys.exit()
if args.gpu: #use_gpu
net = net.cuda()
print("use-gpu")
return net
def get_model_dics(device, model_list= None):
if model_list is None:
model_list = ['densenet121', 'densenet161', 'resnet50', 'resnet152',
'incept_v1', 'incept_v3', 'incept_v4_adv', 'incept_resnet_v2_adv',
'black_densenet161', 'black_resnet50', 'black_incept_v3',
'old_vgg','old_res','old_incept']
models = {}
for model in model_list:
if model=='densenet121':
models['densenet121'] = densenet121(num_classes=110)
load_model(models['densenet121'],"./pre_weights/ep_38_densenet121_val_acc_0.6527.pth",device)
if model=='densenet161':
models['densenet161'] = densenet161(num_classes=110)
load_model(models['densenet161'],"./pre_weights/ep_30_densenet161_val_acc_0.6990.pth",device)
if model=='resnet50':
models['resnet50'] = resnet50(num_classes=110)
load_model(models['resnet50'],"./pre_weights/ep_41_resnet50_val_acc_0.6900.pth",device)
if model=='incept_v3':
models['incept_v3'] = inception_v3(num_classes=110)
load_model(models['incept_v3'],"./pre_weights/ep_36_inception_v3_val_acc_0.6668.pth",device)
if model=='incept_v1':
models['incept_v1'] = googlenet(num_classes=110)
load_model(models['incept_v1'],"./pre_weights/ep_33_googlenet_val_acc_0.7091.pth",device)
#vgg16 = vgg16_bn(num_classes=110)
#load_model(vgg16, "./pre_weights/ep_30_vgg16_bn_val_acc_0.7282.pth",device)
if model=='incept_resnet_v2':
models['incept_resnet_v2'] = InceptionResNetV2(num_classes=110)
load_model(models['incept_resnet_v2'], "./pre_weights/ep_17_InceptionResNetV2_ori_0.8320.pth",device)
if model=='incept_v4':
models['incept_v4'] = InceptionV4(num_classes=110)
load_model(models['incept_v4'],"./pre_weights/ep_17_InceptionV4_ori_0.8171.pth",device)
if model=='incept_resnet_v2_adv':
models['incept_resnet_v2_adv'] = InceptionResNetV2(num_classes=110)
load_model(models['incept_resnet_v2_adv'], "./pre_weights/ep_22_InceptionResNetV2_val_acc_0.8214.pth",device)
if model=='incept_v4_adv':
models['incept_v4_adv'] = InceptionV4(num_classes=110)
load_model(models['incept_v4_adv'],"./pre_weights/ep_24_InceptionV4_val_acc_0.6765.pth",device)
if model=='incept_resnet_v2_adv2':
models['incept_resnet_v2_adv2'] = InceptionResNetV2(num_classes=110)
load_model(models['incept_resnet_v2_adv2'],"./pre_weights/ep_29_InceptionResNetV2_adv2_0.8115.pth",device)
#load_model(models['incept_resnet_v2_adv2'],"../pre_weights/ep_13_InceptionResNetV2_val_acc_0.8889.pth",device)
if model=='incept_v4_adv2':
models['incept_v4_adv2'] = InceptionV4(num_classes=110)
load_model(models['incept_v4_adv2'],"./pre_weights/ep_32_InceptionV4_adv2_0.7579.pth",device)
if model=='resnet152':
models['resnet152'] = resnet152(num_classes=110)
load_model(models['resnet152'],"./pre_weights/ep_14_resnet152_ori_0.6956.pth",device)
if model=='resnet152_adv':
models['resnet152_adv'] = resnet152(num_classes=110)
load_model(models['resnet152_adv'],"./pre_weights/ep_29_resnet152_adv_0.6939.pth",device)
if model=='black_resnet50':
models['black_resnet50'] = resnet50(num_classes=110)
load_model(models['black_resnet50'],"./test_weights/ep_0_resnet50_val_acc_0.7063.pth",device)
if model=='black_densenet161':
models['black_densenet161'] = densenet161(num_classes=110)
load_model(models['black_densenet161'],"./test_weights/ep_4_densenet161_val_acc_0.6892.pth",device)
if model=='black_incept_v3':
models['black_incept_v3']=inception_v3(num_classes=110)
load_model(models['black_incept_v3'],"./test_weights/ep_28_inception_v3_val_acc_0.6680.pth",device)
if model=='old_res':
MainModel = imp.load_source('MainModel', "./models_old/tf_to_pytorch_resnet_v1_50.py")
models['old_res'] = torch.load('./models_old/tf_to_pytorch_resnet_v1_50.pth').to(device)
if model=='old_vgg':
MainModel = imp.load_source('MainModel', "./models_old/tf_to_pytorch_vgg16.py")
models[model] = torch.load('./models_old/tf_to_pytorch_vgg16.pth').to(device)
if model=='old_incept':
MainModel = imp.load_source('MainModel', "./models_old/tf_to_pytorch_inception_v1.py")
models[model] = torch.load('./models_old/tf_to_pytorch_inception_v1.pth').to(device)
return models
def main_worker(gpu, args):
global best_acc1
global minimum_loss
global count
args.gpu = gpu
if not os.path.exists(args.model_dir):
os.makedirs(args.model_dir)
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.arch == "vgg11":
from models.vgg import vgg11
model = vgg11(num_classes=args.num_class, crossCBAM=args.crossCBAM)
elif args.arch == "resnet50":
from models.resnet50 import resnet50
model = resnet50(num_classes=args.num_class,
multitask=args.multitask,
liu=args.liu,
chen=args.chen,
flagCBAM=False,
crossCBAM=args.crossCBAM)
elif args.arch == "resnet34":
from models.resnet50 import resnet34
model = resnet34(num_classes=args.num_class,
multitask=args.multitask,
liu=args.liu,
chen=args.chen,
flagCBAM=False,
crossCBAM=args.crossCBAM)
elif args.arch == "resnet18":
from models.resnet50 import resnet18
model = resnet18(num_classes=args.num_class,
multitask=args.multitask,
liu=args.liu,
chen=args.chen,
flagCBAM=False,
crossCBAM=args.crossCBAM)
elif args.arch == "densenet161":
from models.densenet import densenet161
model = densenet161(num_classes=args.num_class,
multitask=args.multitask,
cosface=False,
liu=args.liu,
chen=args.chen)
elif args.arch == "wired":
from models.wirednetwork import CNN
model = CNN(args, num_classes=args.num_class)
else:
print("no backbone model")
if args.pretrained:
print("==> Load pretrained model")
model_dict = model.state_dict()
pretrain_path = {
"resnet50": "pretrain/resnet50-19c8e357.pth",
"resnet34": "pretrain/resnet34-333f7ec4.pth",
"resnet18": "pretrain/resnet18-5c106cde.pth",
"densenet161": "pretrain/densenet161-8d451a50.pth",
"vgg11": "pretrain/vgg11-bbd30ac9.pth"
}[args.arch]
pretrained_dict = torch.load(pretrain_path)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
pretrained_dict.pop('classifier.weight', None)
pretrained_dict.pop('classifier.bias', None)
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
weight_decay=args.weight_decay)
# optimizer = torch.optim.SGD(model.parameters(), args.base_lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume,
map_location={'cuda:4': 'cuda:0'})
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
exit(0)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
size = 224
# tra = transforms.Compose([
# # transforms.Resize(256),
# transforms.RandomResizedCrop(size),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# # transforms.RandomRotation(90),
# # transforms.ColorJitter(0.05, 0.05, 0.05, 0.05),
# transforms.ToTensor(),
# normalize,
# ])
# tra_test = transforms.Compose([
# transforms.Resize(size+32),
# transforms.CenterCrop(size),
# transforms.ToTensor(),
# normalize])
tra = transforms.Compose([
transforms.Resize(350),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
# transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
transforms.RandomRotation([-180, 180]),
transforms.RandomAffine([-180, 180],
translate=[0.1, 0.1],
scale=[0.7, 1.3]),
transforms.RandomCrop(224),
# transforms.CenterCrop(224),
transforms.ToTensor(),
normalize
])
# tra = transforms.Compose([
# transforms.Resize(350),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
# transforms.ToTensor(),
# normalize])
#
tra_test = transforms.Compose([
transforms.Resize(350),
transforms.CenterCrop(224),
transforms.ToTensor(), normalize
])
#
# from autoaugment import ImageNetPolicy
# tra =transforms.Compose([
# transforms.RandomResizedCrop(224),
# transforms.RandomHorizontalFlip(),
# ImageNetPolicy(),
# transforms.ToTensor(),
# normalize])
# image = PIL.Image.open(path)
# policy = ImageNetPolicy()
# transformed = policy(image)
if args.dataset == 'amd':
from datasets.amd_dataset import traindataset
elif args.dataset == 'pm':
from datasets.pm_dataset import traindataset
elif args.dataset == "drdme":
from datasets.drdme_dataset import traindataset
elif args.dataset == "missidor":
from datasets.missidor import traindataset
else:
print("no dataset")
exit(0)
if args.evaluate:
# result = validate(val_loader, model, args)
result = multi_validate(model, test_times, normalize, traindataset,
args)
print("acc_dr, acc_dme, acc_joint", result)
return
val_dataset = traindataset(root=args.data,
mode='val',
transform=tra_test,
num_class=args.num_class,
multitask=args.multitask)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
train_dataset = traindataset(root=args.data,
mode='train',
transform=tra,
num_class=args.num_class,
multitask=args.multitask)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
worker_init_fn=worker_init_fn)
writer = SummaryWriter()
writer.add_text('Text', str(args))
# from lr_scheduler import LRScheduler
# lr_scheduler = LRScheduler(optimizer, len(train_loader), args)
for epoch in range(args.start_epoch, args.epochs):
is_best = False
lr = adjust_learning_rate(optimizer, epoch, args)
writer.add_scalar("lr", lr, epoch)
# train for one epoch
loss_train = train(train_loader, model, criterion, optimizer, args)
writer.add_scalar('Train loss', loss_train, epoch)
# evaluate on validation set
if epoch % 20 == 0:
acc_dr, acc_dme, joint_acc = validate(val_loader, model, args)
writer.add_scalar("Val acc_dr", acc_dr, epoch)
writer.add_scalar("Val acc_dme", acc_dme, epoch)
writer.add_scalar("Val acc_joint", joint_acc, epoch)
is_best = joint_acc >= best_acc1
best_acc1 = max(joint_acc, best_acc1)
if not args.invalid:
if is_best:
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
},
is_best,
filename="checkpoint" + str(epoch) + ".pth.tar",
save_dir=args.model_dir)
def get_network(args):
""" return given network
"""
if args.task == 'cifar10':
nclass = 10
elif args.task == 'cifar100':
nclass = 100
#Yang added none bn vggs
if args.net == 'vgg16':
from models.vgg import vgg16
net = vgg16(num_classes=nclass)
elif args.net == 'vgg13':
from models.vgg import vgg13
net = vgg13(num_classes=nclass)
elif args.net == 'vgg11':
from models.vgg import vgg11
net = vgg11(num_classes=nclass)
elif args.net == 'vgg19':
from models.vgg import vgg19
net = vgg19(num_classes=nclass)
elif args.net == 'vgg16bn':
from models.vgg import vgg16_bn
net = vgg16_bn(num_classes=nclass)
elif args.net == 'vgg13bn':
from models.vgg import vgg13_bn
net = vgg13_bn(num_classes=nclass)
elif args.net == 'vgg11bn':
from models.vgg import vgg11_bn
net = vgg11_bn(num_classes=nclass)
elif args.net == 'vgg19bn':
from models.vgg import vgg19_bn
net = vgg19_bn(num_classes=nclass)
elif args.net == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif args.net == 'densenet161':
from models.densenet import densenet161
net = densenet161()
elif args.net == 'densenet169':
from models.densenet import densenet169
net = densenet169()
elif args.net == 'densenet201':
from models.densenet import densenet201
net = densenet201()
elif args.net == 'googlenet':
from models.googlenet import googlenet
net = googlenet(num_classes=nclass)
elif args.net == 'inceptionv3':
from models.inceptionv3 import inceptionv3
net = inceptionv3()
elif args.net == 'inceptionv4':
from models.inceptionv4 import inceptionv4
net = inceptionv4()
elif args.net == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
net = inception_resnet_v2()
elif args.net == 'xception':
from models.xception import xception
net = xception(num_classes=nclass)
elif args.net == 'scnet':
from models.sphereconvnet import sphereconvnet
net = sphereconvnet(num_classes=nclass)
elif args.net == 'sphereresnet18':
from models.sphereconvnet import resnet18
net = resnet18(num_classes=nclass)
elif args.net == 'sphereresnet32':
from models.sphereconvnet import sphereresnet32
net = sphereresnet32(num_classes=nclass)
elif args.net == 'plainresnet32':
from models.sphereconvnet import plainresnet32
net = plainresnet32(num_classes=nclass)
elif args.net == 'ynet18':
from models.ynet import resnet18
net = resnet18(num_classes=nclass)
elif args.net == 'ynet34':
from models.ynet import resnet34
net = resnet34(num_classes=nclass)
elif args.net == 'ynet50':
from models.ynet import resnet50
net = resnet50(num_classes=nclass)
elif args.net == 'ynet101':
from models.ynet import resnet101
net = resnet101(num_classes=nclass)
elif args.net == 'ynet152':
from models.ynet import resnet152
net = resnet152(num_classes=nclass)
elif args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18(num_classes=nclass)
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34(num_classes=nclass)
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50(num_classes=nclass)
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101(num_classes=nclass)
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152(num_classes=nclass)
elif args.net == 'preactresnet18':
from models.preactresnet import preactresnet18
net = preactresnet18(num_classes=nclass)
elif args.net == 'preactresnet34':
from models.preactresnet import preactresnet34
net = preactresnet34(num_classes=nclass)
elif args.net == 'preactresnet50':
from models.preactresnet import preactresnet50
net = preactresnet50(num_classes=nclass)
elif args.net == 'preactresnet101':
from models.preactresnet import preactresnet101
net = preactresnet101(num_classes=nclass)
elif args.net == 'preactresnet152':
from models.preactresnet import preactresnet152
net = preactresnet152(num_classes=nclass)
elif args.net == 'resnext50':
from models.resnext import resnext50
net = resnext50(num_classes=nclass)
elif args.net == 'resnext101':
from models.resnext import resnext101
net = resnext101(num_classes=nclass)
elif args.net == 'resnext152':
from models.resnext import resnext152
net = resnext152(num_classes=nclass)
elif args.net == 'shufflenet':
from models.shufflenet import shufflenet
net = shufflenet()
elif args.net == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2()
elif args.net == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet()
elif args.net == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet(num_classes=nclass)
elif args.net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2(num_classes=nclass)
elif args.net == 'nasnet':
from models.nasnet import nasnet
net = nasnet(num_classes=nclass)
elif args.net == 'attention56':
from models.attention import attention56
net = attention56()
elif args.net == 'attention92':
from models.attention import attention92
net = attention92()
elif args.net == 'seresnet18':
from models.senet import seresnet18
net = seresnet18(num_classes=nclass)
elif args.net == 'seresnet34':
from models.senet import seresnet34
net = seresnet34(num_classes=nclass)
elif args.net == 'seresnet50':
from models.senet import seresnet50
net = seresnet50(num_classes=nclass)
elif args.net == 'seresnet101':
from models.senet import seresnet101
net = seresnet101(num_classes=nclass)
elif args.net == 'seresnet152':
from models.senet import seresnet152
net = seresnet152(num_classes=nclass)
else:
print('the network name you have entered is not supported yet')
sys.exit()
if args.gpu: #use_gpu
net = net.cuda()
return net
def get_network(args, use_gpu=True):
""" return given network
"""
if args.net == 'vgg16':
from models.vgg import vgg16_bn
net = vgg16_bn()
elif args.net == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn()
elif args.net == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn()
elif args.net == 'vgg19':
from models.vgg import vgg19_bn
net = vgg19_bn()
elif args.net == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif args.net == 'densenet161':
from models.densenet import densenet161
net = densenet161()
elif args.net == 'densenet169':
from models.densenet import densenet169
net = densenet169()
elif args.net == 'densenet201':
from models.densenet import densenet201
net = densenet201()
elif args.net == 'googlenet':
from models.googlenet import googlenet
net = googlenet()
elif args.net == 'inceptionv3':
from models.inceptionv3 import inceptionv3
net = inceptionv3()
elif args.net == 'inceptionv4':
from models.inceptionv4 import inceptionv4
net = inceptionv4()
elif args.net == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
net = inception_resnet_v2()
elif args.net == 'xception':
from models.xception import xception
net = xception()
elif args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34()
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50()
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101()
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152()
elif args.net == 'preactresnet18':
from models.preactresnet import preactresnet18
net = preactresnet18()
elif args.net == 'preactresnet34':
from models.preactresnet import preactresnet34
net = preactresnet34()
elif args.net == 'preactresnet50':
from models.preactresnet import preactresnet50
net = preactresnet50()
elif args.net == 'preactresnet101':
from models.preactresnet import preactresnet101
net = preactresnet101()
elif args.net == 'preactresnet152':
from models.preactresnet import preactresnet152
net = preactresnet152()
elif args.net == 'resnext50':
from models.resnext import resnext50
net = resnext50()
elif args.net == 'resnext101':
from models.resnext import resnext101
net = resnext101()
elif args.net == 'resnext152':
from models.resnext import resnext152
net = resnext152()
elif args.net == 'shufflenet':
from models.shufflenet import shufflenet
net = shufflenet()
elif args.net == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2()
elif args.net == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet()
elif args.net == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet()
elif args.net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2()
elif args.net == 'nasnet':
from models.nasnet import nasnet
net = nasnet()
elif args.net == 'attention56':
from models.attention import attention56
net = attention56()
elif args.net == 'attention92':
from models.attention import attention92
net = attention92()
elif args.net == 'seresnet18':
from models.senet import seresnet18
net = seresnet18()
elif args.net == 'seresnet34':
from models.senet import seresnet34
net = seresnet34()
elif args.net == 'seresnet50':
from models.senet import seresnet50
net = seresnet50()
elif args.net == 'seresnet101':
from models.senet import seresnet101
net = seresnet101()
elif args.net == 'seresnet152':
from models.senet import seresnet152
net = seresnet152()
else:
print('the network name you have entered is not supported yet')
sys.exit()
if use_gpu:
net = net.cuda()
return net
def get_model(args):
if args.datasets == 'ImageNet':
return models_imagenet.__dict__[args.arch]()
if args.datasets == 'CIFAR10' or args.datasets == 'MNIST':
num_class = 10
elif args.datasets == 'CIFAR100':
num_class = 100
if args.datasets == 'CIFAR100':
if args.arch == 'vgg16':
from models.vgg import vgg16_bn
net = vgg16_bn()
elif args.arch == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn()
elif args.arch == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn()
elif args.arch == 'vgg19':
from models.vgg import vgg19_bn
net = vgg19_bn()
elif args.arch == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif args.arch == 'densenet161':
from models.densenet import densenet161
net = densenet161()
elif args.arch == 'densenet169':
from models.densenet import densenet169
net = densenet169()
elif args.arch == 'densenet201':
from models.densenet import densenet201
net = densenet201()
elif args.arch == 'googlenet':
from models.googlenet import googlenet
net = googlenet()
elif args.arch == 'inceptionv3':
from models.inceptionv3 import inceptionv3
net = inceptionv3()
elif args.arch == 'inceptionv4':
from models.inceptionv4 import inceptionv4
net = inceptionv4()
elif args.arch == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
net = inception_resnet_v2()
elif args.arch == 'xception':
from models.xception import xception
net = xception()
elif args.arch == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif args.arch == 'resnet34':
from models.resnet import resnet34
net = resnet34()
elif args.arch == 'resnet50':
from models.resnet import resnet50
net = resnet50()
elif args.arch == 'resnet101':
from models.resnet import resnet101
net = resnet101()
elif args.arch == 'resnet152':
from models.resnet import resnet152
net = resnet152()
elif args.arch == 'preactresnet18':
from models.preactresnet import preactresnet18
net = preactresnet18()
elif args.arch == 'preactresnet34':
from models.preactresnet import preactresnet34
net = preactresnet34()
elif args.arch == 'preactresnet50':
from models.preactresnet import preactresnet50
net = preactresnet50()
elif args.arch == 'preactresnet101':
from models.preactresnet import preactresnet101
net = preactresnet101()
elif args.arch == 'preactresnet152':
from models.preactresnet import preactresnet152
net = preactresnet152()
elif args.arch == 'resnext50':
from models.resnext import resnext50
net = resnext50()
elif args.arch == 'resnext101':
from models.resnext import resnext101
net = resnext101()
elif args.arch == 'resnext152':
from models.resnext import resnext152
net = resnext152()
elif args.arch == 'shufflenet':
from models.shufflenet import shufflenet
net = shufflenet()
elif args.arch == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2()
elif args.arch == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet()
elif args.arch == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet()
elif args.arch == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2()
elif args.arch == 'nasnet':
from models.nasnet import nasnet
net = nasnet()
elif args.arch == 'attention56':
from models.attention import attention56
net = attention56()
elif args.arch == 'attention92':
from models.attention import attention92
net = attention92()
elif args.arch == 'seresnet18':
from models.senet import seresnet18
net = seresnet18()
elif args.arch == 'seresnet34':
from models.senet import seresnet34
net = seresnet34()
elif args.arch == 'seresnet50':
from models.senet import seresnet50
net = seresnet50()
elif args.arch == 'seresnet101':
from models.senet import seresnet101
net = seresnet101()
elif args.arch == 'seresnet152':
from models.senet import seresnet152
net = seresnet152()
elif args.arch == 'wideresnet':
from models.wideresidual import wideresnet
net = wideresnet()
elif args.arch == 'stochasticdepth18':
from models.stochasticdepth import stochastic_depth_resnet18
net = stochastic_depth_resnet18()
elif args.arch == 'efficientnet':
from models.efficientnet import efficientnet
net = efficientnet(1, 1, 100, bn_momentum=0.9)
elif args.arch == 'stochasticdepth34':
from models.stochasticdepth import stochastic_depth_resnet34
net = stochastic_depth_resnet34()
elif args.arch == 'stochasticdepth50':
from models.stochasticdepth import stochastic_depth_resnet50
net = stochastic_depth_resnet50()
elif args.arch == 'stochasticdepth101':
from models.stochasticdepth import stochastic_depth_resnet101
net = stochastic_depth_resnet101()
else:
net = resnet.__dict__[args.arch](num_classes=num_class)
return net
return resnet.__dict__[args.arch](num_classes=num_class)
def get_model(class_num):
if (MODEL_TYPE == 'alexnet'):
model = alexnet.alexnet(pretrained=FINETUNE)
elif (MODEL_TYPE == 'vgg'):
if (MODEL_DEPTH_OR_VERSION == 11):
model = vgg.vgg11(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 13):
model = vgg.vgg13(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 16):
model = vgg.vgg16(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 19):
model = vgg.vgg19(pretrained=FINETUNE)
else:
print('Error : VGG should have depth of either [11, 13, 16, 19]')
sys.exit(1)
elif (MODEL_TYPE == 'squeezenet'):
if (MODEL_DEPTH_OR_VERSION == 0 or MODEL_DEPTH_OR_VERSION == 'v0'):
model = squeezenet.squeezenet1_0(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 1 or MODEL_DEPTH_OR_VERSION == 'v1'):
model = squeezenet.squeezenet1_1(pretrained=FINETUNE)
else:
print('Error : Squeezenet should have version of either [0, 1]')
sys.exit(1)
elif (MODEL_TYPE == 'resnet'):
if (MODEL_DEPTH_OR_VERSION == 18):
model = resnet.resnet18(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 34):
model = resnet.resnet34(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 50):
model = resnet.resnet50(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 101):
model = resnet.resnet101(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 152):
model = resnet.resnet152(pretrained=FINETUNE)
else:
print(
'Error : Resnet should have depth of either [18, 34, 50, 101, 152]'
)
sys.exit(1)
elif (MODEL_TYPE == 'densenet'):
if (MODEL_DEPTH_OR_VERSION == 121):
model = densenet.densenet121(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 169):
model = densenet.densenet169(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 161):
model = densenet.densenet161(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 201):
model = densenet.densenet201(pretrained=FINETUNE)
else:
print(
'Error : Densenet should have depth of either [121, 169, 161, 201]'
)
sys.exit(1)
elif (MODEL_TYPE == 'inception'):
if (MODEL_DEPTH_OR_VERSION == 3 or MODEL_DEPTH_OR_VERSION == 'v3'):
model = inception.inception_v3(pretrained=FINETUNE)
else:
print('Error : Inception should have version of either [3, ]')
sys.exit(1)
else:
print(
'Error : Network should be either [alexnet / squeezenet / vgg / resnet / densenet / inception]'
)
sys.exit(1)
if (MODEL_TYPE == 'alexnet' or MODEL_TYPE == 'vgg'):
num_ftrs = model.classifier[6].in_features
feature_model = list(model.classifier.children())
feature_model.pop()
feature_model.append(nn.Linear(num_ftrs, class_num))
model.classifier = nn.Sequential(*feature_model)
elif (MODEL_TYPE == 'resnet' or MODEL_TYPE == 'inception'):
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, class_num)
elif (MODEL_TYPE == 'densenet'):
num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, class_num)
return model
def get_network(args, use_gpu=True, num_train=0):
""" return given network
"""
if args.dataset == 'cifar-10':
num_classes = 10
elif args.dataset == 'cifar-100':
num_classes = 100
else:
num_classes = 0
if args.ignoring:
if args.net == 'resnet18':
from models.resnet_ign import resnet18_ign
criterion = nn.CrossEntropyLoss(reduction='none')
net = resnet18_ign(criterion, num_classes=num_classes, num_train=num_train,softmax=args.softmax,isalpha=args.isalpha)
else:
if args.net == 'vgg16':
from models.vgg import vgg16_bn
net = vgg16_bn()
elif args.net == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn()
elif args.net == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn()
elif args.net == 'vgg19':
from models.vgg import vgg19_bn
net = vgg19_bn()
elif args.net == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif args.net == 'densenet161':
from models.densenet import densenet161
net = densenet161()
elif args.net == 'densenet169':
from models.densenet import densenet169
net = densenet169()
elif args.net == 'densenet201':
from models.densenet import densenet201
net = densenet201()
elif args.net == 'googlenet':
from models.googlenet import googlenet
net = googlenet()
elif args.net == 'inceptionv3':
from models.inceptionv3 import inceptionv3
net = inceptionv3()
elif args.net == 'inceptionv4':
from models.inceptionv4 import inceptionv4
net = inceptionv4()
elif args.net == 'inceptionresnetv2':
from models.inceptionv4 import inception_resnet_v2
net = inception_resnet_v2()
elif args.net == 'xception':
from models.xception import xception
net = xception()
elif args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18(num_classes=num_classes)
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34(num_classes=num_classes)
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50(num_classes=num_classes)
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101(num_classes=num_classes)
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152(num_classes=num_classes)
elif args.net == 'preactresnet18':
from models.preactresnet import preactresnet18
net = preactresnet18()
elif args.net == 'preactresnet34':
from models.preactresnet import preactresnet34
net = preactresnet34()
elif args.net == 'preactresnet50':
from models.preactresnet import preactresnet50
net = preactresnet50()
elif args.net == 'preactresnet101':
from models.preactresnet import preactresnet101
net = preactresnet101()
elif args.net == 'preactresnet152':
from models.preactresnet import preactresnet152
net = preactresnet152()
elif args.net == 'resnext50':
from models.resnext import resnext50
net = resnext50()
elif args.net == 'resnext101':
from models.resnext import resnext101
net = resnext101()
elif args.net == 'resnext152':
from models.resnext import resnext152
net = resnext152()
elif args.net == 'shufflenet':
from models.shufflenet import shufflenet
net = shufflenet()
elif args.net == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2()
elif args.net == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet()
elif args.net == 'mobilenet':
from models.mobilenet import mobilenet
net = mobilenet()
elif args.net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2()
elif args.net == 'nasnet':
from models.nasnet import nasnet
net = nasnet()
elif args.net == 'attention56':
from models.attention import attention56
net = attention56()
elif args.net == 'attention92':
from models.attention import attention92
net = attention92()
elif args.net == 'seresnet18':
from models.senet import seresnet18
net = seresnet18()
elif args.net == 'seresnet34':
from models.senet import seresnet34
net = seresnet34()
elif args.net == 'seresnet50':
from models.senet import seresnet50
net = seresnet50()
elif args.net == 'seresnet101':
from models.senet import seresnet101
net = seresnet101()
elif args.net == 'seresnet152':
from models.senet import seresnet152
net = seresnet152()
else:
print('the network name you have entered is not supported yet')
sys.exit()
if use_gpu:
net = net.cuda()
return net
