def objective(space):
block_kernel1, block_stride1, block_kernel2, block_stride2, kernel_size1, stride1, learning_rate = space
block_kernel1 = int(block_kernel1)
block_stride1 = int(block_stride1)
block_kernel2 = int(block_kernel2)
block_stride2 = int(block_stride2)
kernel_size1 = int(kernel_size1)
stride1 = int(stride1)
learning_rate = float(learning_rate)
block = Block(in_planes=64, out_planes=64, block_kernel1=block_kernel1, block_stride1=block_stride1,
block_kernel2=block_kernel2, block_stride2=block_stride2)
net = MobileNet(block, kernel_size1=kernel_size1, stride1=stride1)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
if use_cuda:
net.cuda()
net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)
total_loss = 0.0
for epoch in range(start_epoch, start_epoch + args.num_epochs):
train(epoch, net, optimizer, criterion)
test_loss = test(epoch, net)
total_loss += test_loss
return total_loss / args.num_epochs
def get_net_model(net='alexnet',
pretrained_dataset='imagenet',
dropout=False,
pretrained=True):
if net == 'alexnet':
model = myalexnet(pretrained=(pretrained_dataset == 'imagenet')
and pretrained,
dropout=dropout)
teacher_model = alexnet(pretrained=(pretrained_dataset == 'imagenet'))
elif net == 'mobilenet-imagenet':
model = MobileNet(num_classes=1001, dropout=dropout)
if pretrained and pretrained_dataset == 'imagenet':
model.load_state_dict(torch.load(imagenet_pretrained_mbnet_path))
teacher_model = MobileNet(num_classes=1001)
if os.path.isfile(imagenet_pretrained_mbnet_path):
teacher_model.load_state_dict(
torch.load(imagenet_pretrained_mbnet_path))
else:
warnings.warn('failed to import teacher model!')
elif net == 'erfnet-cityscapes':
model = erfnet(pretrained=(pretrained_dataset == 'cityscapes')
and pretrained,
num_classes=20,
dropout=dropout)
teacher_model = erfnet(pretrained=(pretrained_dataset == 'cityscapes'),
num_classes=20)
else:
raise NotImplementedError
for p in teacher_model.parameters():
p.requires_grad = False
teacher_model.eval()
return model, teacher_model
def run_training(config,
n_classes,
train_loader,
valid_loader,
width=1,
mb_version=1):
"""
Whole training procedure with fine-tune after regular training
"""
# defining model
if width > 1:
model = tvm.resnet18(num_classes=n_classes)
else:
if mb_version == 1:
model = MobileNet(n_classes=n_classes, width_mult=width)
else:
model = MobileNetV2(n_classes=n_classes, width_mult=width)
model = model.to(config['device'])
# print out number of parameters
num_params = 0
for p in model.parameters():
num_params += np.prod(p.size())
print(f"width={width}, num_params {num_params}")
# defining loss criterion, optimizer and learning rate scheduler
criterion = t.nn.CrossEntropyLoss()
opt = t.optim.Adam(model.parameters(), config['lr'])
sched = t.optim.lr_scheduler.MultiStepLR(opt, [3, 6])
# training process with Adam
tr_loss, tr_accuracy, valid_loss, valid_accuracy = train(
config, model, train_loader, valid_loader, criterion, opt, sched)
# training process with SGDR
opt = t.optim.SGD(model.parameters(), config['lr'] / 10, momentum=0.9)
sched = SGDR(opt, 3, 1.2)
tr_loss_finetune, tr_accuracy_finetune, valid_loss_finetune, valid_accuracy_finetune = train(
config, model, train_loader, valid_loader, criterion, opt, sched)
return [
tr_loss + tr_loss_finetune, tr_accuracy + tr_accuracy_finetune,
valid_loss + valid_loss_finetune,
valid_accuracy + valid_accuracy_finetune
]
class KeypointModel(BasicModule):
def __init__(self, opt):
super(KeypointModel, self).__init__(opt)
self.pretrained = MobileNet()
self.trf = nn.Sequential(nn.Conv2d(256, 256, 3, 1, 1),
nn.BatchNorm2d(256), nn.ReLU(True),
nn.Conv2d(256, 128, 3, 1, 1),
nn.BatchNorm2d(128), nn.ReLU(True))
# self.ReturnType = namedtuple('ReturnType',['out1','out2','out3','out4','out5','out6'])
stages = [Stage(128)] + [Stage(169) for _ in range(2, 7)]
self.stages = nn.ModuleList(stages)
def forward(self, img):
img = self.pretrained(img)
#if self.optimizer.param_groups[0]['lr'] == 0:
# img = img.detach()
features = self.trf(img)
output = self.stages[0](features)
outputs = [output]
for ii in range(1, 6):
stage = self.stages[ii]
input = t.cat([features, output], dim=1)
output = stage(input)
outputs.append(output)
return outputs
def get_optimizer(self, lr1, lr2):
param_groups = [{
'params': self.pretrained.parameters(),
'lr': lr1
}, {
'params': self.stages.parameters(),
'lr': lr2
}, {
'params': self.trf.parameters(),
'lr': lr2
}]
self.optimizer = t.optim.Adam(param_groups)
return self.optimizer
def main():
global opt, start_epoch, best_prec1
opt = cfg
opt.gpuids = list(map(int, opt.gpuids))
if opt.cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
model = MobileNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=opt.lr,
momentum=opt.momentum,
weight_decay=opt.weight_decay,
nesterov=True)
start_epoch = 0
ckpt_file = join("model", opt.ckpt)
if opt.cuda:
torch.cuda.set_device(opt.gpuids[0])
with torch.cuda.device(opt.gpuids[0]):
model = model.cuda()
criterion = criterion.cuda()
model = nn.DataParallel(model,
device_ids=opt.gpuids,
output_device=opt.gpuids[0])
cudnn.benchmark = True
# for resuming training
if opt.resume:
if isfile(ckpt_file):
print("==> Loading Checkpoint '{}'".format(opt.ckpt))
if opt.cuda:
checkpoint = torch.load(ckpt_file,
map_location=lambda storage, loc:
storage.cuda(opt.gpuids[0]))
try:
model.module.load_state_dict(checkpoint['model'])
except:
model.load_state_dict(checkpoint['model'])
else:
checkpoint = torch.load(
ckpt_file, map_location=lambda storage, loc: storage)
try:
model.load_state_dict(checkpoint['model'])
except:
# create new OrderedDict that does not contain `module.`
new_state_dict = OrderedDict()
for k, v in checkpoint['model'].items():
if k[:7] == 'module.':
name = k[7:] # remove `module.`
else:
name = k[:]
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
start_epoch = checkpoint['epoch']
optimizer.load_state_dict(checkpoint['optimizer'])
print("==> Loaded Checkpoint '{}' (epoch {})".format(
opt.ckpt, start_epoch))
else:
print("==> no checkpoint found at '{}'".format(opt.ckpt))
return
# Download & Load Dataset
print('==> Preparing data..')
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_val = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
trainset = torchvision.datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=transform_train)
train_loader = torch.utils.data.DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.workers)
valset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform_val)
val_loader = torch.utils.data.DataLoader(valset,
batch_size=opt.test_batch_size,
shuffle=False,
num_workers=opt.workers)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
# for evaluation
if opt.eval:
if isfile(ckpt_file):
print("==> Loading Checkpoint '{}'".format(opt.ckpt))
if opt.cuda:
checkpoint = torch.load(ckpt_file,
map_location=lambda storage, loc:
storage.cuda(opt.gpuids[0]))
try:
model.module.load_state_dict(checkpoint['model'])
except:
model.load_state_dict(checkpoint['model'])
else:
checkpoint = torch.load(
ckpt_file, map_location=lambda storage, loc: storage)
try:
model.load_state_dict(checkpoint['model'])
except:
# create new OrderedDict that does not contain `module.`
new_state_dict = OrderedDict()
for k, v in checkpoint['model'].items():
if k[:7] == 'module.':
name = k[7:] # remove `module.`
else:
name = k[:]
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
start_epoch = checkpoint['epoch']
optimizer.load_state_dict(checkpoint['optimizer'])
print("==> Loaded Checkpoint '{}' (epoch {})".format(
opt.ckpt, start_epoch))
# evaluate on validation set
print("\n===> [ Evaluation ]")
start_time = time.time()
prec1 = validate(val_loader, model, criterion)
elapsed_time = time.time() - start_time
print("====> {:.2f} seconds to evaluate this model\n".format(
elapsed_time))
return
else:
print("==> no checkpoint found at '{}'".format(opt.ckpt))
return
# train...
train_time = 0.0
validate_time = 0.0
for epoch in range(start_epoch, opt.epochs):
adjust_learning_rate(optimizer, epoch)
print('\n==> Epoch: {}, lr = {}'.format(
epoch, optimizer.param_groups[0]["lr"]))
# train for one epoch
print("===> [ Training ]")
start_time = time.time()
train(train_loader, model, criterion, optimizer, epoch)
elapsed_time = time.time() - start_time
train_time += elapsed_time
print(
"====> {:.2f} seconds to train this epoch\n".format(elapsed_time))
# evaluate on validation set
print("===> [ Validation ]")
start_time = time.time()
prec1 = validate(val_loader, model, criterion)
elapsed_time = time.time() - start_time
validate_time += elapsed_time
print("====> {:.2f} seconds to validate this epoch\n".format(
elapsed_time))
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
state = {
'epoch': epoch + 1,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_model(state, epoch, is_best)
avg_train_time = train_time / opt.epochs
avg_valid_time = validate_time / opt.epochs
total_train_time = train_time + validate_time
print("====> average training time per epoch: {}m {:.2f}s".format(
int(avg_train_time // 60), avg_train_time % 60))
print("====> average validation time per epoch: {}m {:.2f}s".format(
int(avg_valid_time // 60), avg_valid_time % 60))
print("====> training time: {}m {:.2f}s".format(int(train_time // 60),
train_time % 60))
print("====> validation time: {}m {:.2f}s".format(int(validate_time // 60),
validate_time % 60))
print("====> total training time: {}m {:.2f}s".format(
int(total_train_time // 60), total_train_time % 60))
