def main():
args = parser.parse_args()
model = None
dtype = torch.cuda.FloatTensor
if args.dataset == 'tiny_imagenet':
dataloader = TinyImagenetDataLoader()
num_classes = 200
elif args.dataset == 'dogs':
dataloader = DogsDataLoader()
num_classes = 121
else:
dataloader = Cifar10DataLoader()
num_classes = 10
if args.model == 'resnet34':
model = ResNet34(num_classes).type(dtype)
elif args.model == 'resnet50':
model = ResNet50(num_classes).type(dtype)
elif args.model == 'resnet50-preact':
model = ResNet50(num_classes, pre_activation=True).type(dtype)
elif args.model == 'resnet50-imported':
model = ResNetImported(num_classes).type(dtype)
agent = Agent(model, dataloader)
agent.load_checkpoint(args.checkpoint)
# agent.check_accuracy()
agent.print_accuracy()
def build_model():
# Model
print('\n==> Building model..')
# net = VGG('VGG19')
# net = ResNet50()
# net = PreActResNet18()
# net = GoogLeNet()
# net = DenseNet121()
# net = ResNeXt29_2x64d()
# net = MobileNet()
# net = MobileNetV2()
# net = DPN92()
# net = ShuffleNetG2()
# net = SENet18()
# net = ShuffleNetV2(1)
if args.arch == 'resnet50':
net = ResNet50()
elif args.arch == 'densenet121':
net = DenseNet121()
print('Select Model {} ...'.format(args.arch))
net = net.to(device)
if device == 'cuda':
cudnn.benchmark = True
return net
def get_model(name, config):
if name == "densenet121":
return DenseNet(config.num_classes, config.densenet_weights, config)
elif name == "ResNet50":
return ResNet50(config.num_classes, config.resnet_weights, config)
else:
raise NotImplementedError
def get_model(arch, num_classes, channels=3):
"""
Args:
arch: string, Network architecture
num_classes: int, Number of classes
channels: int, Number of input channels
Returns:
model, nn.Module, generated model
"""
if arch.lower() == "resnet18":
model = ResNet18(channels, num_classes)
elif arch.lower() == "resnet34":
model = ResNet34(channels, num_classes)
elif arch.lower() == "resnet50":
model = ResNet50(channels, num_classes)
elif arch.lower() == "resnet101":
model = ResNet101(channels, num_classes)
elif arch.lower() == "resnet152":
model = ResNet152(channels, num_classes)
elif arch.lower() == "mobilenet_v1":
model = MobileNetV1(num_classes, channels)
elif arch.lower() == "mobilenet_v2":
model = MobileNetV2(num_classes, channels)
else:
raise NotImplementedError(
f"{arch} not implemented. "
f"For supported architectures see documentation")
return model
def initialize_model(model_name):
'''
Initialise a model with a custom head to predict both sequence length and digits
Parameters
----------
model_name : str
Model Name can be either:
ResNet
VGG
BaselineCNN
ConvNet
BaselineCNN_dropout
Returns
-------
model : object
The model to be initialize
'''
if model_name[:3] == "VGG":
model = VGG(model_name, num_classes=7)
model.classifier = CustomHead(512)
elif model_name[:6] == "ResNet":
if model_name == "ResNet18":
model = ResNet18(num_classes=7)
model.linear = CustomHead(512)
elif model_name == "ResNet34":
model = ResNet18(num_classes=7)
model.linear = CustomHead(512)
elif model_name == "ResNet50":
model = ResNet50(num_classes=7)
model.linear = CustomHead(512 * 4)
elif model_name == "ResNet101":
model = ResNet101(num_classes=7)
model.linear = CustomHead(512 * 4)
elif model_name == "ResNet152":
model = ResNet152(num_classes=7)
model.linear = CustomHead(512 * 4)
elif model_name == "BaselineCNN":
model = BaselineCNN(num_classes=7)
model.fc2 = CustomHead(4096)
elif model_name == "BaselineCNN_dropout":
model = BaselineCNN_dropout(num_classes=7, p=0.5)
model.fc2 = CustomHead(4096)
return model
def get_new_model(tmp_scale = True, num_classes = args.num_classes):
if args.model == 'resnet18':
return ResNet18(tmp_scale = tmp_scale, num_classes = num_classes)
elif args.model == 'resnet50':
return ResNet50(tmp_scale = tmp_scale, num_classes = num_classes)
elif args.model == 'resnet101':
return ResNet101(tmp_scale = tmp_scale, num_classes = num_classes)
elif args.model == 'inceptionv4':
return inceptionv4(tmp_scale = tmp_scale, num_classes = num_classes)
elif args.model == 'densenet':
return DenseNet(tmp_scale = tmp_scale)
def main():
args = parser.parse_args()
model = None
scheduler = None
if args.dataset == 'tiny_imagenet':
dataloader = TinyImagenetDataLoader()
num_classes = 200
elif args.dataset == 'dogs':
dataloader = DogsDataLoader()
num_classes = 121
else:
dataloader = Cifar10DataLoader()
num_classes = 10
if args.model == 'resnet34':
model = ResNet34(num_classes).cuda()
elif args.model == 'resnet50':
model = ResNet50(num_classes).cuda()
elif args.model == 'resnet50-preact':
model = ResNet50(num_classes, pre_activation=True).cuda()
elif args.model == 'resnet50-imported':
model = ResNetImported(num_classes).cuda()
elif args.model == 'resnext29':
# model = ResNext29(num_classes).cuda()
model = resnext29_8x64d(num_classes).cuda()
if args.optimizer == 'momentum':
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=5,
gamma=0.1)
else:
optimizer = optim.Adam(model.parameters(), lr=args.lr)
loss_fn = nn.CrossEntropyLoss().cuda()
agent = Agent(model, dataloader)
agent.train(loss_fn, args.epochs, optimizer, scheduler, args.log_and_save)
def get_model_by_name(name='resnet50'):
model_dict = {
'resnet50': ResNet50(),
'facenet': facenet(),
'vargfacenet': varGFaceNet(),
'inception_resnet': inceptionresnet2(),
'wrn': wrn50()
}
print("use model {}".format(name))
if name in model_dict.keys():
model = model_dict[name]
else:
raise ValueError("no model like this")
return model
def get_encoder(config):
if config['encoder_type'] == "vgg19":
encoder = VGG19()
elif config['encoder_type'] == "vgg19_bn":
encoder = VGG19_bn()
elif config['encoder_type'] == "resnet152":
encoder = ResNet152()
elif config['encoder_type'] == "resnet34":
encoder = ResNet34()
elif config['encoder_type'] == "resnet50":
encoder = ResNet50()
else:
raise RuntimeError("invalid encoder type")
return encoder
def get_classifier(mode, n_classes=10):
if mode == 'resnet18':
classifier = ResNet18(num_classes=n_classes)
elif mode == 'resnet34':
classifier = ResNet34(num_classes=n_classes)
elif mode == 'resnet50':
classifier = ResNet50(num_classes=n_classes)
elif mode == 'resnet18_imagenet':
classifier = resnet18(num_classes=n_classes)
elif mode == 'resnet50_imagenet':
classifier = resnet50(num_classes=n_classes)
else:
raise NotImplementedError()
return classifier
def get_classifier(mode, n_classes=10):
if mode == 'resnet18':
classifier = ResNet18(num_classes=n_classes)
elif mode == 'resnet34':
classifier = ResNet34(num_classes=n_classes)
elif mode == 'resnet50':
classifier = ResNet50(num_classes=n_classes)
elif mode == 'resnet18_imagenet':
classifier = resnet18(num_classes=n_classes)
elif mode == 'resnet50_imagenet':
classifier = resnet50(num_classes=n_classes)
elif mode == 'live':
classifier = FeatherNet(input_size=128, se=True, avgdown=True)
# classifier = ResNet18(num_classes=n_classes)
# classifier = LiveModel()
else:
raise NotImplementedError()
return classifier
def get_model(args):
if args.dataset == 'cifar-10':
num_classes = 10
elif args.dataset == 'cifar-100':
num_classes = 100
else:
raise NotImplementedError
if 'contrastive' in args.train_type or 'linear_eval' in args.train_type:
contrastive_learning = True
else:
contrastive_learning = False
if args.model == 'ResNet18':
model = ResNet18(num_classes, contrastive_learning)
print('ResNet18 is loading ...')
elif args.model == 'ResNet50':
model = ResNet50(num_classes, contrastive_learning)
print('ResNet 50 is loading ...')
return model
def get_model(model):
model_path = '../saved'
if model == 'LeNet-5':
net = LeNet()
model_name = 'lenet.pth'
elif model == 'VGG-16':
net = Vgg16_Net()
model_name = 'vgg16.pth'
elif model == 'ResNet18':
net = ResNet18()
model_name = 'resnet18.pth'
elif model == 'ResNet34':
net = ResNet34()
model_name = 'resnet34.pth'
elif model == 'ResNet50':
net = ResNet50()
model_name = 'resnet50.pth'
else:
net = ResNet101()
model_name = 'resnet101.pth'
return net, os.path.join(model_path, model_name)
def model_init(self, args):
# Network
if args.dataset == 'MNIST':
print("MNIST")
self.net = cuda(ToyNet_MNIST(y_dim=self.y_dim), self.cuda)
elif args.dataset == 'CIFAR10':
print("Dataset used CIFAR10")
if args.network_choice == 'ToyNet':
self.net = cuda(ToyNet_CIFAR10(y_dim=self.y_dim), self.cuda)
elif args.network_choice == 'ResNet18':
self.net = cuda(ResNet18(), self.cuda)
elif args.network_choice == 'ResNet34':
self.net = cuda(ResNet34(), self.cuda)
elif args.network_choice == 'ResNet50':
self.net = cuda(ResNet50(), self.cuda)
self.net.weight_init(_type='kaiming')
# setup optimizer
self.optim = optim.Adam([{
'params': self.net.parameters(),
'lr': self.lr
}],
betas=(0.5, 0.999))
def get_classifier(self, is_multi_class):
if self.params.model == 'resnet18':
from models.resnet import ResNet18
classifier = ResNet18(num_classes=self.params.n_classes)
elif self.params.model == 'resnet34':
from models.resnet import ResNet34
classifier = ResNet34(num_classes=self.params.n_classes)
elif self.params.model == 'resnet50':
from models.resnet import ResNet50
classifier = ResNet50(num_classes=self.params.n_classes)
elif self.params.model == 'resnet18_imagenet':
if is_multi_class:
from models.resnet_imagenet_multiclass_infer import resnet18
else:
from models.resnet_imagenet import resnet18
classifier = resnet18(num_classes=self.params.n_classes)
elif self.params.model == 'resnet50_imagenet':
from models.resnet_imagenet import resnet50
classifier = resnet50(num_classes=self.params.n_classes)
else:
raise NotImplementedError()
return classifier
def select_model(model_name: str,
pretrain: bool,
n_class: int,
onehot: int,
onehot2=0):
if model_name == 'resnet50':
model = ResNet50(onehot=onehot, onehot2=onehot2)
elif model_name == "resnext":
model = resnext50_32x4d(onehot=onehot, onehot2=onehot2)
elif model_name == "resnet101":
model = resnet101(onehot=onehot, onehot2=onehot2)
elif model_name == "resnext101":
model = resnext101_32x8d(onehot=onehot, onehot2=onehot2)
elif model_name == "resnext101_32x16d":
model = resnext101_32x16d(onehot=onehot, onehot2=onehot2)
elif model_name == 'densenet':
model = DenseNet121(onehot=onehot, onehot2=onehot2)
elif model_name == 'densenet201':
model = densenet201(onehot=onehot, onehot2=onehot2)
elif model_name == "nest200":
model = resnest200(onehot=onehot, onehot2=onehot2)
elif model_name == "nest101":
model = resnest101(onehot=onehot, onehot2=onehot2)
elif model_name == "efficient_b2":
model = EfficientNet_B2(onehot=onehot, onehot2=onehot2)
elif model_name == "efficient_b5":
model = EfficientNet_B5(onehot=onehot, onehot2=onehot2)
elif model_name == "efficient_b6":
model = EfficientNet_B6(onehot=onehot, onehot2=onehot2)
elif model_name == "xception":
model = xception(onehot=onehot, onehot2=onehot2)
else:
raise NotImplementedError(
'Please select in [resnext101, resnext101_32x16d, nest200, densenet201]'
)
return model
def select(self, model, path_fc=False, upsample='pixel'):
if model == 'cnn':
net = SimpleModel(
in_shape=self.in_shape,
activation=self.activation,
num_classes=self.num_classes,
filters=self.filters,
)
else:
assert (self.dataset != 'MNIST' and self.dataset != 'Fashion-MNIST'
), "Cannot use resnet or densenet for mnist style data"
if model == 'resnet':
assert self.resdepth in [
18, 34, 50, 101, 152
], "Non-standard and unsupported resnet depth ({})".format(
self.resdepth)
if self.resdepth == 18:
net = ResNet18()
elif self.resdepth == 34:
net = ResNet34()
elif self.resdepth == 50:
net = ResNet50()
elif self.resdepth == 101:
net = ResNet101()
else:
net = ResNet152()
elif model == 'densenet':
assert self.resdepth in [
121, 161, 169, 201
], "Non-standard and unsupported densenet depth ({})".format(
self.resdepth)
if self.resdepth == 121:
net = DenseNet121()
elif self.resdepth == 161:
net = DenseNet161()
elif self.resdepth == 169:
net = DenseNet169()
else:
net = DenseNet201()
elif model == 'preact_resnet':
assert self.resdepth in [
10, 18, 34, 50, 101, 152
], "Non-standard and unsupported preact resnet depth ({})".format(
self.resdepth)
if self.resdepth == 10:
net = PreActResNet10(path_fc=path_fc,
num_classes=self.num_classes,
upsample=upsample)
elif self.resdepth == 18:
net = PreActResNet18()
elif self.resdepth == 34:
net = PreActResNet34()
elif self.resdepth == 50:
net = PreActResNet50()
elif self.resdepth == 101:
net = PreActResNet101()
else:
net = PreActResNet152()
elif model == 'wresnet':
assert (
(self.resdepth - 4) % 6 == 0
), "Wideresnet depth of {} not supported, must fulfill: (depth - 4) % 6 = 0".format(
self.resdepth)
net = WideResNet(depth=self.resdepth,
num_classes=self.num_classes,
widen_factor=self.widen_factor)
return net
def get_network(name: str, num_classes: int) -> None:
return \
AlexNet(
num_classes=num_classes) if name == 'AlexNet' else\
DenseNet201(
num_classes=num_classes) if name == 'DenseNet201' else\
DenseNet169(
num_classes=num_classes) if name == 'DenseNet169' else\
DenseNet161(
num_classes=num_classes) if name == 'DenseNet161' else\
DenseNet121(
num_classes=num_classes) if name == 'DenseNet121' else\
DenseNet121CIFAR(
num_classes=num_classes) if name == 'DenseNet121CIFAR' else\
GoogLeNet(
num_classes=num_classes) if name == 'GoogLeNet' else\
InceptionV3(
num_classes=num_classes) if name == 'InceptionV3' else\
MNASNet_0_5(
num_classes=num_classes) if name == 'MNASNet_0_5' else\
MNASNet_0_75(
num_classes=num_classes) if name == 'MNASNet_0_75' else\
MNASNet_1(
num_classes=num_classes) if name == 'MNASNet_1' else\
MNASNet_1_3(
num_classes=num_classes) if name == 'MNASNet_1_3' else\
MobileNetV2(
num_classes=num_classes) if name == 'MobileNetV2' else\
ResNet18(
num_classes=num_classes) if name == 'ResNet18' else\
ResNet34(
num_classes=num_classes) if name == 'ResNet34' else\
ResNet34CIFAR(
num_classes=num_classes) if name == 'ResNet34CIFAR' else\
ResNet50CIFAR(
num_classes=num_classes) if name == 'ResNet50CIFAR' else\
ResNet101CIFAR(
num_classes=num_classes) if name == 'ResNet101CIFAR' else\
ResNet18CIFAR(
num_classes=num_classes) if name == 'ResNet18CIFAR' else\
ResNet50(
num_classes=num_classes) if name == 'ResNet50' else\
ResNet101(
num_classes=num_classes) if name == 'ResNet101' else\
ResNet152(
num_classes=num_classes) if name == 'ResNet152' else\
ResNeXt50(
num_classes=num_classes) if name == 'ResNext50' else\
ResNeXtCIFAR(
num_classes=num_classes) if name == 'ResNeXtCIFAR' else\
ResNeXt101(
num_classes=num_classes) if name == 'ResNext101' else\
WideResNet50(
num_classes=num_classes) if name == 'WideResNet50' else\
WideResNet101(
num_classes=num_classes) if name == 'WideResNet101' else\
ShuffleNetV2_0_5(
num_classes=num_classes) if name == 'ShuffleNetV2_0_5' else\
ShuffleNetV2_1(
num_classes=num_classes) if name == 'ShuffleNetV2_1' else\
ShuffleNetV2_1_5(
num_classes=num_classes) if name == 'ShuffleNetV2_1_5' else\
ShuffleNetV2_2(
num_classes=num_classes) if name == 'ShuffleNetV2_2' else\
SqueezeNet_1(
num_classes=num_classes) if name == 'SqueezeNet_1' else\
SqueezeNet_1_1(
num_classes=num_classes) if name == 'SqueezeNet_1_1' else\
VGG11(
num_classes=num_classes) if name == 'VGG11' else\
VGG11_BN(
num_classes=num_classes) if name == 'VGG11_BN' else\
VGG13(
num_classes=num_classes) if name == 'VGG13' else\
VGG13_BN(
num_classes=num_classes) if name == 'VGG13_BN' else\
VGG16(
num_classes=num_classes) if name == 'VGG16' else\
VGG16_BN(
num_classes=num_classes) if name == 'VGG16_BN' else\
VGG19(
num_classes=num_classes) if name == 'VGG19' else\
VGG19_BN(
num_classes=num_classes) if name == 'VGG19_BN' else \
VGGCIFAR('VGG16',
num_classes=num_classes) if name == 'VGG16CIFAR' else \
EfficientNetB4(
num_classes=num_classes) if name == 'EfficientNetB4' else \
EfficientNetB0CIFAR(
num_classes=num_classes) if name == 'EfficientNetB0CIFAR' else\
None
classes = ('plane', 'car', 'bird', 'cat', 'deeer',
'dog', 'frog', 'horse', 'ship', 'truck')
# setup device for training
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# setup Tensorboard file path
writer = SummaryWriter('experiments/teachers/resnet/resnet50')
# Setup best accuracy for comparing and model checkpoints
best_accuracy = 0.0
# Configure the Network
# You can swap out any kind of architectire from /models in here
model_fn = ResNet50()
model_fn = model_fn.to(device)
cudnn.benchmark = True
# print summary of model
summary(model_fn, (3, 32, 32))
# Setup the loss function
criterion = nn.CrossEntropyLoss()
# Setup the optimizer method for all the parameters
optimizer_fn = optim.SGD(model_fn.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
# setup learning rate scheduler
scheduler = MultiStepLR(optimizer_fn, milestones=[150, 225, 270], gamma=0.1)
train_and_evaluate(model=model_fn, train_dataloader=trainloader, test_dataloader=testloader,
# This is not very elegant solution: Please refine it if you can.
trainloader = torch.utils.data.DataLoader(trainset, batch_size=1000,
shuffle=False, num_workers=4)
testset = torchvision.datasets.CIFAR10(root="/home/htut/Desktop/Knowledge_Distillation_Pytorch/datasets", train=False,
download=True, transform=transforms)
testloader = torch.utils.data.DataLoader(testset, batch_size=1000,
shuffle=False, num_workers=4)
# setup device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
weights_path = "checkpoints/teachers/resnet/resnet50_acc:95.370_loss:0.188.pt"
train_logits_path = "logits/resnet/resnet50/train_logits_resnet50.npy"
test_logits_path = "logits/resnet/resnet50/test_logits_resnet50.npy"
model = ResNet50()
model.to(device)
# Load the model
model.load_state_dict(torch.load(weights_path))
# Set the model into evaluation mode
model.eval()
with torch.no_grad():
all_train_logits = []
all_test_logits = []
for train_data in trainloader:
train_images, _ = train_data
train_images = train_images.to(device)
train_logits = model(train_images)
def setup_and_prune(cmd_args,
params,
main_logger,
*,
prune_type="single",
model_name=None):
"""compress a model
cmd_args: result from argparse
params: parameters used to build a pruner
"""
assert prune_type in ["single", "multiple"]
assert model_name in ["vgg16", "resnet50", "resnet56"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_loader, test_loader = get_data_loaders(cmd_args)
if model_name == "vgg16":
model = VGG("VGG16").to(device=device)
elif model_name == "resnet50":
model = ResNet50().to(device=device)
elif model_name == "resnet56":
model = resnet56().to(device=device)
else:
raise ValueError(f"Model {model_name} is wrong.")
main_logger.info("Loading pretrained model {} ...".format(
cmd_args.pretrained))
try:
model.load_state_dict(torch.load(cmd_args.pretrained))
except FileNotFoundError:
print("Pretrained model doesn't exist")
try:
sys.exit(0)
except SystemExit:
os._exit(0)
main_logger.info("Testing pretrained model...")
test(cmd_args, model, device, test_loader, main_logger)
main_logger.info("start model pruning...")
if isinstance(model, nn.DataParallel):
model = model.module
optimizer_finetune = torch.optim.SGD(model.parameters(),
lr=cmd_args.finetune_lr,
momentum=0.9,
weight_decay=1e-4)
best_top1 = 0
if model_name == "vgg16":
if prune_type == "single":
prune_stats = do_prune(model, params)
else:
prune_stats = do_prune_multiple(model, params)
elif model_name == "resnet50":
if prune_type == "multiple":
prune_stats = do_prune_multiple_resnet50(model, params)
else:
raise ValueError(f"prune type {prune_type} is not implemented")
elif model_name == "resnet56":
if prune_type == "multiple":
prune_stats = do_prune_resnet56(model, params)
else:
raise ValueError(f"prune type {prune_type} is not implemented")
if cmd_args.multi_gpu and torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
main_logger.info("Testing pruned model before finetune...")
test(cmd_args, model, device, test_loader, main_logger)
for epoch in range(cmd_args.prune_epochs):
main_logger.info("# Finetune Epoch {} #".format(epoch + 1))
train(
cmd_args,
model,
device,
train_loader,
optimizer_finetune,
epoch,
main_logger,
)
main_logger.info("Testing finetuned model after pruning...")
top1 = test(cmd_args, model, device, test_loader, main_logger)
if top1 > best_top1:
best_top1 = top1
# export finetuned model
info = {}
info["top1"] = best_top1 / 100
info["sparsity"] = prune_stats["sparsity"]
info["value"] = -(best_top1 / 100 + prune_stats["sparsity"])
info["value_sigma"] = 0.25
return info
np.random.shuffle(indices)
train_indices, val_indices = indices[split:], indices[:split]
train_sampler = SubsetRandomSampler(train_indices)
validation_sampler = SubsetRandomSampler(val_indices)
train_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, num_workers=1, sampler=train_sampler)
validation_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, num_workers=1, sampler=validation_sampler)
return train_loader,validation_loader
#Model
if args.dataset=='cifar10':
numberclass=10
if args.dataset=='cifar100':
numberclass=100
net=ResNet50(classnum=args.num_classes)
rand_input=torch.rand(args.batch_size,3,64,64)
writer.add_graph(net,(rand_input,))
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
scheduler=optim.lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.1)
def main(args):
check_path(args)
# CIFAR-10的全部类别，一共10类
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
# 数据集
data_builder = DataBuilder(args)
dataSet = DataSet(data_builder.train_builder(),
data_builder.test_builder(), classes)
# 选择模型
if args.lenet:
net = LeNet()
model_name = args.name_le
elif args.vgg:
net = Vgg16_Net()
model_name = args.name_vgg
elif args.resnet18:
net = ResNet18()
model_name = args.name_res18
elif args.resnet34:
net = ResNet34()
model_name = args.name_res34
elif args.resnet50:
net = ResNet50()
model_name = args.name_res50
elif args.resnet101:
net = ResNet101()
model_name = args.name_res101
elif args.resnet152:
net = ResNet152()
model_name = args.name_res152
# 交叉熵损失函数
criterion = nn.CrossEntropyLoss()
# SGD优化器
optimizer = optim.SGD(net.parameters(),
lr=args.learning_rate,
momentum=args.sgd_momentum,
weight_decay=args.weight_decay)
# 余弦退火调整学习率
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=150)
# 模型的参数保存路径
model_path = os.path.join(args.model_path, model_name)
# 启动训练
if args.do_train:
print("Training...")
trainer = Trainer(net, criterion, optimizer, scheduler,
dataSet.train_loader, dataSet.test_loader,
model_path, args)
trainer.train(epochs=args.epoch)
# t.save(net.state_dict(), model_path)
# 启动测试，如果--do_train也出现，则用刚刚训练的模型进行测试
# 否则就使用已保存的模型进行测试
if args.do_eval:
if not args.do_train and not os.path.exists(model_path):
print(
"Sorry, there's no saved model yet, you need to train first.")
return
# --do_eval
if not args.do_train:
checkpoint = t.load(model_path)
net.load_state_dict(checkpoint['net'])
accuracy = checkpoint['acc']
epoch = checkpoint['epoch']
print("Using saved model, accuracy : %f epoch: %d" %
(accuracy, epoch))
tester = Tester(dataSet.test_loader, net, args)
tester.test()
if args.show_model:
if not os.path.exists(model_path):
print(
"Sorry, there's no saved model yet, you need to train first.")
return
show_model(args)
if args.do_predict:
device = t.device("cuda" if t.cuda.is_available() else "cpu")
checkpoint = t.load(model_path, map_location=device)
net.load_state_dict(checkpoint['net'])
predictor = Predictor(net, classes)
img_path = 'test'
img_name = [os.path.join(img_path, x) for x in os.listdir(img_path)]
for img in img_name:
predictor.predict(img)
def select(self, model, args):
"""
Selector utility to create models from model directory
:param model: which model to select. Currently choices are: (cnn | resnet | preact_resnet | densenet | wresnet)
:return: neural network to be trained
"""
if model == 'cnn':
net = SimpleModel(in_shape=self.in_shape,
activation=args.activation,
num_classes=self.num_classes,
filters=args.filters,
strides=args.strides,
kernel_sizes=args.kernel_sizes,
linear_widths=args.linear_widths,
use_batch_norm=args.use_batch_norm)
else:
assert (args.dataset != 'MNIST' and args.dataset != 'Fashion-MNIST'), \
"Cannot use resnet or densenet for mnist style data"
if model == 'resnet':
assert args.resdepth in [18, 34, 50, 101, 152], \
"Non-standard and unsupported resnet depth ({})".format(args.resdepth)
if args.resdepth == 18:
net = ResNet18(self.num_classes)
elif args.resdepth == 34:
net = ResNet34(self.num_classes)
elif args.resdepth == 50:
net = ResNet50(self.num_classes)
elif args.resdepth == 101:
net = ResNet101(self.num_classes)
else:
net = ResNet152()
elif model == 'densenet':
assert args.resdepth in [121, 161, 169, 201], \
"Non-standard and unsupported densenet depth ({})".format(args.resdepth)
if args.resdepth == 121:
net = DenseNet121(
growth_rate=12, num_classes=self.num_classes
) # NB NOTE: growth rate controls cifar implementation
elif args.resdepth == 161:
net = DenseNet161(growth_rate=12,
num_classes=self.num_classes)
elif args.resdepth == 169:
net = DenseNet169(growth_rate=12,
num_classes=self.num_classes)
else:
net = DenseNet201(growth_rate=12,
num_classes=self.num_classes)
elif model == 'preact_resnet':
assert args.resdepth in [18, 34, 50, 101, 152], \
"Non-standard and unsupported preact resnet depth ({})".format(args.resdepth)
if args.resdepth == 18:
net = PreActResNet18(self.num_classes)
elif args.resdepth == 34:
net = PreActResNet34(self.num_classes)
elif args.resdepth == 50:
net = PreActResNet50(self.num_classes)
elif args.resdepth == 101:
net = PreActResNet101(self.num_classes)
else:
net = PreActResNet152()
elif model == 'wresnet':
assert ((args.resdepth - 4) % 6 == 0), \
"Wideresnet depth of {} not supported, must fulfill: (depth - 4) % 6 = 0".format(args.resdepth)
net = WideResNet(depth=args.resdepth,
num_classes=self.num_classes,
widen_factor=args.widen_factor)
else:
raise NotImplementedError(
'Model {} not supported'.format(model))
return net
# 如果您有两组分别用于训练和验证的文件，
# 则可以使用 tf.placeholder(tf.string) 来表示文件名，
# 并使用适当的文件名初始化迭代器：
filenames = tf.placeholder(tf.string, shape=[None])
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(_parse_function) # Parse the record into tensors.
dataset = dataset.shuffle(buffer_size=10000)
# dataset = dataset.repeat() # Repeat the input indefinitely.
dataset = dataset.batch(batch_size)
iterator = dataset.make_initializable_iterator()
next_img, next_label = iterator.get_next()
print(next_img.get_shape())
print(next_label.get_shape())
model = ResNet50(next_img, class_num)
logits = model.logits
print(logits)
with tf.name_scope('loss'):
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=next_label, logits=logits, name='xentropy')
loss = tf.reduce_mean(cross_entropy, name='train_loss')
l2 = tf.add_n([tf.nn.l2_loss(var) for var in tf.trainable_variables()])
with tf.name_scope('train'):
train_op = tf.train.MomentumOptimizer(learning_rate=learning_rate,
momentum=0.9).minimize(loss=loss +
l2 * 1e-5)
with tf.name_scope('accuracy'):
trainsubset = get_subsample_dataset_label_noise(trainset,
permute_index[trial],
noise_size=args.noise_size)
trainloader = torch.utils.data.DataLoader(trainsubset,
batch_size=128,
shuffle=True)
##########################################
# set up model and optimizer
##########################################
if args.arch == 'resnet18':
net = ResNet18(width=args.width).cuda()
elif args.arch == 'resnet34':
net = ResNet34(width=args.width).cuda()
elif args.arch == 'resnet50':
net = ResNet50(width=args.width).cuda()
elif args.arch == 'resnext':
net = ResNeXt29(width=args.width).cuda()
elif args.arch == 'resnext_1d':
net = ResNeXt29_1d(width=args.width).cuda()
elif args.arch == 'vgg':
net = VGG11(width=args.width).cuda()
elif args.arch == 'resnet26_bottle':
net = ResNet26_bottle(width=args.width).cuda()
elif args.arch == 'resnet38_bottle':
net = ResNet38_bottle(width=args.width).cuda()
elif args.arch == 'resnet50_bottle':
net = ResNet50_bottle(width=args.width).cuda()
else:
print('no available arch')
raise RuntimeError
PGD_attack = PGD_l2(model=model, eps=args.eps * normalization_factor, iters=args.iters,
alpha=args.step_size * normalization_factor)
temp_hist = calculate_thickness(train_loader, model, PGD_attack, num_classes, args)
hist_list_epochs[epoch_num] = temp_hist
else:
model_list = {
'resnet20': resnet20_cifar(),
'resnet32': resnet32_cifar(),
'resnet44': resnet44_cifar(),
'resnet56': resnet56_cifar(),
'resnet110': resnet110_cifar(),
'ResNet18': ResNet18(),
'ResNet50': ResNet50(),
'DenseNet': densenet_cifar(),
'VGG13': VGG('VGG13'),
'VGG19': VGG('VGG19')
}
model = model_list[args.arch]
if args.noise_type == "None":
Noisy_mixup = False
elif args.noise_type == "Noisy":
Noisy_mixup = True
if args.arch=="ResNet18" and Noisy_mixup:
model.linear = nn.Linear(in_features=512, out_features=num_classes+1, bias=True)
if __name__ == '__main__':
IMG_HEIGHT = 224
IMG_WIDTH = 224
IMG_DEPTH = 1
num_class = 11
batch_size = 16
learning_rate = 0.01
graph = tf.Graph()
with graph.as_default():
xs = tf.placeholder("float",
shape=[None, IMG_HEIGHT, IMG_WIDTH, IMG_DEPTH])
ys = tf.placeholder("float", shape=[None, num_class])
model = ResNet50(xs, num_class)
# model = DenseNet_BC(xs,num_class,k=32,layers=[6,12,24,16])
output = model.logits
# `logits` and `labels` must have the same shape,
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=ys,
logits=output)
loss = tf.reduce_mean(cross_entropy, name='loss')
# ==================注意：======不同的网络有不同的训练方法===================
# densenetBC:we use a weight decay of 10−4 and a Nesterov momentum of 0.9 without dampening.
l2 = tf.add_n([tf.nn.l2_loss(var) for var in tf.trainable_variables()])
train_op = tf.train.MomentumOptimizer(
learning_rate=learning_rate, momentum=0.9,
use_nesterov=True).minimize(loss=loss + l2 * 1e-4)
correct_prediction = tf.equal(tf.argmax(output, 1), tf.argmax(ys, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
def __init__(self, base_model, out_dim):
super(ResNetSimCLR, self).__init__()
self.resnet_dict = {"resnet18": ResNet18(num_classes=out_dim),
"resnet50": ResNet50(num_classes=out_dim)}
self.backbone = self._get_basemodel(base_model)
def load_model(model_name, training_type, configs):
"""
Loads model.
"""
dataset = configs.dataset.lower()
# set the input channels and num_classes
if dataset == "mnist" or dataset == "fashionmnist":
configs.num_classes = 10
configs.input_channels = 1
elif dataset == "cifar-100":
configs.num_classes = 100
configs.input_channels = 3
elif dataset == "imagenet":
configs.num_classes = 1000
configs.input_channels = 3
else:
configs.num_classes = 10
configs.input_channels = 3
# pick model
if model_name == "Resnet18":
# load weights
if training_type == "pretrained":
print("Loading pretrained Resnet18")
model = torchvision.models.resnet18(pretrained=True)
model.fc.Linear = nn.Linear(model.fc.in_features,
configs.num_classes)
elif training_type == "untrained":
print("Loading untrained Resnet18")
model = ResNet18(num_classes=configs.num_classes,
input_channels=configs.input_channels)
elif model_name == "Resnet50":
# load weights
if training_type == "pretrained":
print(f"Loading pretrained {model_name}")
model = torchvision.models.resnet50(pretrained=True)
model.fc.Linear = nn.Linear(model.fc.in_features,
configs.num_classes)
elif training_type == "untrained":
print(f"Loading untrained {model_name}")
model = ResNet50(num_classes=configs.num_classes,
input_channels=configs.input_channels)
elif model_name == "Resnet101":
if training_type == 'pretrained':
print(f"Loading pretrained {model_name}")
model = torchvision.models.resnet101(pretrained=True)
elif training_type == "untrained":
print(f"Loading untrained {model_name}")
model = torchvision.models.resnet101()
elif model_name == "VGG19":
# load weights
if training_type == "pretrained":
print(f"Loading pretrained {model_name}")
model = torchvision.models.vgg19(pretrained=True)
model.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 512),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(512, 512),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(512, 512),
nn.ReLU(True),
nn.Linear(512, 10),
)
elif training_type == "untrained":
print(f"Loading untrained {model_name}")
model = vgg19_bn(in_channels=configs.input_channels,
num_classes=configs.num_classes)
elif "efficientnet" in model_name:
if training_type == 'pretrained':
print(f"Loading pretrained {model_name}")
model = load_efficientnet(model_name, configs.num_classes,
configs.input_channels, True)
elif training_type == "untrained":
print(f"Loading untrained {model_name}")
model = load_efficientnet(model_name, configs.num_classes,
configs.input_channels, False)
else:
print("Please provide a model")
# push model to cuda
if torch.cuda.device_count() > 1:
print(f"Number of GPUs available are {torch.cuda.device_count()}")
model = nn.DataParallel(model)
print("\nModel moved to Data Parallel")
model.cuda()
return model
