def main():
loader = get_loader('test', batch_size=1, shuffle=False, num_data=-1)
decoder = SeqModel(False).to(DEVICE)
decoder = SelfAttention().to(DEVICE)
if args.load_model:
dic = torch.load(args.load_model)
state_dict = dic["state_dict"]
decoder.load_state_dict(state_dict)
print('Loss of loaded model: {:.4f}'.format(dic['val_loss']))
#encoder = resnet18(pretrained=False, num_classes=1, in_channels=2).to(DEVICE)
#encoder = vgg11(pretrained=False, num_classes=1).to(DEVICE)
encoder = vgg11_bn(pretrained=False, num_classes=1).to(DEVICE)
dic = torch.load("saved_model/farnback_vgg.tar")
state_dict = dic["state_dict"]
encoder.load_state_dict(state_dict)
modules = list(encoder.children())[:-2]
encoder = nn.Sequential(*modules)
criterion = nn.BCELoss(reduction='mean')
test(encoder, decoder, criterion, loader, DEVICE)
def get_config():
print('=> Building model..')
if args.model == 'vgg16_bn':
from models.vgg import vgg16_bn
net, ratios, policy = vgg16_bn(), vgg16_ratios, vgg16_pruning_policy
elif args.model == 'vgg11_bn':
from models.vgg import vgg11_bn
net, ratios, policy = vgg11_bn(), {20: 0.5}, vgg16_pruning_policy
else:
print("Not support model {}".format(args.model))
raise NotImplementedError
return net, ratios, policy
def main():
# Initialize tensorboard
tensorboard_path = 'runs/{}'.format(args.saved_model_name)
tb_writer = TensorboardWriter(tensorboard_path)
# Initialze early stop
early_stop = EarlyStop(patience=5, mode='max')
# Get dataloader
train_loader = get_loader('train', batch_size=1, shuffle=True, num_data=-1)
val_loader = get_loader('val', batch_size=1, shuffle=False, num_data=-1)
loader = dict(train=train_loader, val=val_loader)
#encoder = CNN_Encoder(in_channels=2).to(DEVICE)
#encoder = resnet18(pretrained=False, num_classes=1, in_channels=2).to(DEVICE)
encoder = vgg11_bn(pretrained=False, num_classes=1,
in_channels=2).to(DEVICE)
if args.load_model: # loader pretrained encoder
dic = torch.load(args.load_model)
state_dict = dic["state_dict"]
encoder.load_state_dict(state_dict)
modules = list(encoder.children())[:-2]
encoder = nn.Sequential(*modules)
#decoder = SeqModel().to(DEVICE)
decoder = SelfAttention().to(DEVICE)
criterion = nn.BCELoss(reduction='mean')
optimizer = optim.Adam(decoder.parameters(), lr=1e-4)
if args.finetune:
optimizer = optim.Adam([{
'params': encoder.parameters(),
'lr': 1e-5
}, {
'params': decoder.parameters()
}],
lr=1e-4)
for epoch in range(1, 100):
train_loss, val_loss, train_acc, val_acc = train_val(
encoder, decoder, criterion, optimizer, loader, epoch, DEVICE)
tb_writer.tensorboard_write(epoch, train_loss, val_loss, train_acc,
val_acc)
saved_model_path = os.path.join("saved_model",
"{}".format(args.saved_model_name))
if early_stop.stop(val_acc, encoder, decoder, epoch, saved_model_path):
break
tb_writer.close()
def get_model(model, dataset, classify=True):
"""
VGG Models
"""
if model == 'vgg11':
model = vgg.vgg11_bn(dataset=dataset, classify=classify)
if model == 'vgg13':
model = vgg.vgg13_bn(dataset=dataset, classify=classify)
if model == 'vgg16':
model = vgg.vgg16_bn(dataset=dataset, classify=classify)
if model == 'vgg19':
model = vgg.vgg19_bn(dataset=dataset, classify=classify)
"""
CyVGG Models
"""
if model == 'cyvgg11':
model = cyvgg.cyvgg11_bn(dataset=dataset, classify=classify)
if model == 'cyvgg13':
model = cyvgg.cyvgg13_bn(dataset=dataset, classify=classify)
if model == 'cyvgg16':
model = cyvgg.cyvgg16_bn(dataset=dataset, classify=classify)
if model == 'cyvgg19':
model = cyvgg.cyvgg19_bn(dataset=dataset, classify=classify)
"""
Resnet Models
"""
if model == 'resnet20':
model = resnet.resnet20(dataset=dataset)
if model == 'resnet32':
model = resnet.resnet32(dataset=dataset)
if model == 'resnet44':
model = resnet.resnet44(dataset=dataset)
if model == 'resnet56':
model = resnet.resnet56(dataset=dataset)
"""
CyResnet Models
"""
if model == 'cyresnet20':
model = cyresnet.cyresnet20(dataset=dataset)
if model == 'cyresnet32':
model = cyresnet.cyresnet32(dataset=dataset)
if model == 'cyresnet44':
model = cyresnet.cyresnet44(dataset=dataset)
if model == 'cyresnet56':
model = cyresnet.cyresnet56(dataset=dataset)
return model
def get_network(args,cfg):
""" return given network
"""
# pdb.set_trace()
if args.net == 'lenet5':
net = LeNet5().cuda()
elif args.net == 'alexnet':
net = alexnet(pretrained=args.pretrain, num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg16':
net = vgg16(pretrained=args.pretrain, num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg13':
net = vgg13(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg11':
net = vgg11(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg19':
net = vgg19(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg16_bn':
net = vgg16_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg13_bn':
net = vgg13_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg11_bn':
net = vgg11_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg19_bn':
net = vgg19_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net =='inceptionv3':
net = inception_v3().cuda()
# elif args.net == 'inceptionv4':
# net = inceptionv4().cuda()
# elif args.net == 'inceptionresnetv2':
# net = inception_resnet_v2().cuda()
elif args.net == 'resnet18':
net = resnet18(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda(args.gpuid)
elif args.net == 'resnet34':
net = resnet34(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'resnet50':
net = resnet50(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda(args.gpuid)
elif args.net == 'resnet101':
net = resnet101(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'resnet152':
net = resnet152(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'squeezenet':
net = squeezenet1_0().cuda()
else:
print('the network name you have entered is not supported yet')
sys.exit()
return net
def load_model():
print("Loading pretrained model ...")
decoder = SelfAttention().to(DEVICE)
dic = torch.load("saved_model/decoder.tar")
state_dict = dic["state_dict"]
decoder.load_state_dict(state_dict)
encoder = vgg11_bn(pretrained=False, num_classes=1).to(DEVICE)
dic = torch.load("saved_model/encoder.tar")
state_dict = dic["state_dict"]
encoder.load_state_dict(state_dict)
modules = list(encoder.children())[:-2]
encoder = nn.Sequential(*modules)
encoder.eval()
decoder.eval()
return encoder, decoder
def get_network(net, num_classes, input_channels=3, use_gpu=True):
""" return given network
"""
if net == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn()
elif net == 'resnet18':
from models.resnet import resnet18
net = resnet18(num_classes=num_classes, input_channels=input_channels)
elif net == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2()
elif net == 'resnet18_modx':
from models.resnet_modx import resnet18
net = resnet18(num_classes=num_classes, input_channels=input_channels)
else:
print('the network name you have entered is not supported yet')
sys.exit()
if use_gpu:
net = net.cuda()
return net
def get_network(args):
""" return given network
"""
if 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 == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50()
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=False):
""" 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()
else:
print('the network name you have entered is not supported yet')
sys.exit()
if 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
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
input_size = 224
train_transform = transforms.Compose([
transforms.RandomResizedCrop(input_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
img1k_calib = Imagenet1kCalib("/home/lqy/data/imagenet1k_calib", train_transform)
# for img in img1k_calib:
# print(img.size())
train_loader = torch.utils.data.DataLoader(img1k_calib,
batch_size=50,
num_workers=1,
pin_memory=True)
vgg11 = vgg11_bn()
lassopruner_config = LassoPruneConfig(vgg11,
"/home/lqy/workshop/amc/checkpoints/vgg11_bn-6002323d.pth",
train_dataloader=train_loader)
lassopruner = LassoPruner(lassopruner_config)
lassopruner.prune_layer(27, 0.5)
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 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 train(opt):
epochs = opt.epochs
batch_size = opt.batch_size
name = opt.name
# Tensorboard settins
log_dir = Path('logs') / name
tb_writer = SummaryWriter(log_dir=log_dir)
# Augmentation
train_transforms = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2471, 0.2435, 0.2616))
])
val_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2471, 0.2435, 0.2616))
])
# Train dataset
train_dataset = torchvision.datasets.CIFAR10('./data',
train=True,
download=True,
transform=train_transforms)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
num_workers = min([os.cpu_count(), batch_size])
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
drop_last=True)
# Validation dataset
val_dataset = torchvision.datasets.CIFAR10('./data',
train=False,
download=True,
transform=val_transforms)
val_dataloader = torch.utils.data.DataLoader(val_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
drop_last=True)
# Network model
model = vgg11_bn()
#model = MySimpleModel()
# GPU-support
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1: # Multi-GPU
model = torch.nn.DataParallel(model)
model.to(device)
# Loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# apex
model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
# Learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=200)
# Loading a weight file (if exists)
weight_file = Path('weights') / (name + '.pth')
best_accuracy = 0.0
start_epoch, end_epoch = (0, epochs)
if os.path.exists(weight_file):
checkpoint = torch.load(weight_file)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
amp.load_state_dict(checkpoint['amp'])
start_epoch = checkpoint['epoch'] + 1
best_accuracy = checkpoint['best_accuracy']
print('resumed from epoch %d' % start_epoch)
for epoch in range(start_epoch, end_epoch):
print('epoch: %d/%d' % (epoch, end_epoch - 1))
# Training
t0 = time.time()
epoch_loss = train_one_epoch(train_dataloader, model, loss_fn,
optimizer, device)
t1 = time.time()
print('loss=%.4f (took %.2f sec)\n' % (epoch_loss, t1 - t0))
lr_scheduler.step()
# validation
val_epoch_loss, accuracy = val_one_epoch(val_dataloader, model,
loss_fn, device)
print('[validation] loss=%.4f, accuracy=%.4f' %
(val_epoch_loss, accuracy))
# Saving the best status into a weight file
if accuracy > best_accuracy:
best_weight_file = Path('weights') / (name + '_best.pth')
best_accuracy = accuracy
state = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'amp': amp.state_dict(),
'epoch': epoch,
'best_accuracy': best_accuracy
}
torch.save(state, best_weight_file)
print('best accuracy=>saved\n')
# saving the current status into a weight file
state = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'amp': amp.state_dict(),
'epoch': epoch,
'best_accuracy': best_accuracy
}
torch.save(state, weight_file)
# tensorboard logging
tb_writer.add_scalar('train_epoch_loss', epoch_loss, epoch)
tb_writer.add_scalar('val_epoch_loss', val_epoch_loss, epoch)
tb_writer.add_scalar('val_accuracy', accuracy, epoch)
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
def get_model(args, model_path=None):
"""
:param args: super arguments
:param model_path: if not None, load already trained model parameters.
:return: model
"""
if args.scratch: # train model from scratch
pretrained = False
model_dir = None
print("=> Loading model '{}' from scratch...".format(args.model))
else: # train model with pretrained model
pretrained = True
model_dir = os.path.join(args.root_path, args.pretrained_models_path)
print("=> Loading pretrained model '{}'...".format(args.model))
if args.model.startswith('resnet'):
if args.model == 'resnet18':
model = resnet18(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'resnet34':
model = resnet34(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'resnet50':
model = resnet50(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'resnet101':
model = resnet101(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'resnet152':
model = resnet152(pretrained=pretrained, model_dir=model_dir)
model.fc = nn.Linear(model.fc.in_features, args.num_classes)
elif args.model.startswith('vgg'):
if args.model == 'vgg11':
model = vgg11(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'vgg11_bn':
model = vgg11_bn(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'vgg13':
model = vgg13(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'vgg13_bn':
model = vgg13_bn(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'vgg16':
model = vgg16(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'vgg16_bn':
model = vgg16_bn(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'vgg19':
model = vgg19(pretrained=pretrained, model_dir=model_dir)
elif args.model == 'vgg19_bn':
model = vgg19_bn(pretrained=pretrained, model_dir=model_dir)
model.classifier[6] = nn.Linear(model.classifier[6].in_features, args.num_classes)
elif args.model == 'alexnet':
model = alexnet(pretrained=pretrained, model_dir=model_dir)
model.classifier[6] = nn.Linear(model.classifier[6].in_features, args.num_classes)
# Load already trained model parameters and go on training
if model_path is not None:
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['model'])
return model
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
if __name__ == '__main__':
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
else:
print("=> creating model '{}'".format(args.arch))
if args.arch == 'vgg11_cbn':
model = vgg11_cbn(pretrained=args.pretrained)
elif args.arch == 'vgg16_cbn':
model = vgg16_cbn()
elif args.arch == 'vgg19_cbn':
model = vgg19_cbn()
elif args.arch == 'vgg11_bn':
model = vgg11_bn(pretrained=args.pretrained)
elif args.arch == 'vgg16_bn':
model = vgg16_bn()
elif args.arch == 'vgg19_bn':
model = vgg19_bn()
elif args.arch == 'resnet18':
model = resnet18(pretrained=args.pretrained)
elif args.arch == 'resnet50':
model = resnet50(pretrained=args.pretrained)
elif args.arch == 'resnet101':
model = resnet101(pretrained=args.pretrained)
elif args.arch == 'resnet18_cbn':
model = resnet18_cbn()
elif args.arch == 'resnet50_cbn':
model = resnet50_cbn()
elif args.arch == 'resnet101_cbn':
def main():
global args, best_err1
args = parser.parse_args()
# TensorBoard configure
if args.tensorboard:
configure('%s_checkpoints/%s'%(args.dataset, args.expname))
# CUDA
os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_ids)
if torch.cuda.is_available():
cudnn.benchmark = True # https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936
kwargs = {'num_workers': 2, 'pin_memory': True}
else:
kwargs = {'num_workers': 2}
# Data loading code
if args.dataset == 'cifar10':
normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
std=[0.2023, 0.1994, 0.2010])
elif args.dataset == 'cifar100':
normalize = transforms.Normalize(mean=[0.5071, 0.4865, 0.4409],
std=[0.2634, 0.2528, 0.2719])
elif args.dataset == 'cub':
normalize = transforms.Normalize(mean=[0.4862, 0.4973, 0.4293],
std=[0.2230, 0.2185, 0.2472])
elif args.dataset == 'webvision':
normalize = transforms.Normalize(mean=[0.49274242, 0.46481857, 0.41779366],
std=[0.26831809, 0.26145372, 0.27042758])
else:
raise Exception('Unknown dataset: {}'.format(args.dataset))
# Transforms
if args.augment:
train_transform = transforms.Compose([
transforms.RandomResizedCrop(args.train_image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
train_transform = transforms.Compose([
transforms.RandomResizedCrop(args.train_image_size),
transforms.ToTensor(),
normalize,
])
val_transform = transforms.Compose([
transforms.Resize(args.test_image_size),
transforms.CenterCrop(args.test_crop_image_size),
transforms.ToTensor(),
normalize
])
# Datasets
num_classes = 10 # default 10 classes
if args.dataset == 'cifar10':
train_dataset = datasets.CIFAR10('./data/', train=True, download=True, transform=train_transform)
val_dataset = datasets.CIFAR10('./data/', train=False, download=True, transform=val_transform)
num_classes = 10
elif args.dataset == 'cifar100':
train_dataset = datasets.CIFAR100('./data/', train=True, download=True, transform=train_transform)
val_dataset = datasets.CIFAR100('./data/', train=False, download=True, transform=val_transform)
num_classes = 100
elif args.dataset == 'cub':
train_dataset = datasets.ImageFolder('/media/ouc/30bd7817-d3a1-4e83-b7d9-5c0e373ae434/DuAngAng/datasets/CUB-200-2011/train/',
transform=train_transform)
val_dataset = datasets.ImageFolder('/media/ouc/30bd7817-d3a1-4e83-b7d9-5c0e373ae434/DuAngAng/datasets/CUB-200-2011/test/',
transform=val_transform)
num_classes = 200
elif args.dataset == 'webvision':
train_dataset = datasets.ImageFolder('/media/ouc/30bd7817-d3a1-4e83-b7d9-5c0e373ae434/LiuJing/WebVision/info/train',
transform=train_transform)
val_dataset = datasets.ImageFolder('/media/ouc/30bd7817-d3a1-4e83-b7d9-5c0e373ae434/LiuJing/WebVision/info/val',
transform=val_transform)
num_classes = 1000
else:
raise Exception('Unknown dataset: {}'.format(args.dataset))
# Data Loader
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batchsize, shuffle=True, **kwargs)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False, **kwargs)
# Create model
if args.model == 'AlexNet':
model = alexnet(pretrained=False, num_classes=num_classes)
elif args.model == 'VGG':
use_batch_normalization = True # default use Batch Normalization
if use_batch_normalization:
if args.depth == 11:
model = vgg11_bn(pretrained=False, num_classes=num_classes)
elif args.depth == 13:
model = vgg13_bn(pretrained=False, num_classes=num_classes)
elif args.depth == 16:
model = vgg16_bn(pretrained=False, num_classes=num_classes)
elif args.depth == 19:
model = vgg19_bn(pretrained=False, num_classes=num_classes)
else:
raise Exception('Unsupport VGG detph: {}, optional depths: 11, 13, 16 or 19'.format(args.depth))
else:
if args.depth == 11:
model = vgg11(pretrained=False, num_classes=num_classes)
elif args.depth == 13:
model = vgg13(pretrained=False, num_classes=num_classes)
elif args.depth == 16:
model = vgg16(pretrained=False, num_classes=num_classes)
elif args.depth == 19:
model = vgg19(pretrained=False, num_classes=num_classes)
else:
raise Exception('Unsupport VGG detph: {}, optional depths: 11, 13, 16 or 19'.format(args.depth))
elif args.model == 'Inception':
model = inception_v3(pretrained=False, num_classes=num_classes)
elif args.model == 'ResNet':
if args.depth == 18:
model = resnet18(pretrained=False, num_classes=num_classes)
elif args.depth == 34:
model = resnet34(pretrained=False, num_classes=num_classes)
elif args.depth == 50:
model = resnet50(pretrained=False, num_classes=num_classes)
elif args.depth == 101:
model = resnet101(pretrained=False, num_classes=num_classes)
elif args.depth == 152:
model = resnet152(pretrained=False, num_classes=num_classes)
else:
raise Exception('Unsupport ResNet detph: {}, optional depths: 18, 34, 50, 101 or 152'.format(args.depth))
elif args.model == 'MPN-COV-ResNet':
if args.depth == 18:
model = mpn_cov_resnet18(pretrained=False, num_classes=num_classes)
elif args.depth == 34:
model = mpn_cov_resnet34(pretrained=False, num_classes=num_classes)
elif args.depth == 50:
model = mpn_cov_resnet50(pretrained=False, num_classes=num_classes)
elif args.depth == 101:
model = mpn_cov_resnet101(pretrained=False, num_classes=num_classes)
elif args.depth == 152:
model = mpn_cov_resnet152(pretrained=False, num_classes=num_classes)
else:
raise Exception('Unsupport MPN-COV-ResNet detph: {}, optional depths: 18, 34, 50, 101 or 152'.format(args.depth))
else:
raise Exception('Unsupport model'.format(args.model))
# Get the number of model parameters
print('Number of model parameters: {}'.format(sum([p.data.nelement() for p in model.parameters()])))
if torch.cuda.is_available():
model = model.cuda()
# 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)
args.start_epoch = checkpoint['epoch']
best_err1 = checkpoint['best_err1']
model.load_state_dict(checkpoint['state_dict'])
print("==> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("==> no checkpoint found at '{}'".format(args.resume))
print(model)
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss()
if torch.cuda.is_available():
criterion = criterion.cuda()
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# Train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# Evaluate on validation set
err1 = validate(val_loader, model, criterion, epoch)
# Remember best err1 and save checkpoint
is_best = (err1 <= best_err1)
best_err1 = min(err1, best_err1)
print("Current best accuracy (error):", best_err1)
save_checkpoint({
'epoch': epoch+1,
'state_dict': model.state_dict(),
'best_err1': best_err1,
}, is_best)
print("Best accuracy (error):", best_err1)
def get_model(args):
assert args.model in [
'derpnet', 'alexnet', 'resnet', 'vgg', 'vgg_attn', 'inception'
]
if args.model == 'alexnet':
model = alexnet.alexnet(pretrained=args.pretrained,
n_channels=args.n_channels,
num_classes=args.n_classes)
elif args.model == 'inception':
model = inception.inception_v3(pretrained=args.pretrained,
aux_logits=False,
progress=True,
num_classes=args.n_classes)
elif args.model == 'vgg':
assert args.model_depth in [11, 13, 16, 19]
if args.model_depth == 11:
model = vgg.vgg11_bn(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
if args.model_depth == 13:
model = vgg.vgg13_bn(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
if args.model_depth == 16:
model = vgg.vgg16_bn(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
if args.model_depth == 19:
model = vgg.vgg19(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
elif args.model == 'vgg_attn':
assert args.model_depth in [11, 13, 16, 19]
if args.model_depth == 11:
model = vgg_attn.vgg11_bn(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
if args.model_depth == 13:
model = vgg_attn.vgg11_bn(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
if args.model_depth == 16:
model = vgg_attn.vgg11_bn(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
if args.model_depth == 19:
model = vgg_attn.vgg11_bn(pretrained=args.pretrained,
progress=True,
num_classes=args.n_classes)
elif args.model == 'derpnet':
model = derp_net.Net(n_channels=args.n_channels,
num_classes=args.n_classes)
elif args.model == 'resnet':
assert args.model_depth in [10, 18, 34, 50, 101, 152, 200]
if args.model_depth == 10:
model = resnet.resnet10(pretrained=args.pretrained,
num_classes=args.n_classes)
elif args.model_depth == 18:
model = resnet.resnet18(pretrained=args.pretrained,
num_classes=args.n_classes)
elif args.model_depth == 34:
model = resnet.resnet34(pretrained=args.pretrained,
num_classes=args.n_classes)
elif args.model_depth == 50:
model = resnet.resnet50(pretrained=args.pretrained,
num_classes=args.n_classes)
elif args.model_depth == 101:
model = resnet.resnet101(pretrained=args.pretrained,
num_classes=args.n_classes)
elif args.model_depth == 152:
model = resnet.resnet152(pretrained=args.pretrained,
num_classes=args.n_classes)
elif args.model_depth == 200:
model = resnet.resnet200(pretrained=args.pretrained,
num_classes=args.n_classes)
if args.pretrained and args.pretrain_path and not args.model == 'alexnet' and not args.model == 'vgg' and not args.model == 'resnet':
print('loading pretrained model {}'.format(args.pretrain_path))
pretrain = torch.load(args.pretrain_path)
assert args.arch == pretrain['arch']
# here all the magic happens: need to pick the parameters which will be adjusted during training
# the rest of the params will be frozen
pretrain_dict = {
key[7:]: value
for key, value in pretrain['state_dict'].items()
if key[7:9] != 'fc'
}
from collections import OrderedDict
pretrain_dict = OrderedDict(pretrain_dict)
# https://stackoverflow.com/questions/972/adding-a-method-to-an-existing-object-instance
import types
model.load_state_dict = types.MethodType(load_my_state_dict, model)
old_dict = copy.deepcopy(
model.state_dict()) # normal copy() just gives a reference
model.load_state_dict(pretrain_dict)
new_dict = model.state_dict()
num_features = model.fc.in_features
if args.model == 'densenet':
model.classifier = nn.Linear(num_features, args.n_classes)
else:
#model.fc = nn.Sequential(nn.Linear(num_features, 1028), nn.ReLU(), nn.Dropout(0.5), nn.Linear(1028, args.n_finetune_classes))
model.fc = nn.Linear(num_features, args.n_classes)
# parameters = get_fine_tuning_parameters(model, args.ft_begin_index)
parameters = model.parameters() # fine-tunining EVERYTHIIIIIANG
# parameters = model.fc.parameters() # fine-tunining ONLY FC layer
return model, parameters
return model, model.parameters()
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
def get_network(key, num_cls=2, use_gpu=False):
""" return given network
"""
if key == 'vgg16':
from models.vgg import vgg16_bn
net = vgg16_bn(num_cls)
elif key == 'vgg13':
from models.vgg import vgg13_bn
net = vgg13_bn(num_cls)
elif key == 'vgg11':
from models.vgg import vgg11_bn
net = vgg11_bn(num_cls)
elif key == 'vgg19':
from models.vgg import vgg19_bn
net = vgg19_bn(num_cls)
elif key == 'resnext':
print('we will continue')
elif key == 'efficientNetb0':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'efficientNetb1':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'efficientNetb2':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'efficientNetb3':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'efficientNetb4':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'efficientNetb5':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'efficientNetb6':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'efficientNetb7':
from models.torchefficient import make_model
net = make_model(key, num_cls)
elif key == 'resnext50_32x8d':
from models.resnext import make_model
net = make_model(key)
elif key == 'resnext101_32x8d':
from models.resnext import make_model
net = make_model(key)
elif key == 'resnet50':
from models.resnet import make_model
net = make_model(key)
elif key == 'resnet18':
from models.resnet import make_model
net = make_model(key)
elif key == 'resnet34':
from models.resnet import make_model
net = make_model(key)
elif key == 'resnet101':
from models.resnet import make_model
net = make_model(key)
else:
print('the network name you have entered is not supported yet')
sys.exit()
if use_gpu:
net = net.cuda()
return net
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(netname, use_gpu=True):
""" return given network
"""
if netname == 'vgg16':
from models.vgg import vgg16_bn #!
net = vgg16_bn()
elif netname == 'vgg16_cbn':
from models.vgg_nobn import vgg16_cbn #!
net = vgg16_cbn()
elif netname == 'vgg11':
from models.vgg import vgg11_bn #!
net = vgg11_bn()
elif netname == 'vgg11_cbn':
from models.vgg_nobn import vgg11_cbn #!
net = vgg11_cbn()
elif netname == 'vgg11_nobn':
from models.vgg_nobn import vgg11_nobn #!
net = vgg11_nobn()
elif netname == 'vgg16_nobn':
from models.vgg_nobn import vgg16_nobn #!
net = vgg16_nobn()
elif netname == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif netname == 'resnet50':
from models.resnet import resnet50
net = resnet50()
elif netname == 'resnet101':
from models.resnet import resnet101
net = resnet101()
elif netname == 'resnet18_nobn':
from models.resnet_nobn import resnet18_nobn
net = resnet18_nobn()
elif netname == 'resnet18_fixup':
from models.resnet_fixup import resnet18
net = resnet18()
elif netname == 'resnet50_fixup':
from models.resnet_fixup import resnet50
net = resnet50()
elif netname == 'resnet18_cbn':
from models.resnet_nobn import resnet18_cbn
net = resnet18_cbn()
elif netname == 'resnet50_cbn':
from models.resnet_nobn import resnet50_cbn
net = resnet50_cbn()
elif netname == 'resnet50_nobn':
from models.resnet_nobn import resnet50_nobn
net = resnet50_nobn()
elif netname == 'resnet101_cbn':
from models.resnet_nobn import resnet101_cbn
net = resnet101_cbn()
elif netname == 'densenet121':
from models.densenet import densenet121
net = densenet121()
elif netname == 'densenet121_cbn':
from models.densenet_nobn import densenet121
net = densenet121()
elif netname == 'shufflenetv2':
from models.shufflenetv2 import shufflenetv2
net = shufflenetv2()
elif netname == 'shufflenetv2_cbn':
from models.shufflenetv2_nobn import shufflenetv2_cbn
net = shufflenetv2_cbn()
elif netname == 'shufflenetv2_nobn':
from models.shufflenetv2_nobn import shufflenetv2_nobn
net = shufflenetv2_nobn()
elif netname == 'squeezenet':
from models.squeezenet import squeezenet
net = squeezenet()
elif netname == 'squeezenet_nobn':
from models.squeezenet_nobn import squeezenet_nobn
net = squeezenet_nobn()
elif netname == 'squeezenet_cbn':
from models.squeezenet_nobn import squeezenet_cbn
net = squeezenet_cbn()
elif netname == 'seresnet18':
from models.senet import seresnet18
net = seresnet18()
elif netname == 'seresnet50':
from models.senet import seresnet50
net = seresnet50()
elif netname == 'seresnet18_cbn':
from models.senet_nobn import seresnet18
net = seresnet18()
elif netname == 'seresnet50_cbn':
from models.senet_nobn import seresnet50
net = seresnet50()
elif netname == 'fixup_cbn':
from models.fixup_resnet_cifar import fixup_resnet56
net = fixup_resnet56(cbn=True)
elif netname == 'fixup':
from models.fixup_resnet_cifar import fixup_resnet56
net = fixup_resnet56()
elif netname == 'mobilenetv2':
from models.mobilenetv2 import mobilenetv2
net = mobilenetv2()
elif netname == 'mobilenetv2_cbn':
from models.mobilenetv2_nobn import mobilenetv2
net = mobilenetv2()
else:
print(netname)
print('the network name you have entered is not supported yet')
sys.exit()
if use_gpu:
# net = torch.nn.parallel.DataParallel(net)
net = net.cuda()
return net
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_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)
