x = self.down_sample(x, filters=128)
dilation_rate = [1, 2, 5, 9, 2, 5, 9, 17]
for dilation in dilation_rate:
x = self.ss_bt(x, dilation=dilation)
return x
def decoder(self, x):
x = self.apn_module(x)
x = layers.UpSampling2D(size=8, interpolation='bilinear')(x)
x = layers.Conv2D(self.classes, kernel_size=3, padding='same')(x)
x = layers.BatchNormalization()(x)
x = layers.Activation('softmax')(x)
return x
def model(self):
inputs = layers.Input(shape=self.input_shape)
encoder_out = self.encoder(inputs)
output = self.decoder(encoder_out)
return models.Model(inputs, output)
if __name__ == '__main__':
from flops import get_flops
model = LEDNet(2, 3, (256, 256, 3)).model()
model.summary()
get_flops(model)
def main():
args = get_args()
# Log
log_format = '[%(asctime)s] %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
format=log_format, datefmt='%d %I:%M:%S')
t = time.time()
local_time = time.localtime(t)
if not os.path.exists('./log'):
os.mkdir('./log')
fh = logging.FileHandler(os.path.join('log/train-{}{:02}{}'.format(local_time.tm_year % 2000, local_time.tm_mon, t)))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
use_gpu = False
if torch.cuda.is_available():
use_gpu = True
assert os.path.exists(args.train_dir)
train_dataset = datasets.ImageFolder(
args.train_dir,
transforms.Compose([
transforms.RandomResizedCrop(args.im_size),
transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
transforms.RandomHorizontalFlip(0.5),
ToBGRTensor(),
])
)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=8, pin_memory=use_gpu)
train_dataprovider = DataIterator(train_loader)
assert os.path.exists(args.val_dir)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(args.val_dir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(args.im_size),
ToBGRTensor(),
])),
batch_size=200, shuffle=False,
num_workers=8, pin_memory=use_gpu
)
val_dataprovider = DataIterator(val_loader)
print('load data successfully')
arch_path='arch.pkl'
if os.path.exists(arch_path):
with open(arch_path,'rb') as f:
architecture=pickle.load(f)
else:
raise NotImplementedError
channels_scales = (1.0,)*20
model = ShuffleNetV2_OneShot(architecture=architecture, channels_scales=channels_scales, n_class=args.num_classes, input_size=args.im_size)
print('flops:',get_flops(model))
optimizer = torch.optim.SGD(get_parameters(model),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion_smooth = CrossEntropyLabelSmooth(args.num_classes, 0.1)
if use_gpu:
# model = nn.DataParallel(model)
loss_function = criterion_smooth.cuda()
device = torch.device("cuda")
else:
loss_function = criterion_smooth
device = torch.device("cpu")
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer,
lambda step : (1.0-step/args.total_iters) if step <= args.total_iters else 0, last_epoch=-1)
# model = model.to(device)
model = model.cuda()
all_iters = 0
if args.auto_continue:
lastest_model, iters = get_lastest_model()
if lastest_model is not None:
all_iters = iters
checkpoint = torch.load(lastest_model, map_location=None if use_gpu else 'cpu')
model.load_state_dict(checkpoint['state_dict'], strict=True)
print('load from checkpoint')
for i in range(iters):
scheduler.step()
args.optimizer = optimizer
args.loss_function = loss_function
args.scheduler = scheduler
args.train_dataprovider = train_dataprovider
args.val_dataprovider = val_dataprovider
if args.eval:
if args.eval_resume is not None:
checkpoint = torch.load(args.eval_resume, map_location=None if use_gpu else 'cpu')
model.load_state_dict(checkpoint, strict=True)
validate(model, device, args, all_iters=all_iters)
exit(0)
t = time.time()
while all_iters < args.total_iters:
all_iters = train(model, device, args, val_interval=args.val_interval, bn_process=False, all_iters=all_iters)
validate(model, device, args, all_iters=all_iters)
# all_iters = train(model, device, args, val_interval=int(1280000/args.batch_size), bn_process=True, all_iters=all_iters)
validate(model, device, args, all_iters=all_iters)
save_checkpoint({'state_dict': model.state_dict(),}, args.total_iters, tag='bnps-')
print("Finished {} iters in {:.3f} seconds".format(all_iters, time.time()-t))
o = UpSampling2D((2, 2))(o)
o = concatenate([o, f3], axis=-1)
o = Conv2D(64, (3, 3), padding='same')(o)
o = BatchNormalization()(o)
o = UpSampling2D((2, 2))(o)
o = concatenate([o, f2], axis=-1)
o = Conv2D(32, (3, 3), padding='same')(o)
o = BatchNormalization()(o)
o = UpSampling2D((2, 2))(o)
o = concatenate([o, f1], axis=-1)
o = Conv2D(16, (3, 3), padding='same')(o)
o = BatchNormalization()(o)
o = Conv2D(cls_num, (3, 3), padding='same')(o)
o = UpSampling2D((2, 2))(o)
o = Activation('softmax')(o)
return Model(inputs, o)
if __name__ == '__main__':
from flops import get_flops
model = MobileNet(input_shape=(256, 256, 3), cls_num=3)
model.summary()
get_flops(model, True)
def main():
args = get_args()
# Log
log_format = '[%(asctime)s] %(message)s'
logging.basicConfig(stream=sys.stdout,
level=logging.INFO,
format=log_format,
datefmt='%d %I:%M:%S')
t = time.time()
local_time = time.localtime(t)
if not os.path.exists('./log'):
os.mkdir('./log')
fh = logging.FileHandler(
os.path.join('log/train-{}{:02}{}'.format(local_time.tm_year % 2000,
local_time.tm_mon, t)))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
use_gpu = False
if torch.cuda.is_available():
use_gpu = True
assert os.path.exists(args.train_dir)
train_dataset = datasets.ImageFolder(
args.train_dir,
transforms.Compose([
transforms.RandomResizedCrop(96),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4),
transforms.RandomHorizontalFlip(0.5),
ToBGRTensor(),
]))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
pin_memory=use_gpu)
train_dataprovider = DataIterator(train_loader)
assert os.path.exists(args.val_dir)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(
args.val_dir,
transforms.Compose([
OpencvResize(96),
# transforms.CenterCrop(96),
ToBGRTensor(),
])),
batch_size=200,
shuffle=False,
num_workers=4,
pin_memory=use_gpu)
val_dataprovider = DataIterator(val_loader)
arch_path = 'cl400.p'
if os.path.exists(arch_path):
with open(arch_path, 'rb') as f:
architectures = pickle.load(f)
else:
raise NotImplementedError
channels_scales = (1.0, ) * 20
cands = {}
splits = [(i, 10 + i) for i in range(0, 400, 10)]
architectures = np.array(architectures)
architectures = architectures[
splits[args.split_num][0]:splits[args.split_num][1]]
print(len(architectures))
logging.info("Training and Validating arch: " +
str(splits[args.split_num]))
for architecture in architectures:
architecture = tuple(architecture.tolist())
model = ShuffleNetV2_OneShot(architecture=architecture,
channels_scales=channels_scales,
n_class=10,
input_size=96)
print('flops:', get_flops(model))
optimizer = torch.optim.SGD(get_parameters(model),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion_smooth = CrossEntropyLabelSmooth(1000, 0.1)
if use_gpu:
model = nn.DataParallel(model)
loss_function = criterion_smooth.cuda()
device = torch.device("cuda")
else:
loss_function = criterion_smooth
device = torch.device("cpu")
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lambda step: (1.0 - step / args.total_iters)
if step <= args.total_iters else 0,
last_epoch=-1)
model = model.to(device)
all_iters = 0
if args.auto_continue:
lastest_model, iters = get_lastest_model()
if lastest_model is not None:
all_iters = iters
checkpoint = torch.load(
lastest_model, map_location=None if use_gpu else 'cpu')
model.load_state_dict(checkpoint['state_dict'], strict=True)
print('load from checkpoint')
for i in range(iters):
scheduler.step()
args.optimizer = optimizer
args.loss_function = loss_function
args.scheduler = scheduler
args.train_dataprovider = train_dataprovider
args.val_dataprovider = val_dataprovider
# print("BEGIN VALDATE: ", args.eval, args.eval_resume)
if args.eval:
if args.eval_resume is not None:
checkpoint = torch.load(
args.eval_resume, map_location=None if use_gpu else 'cpu')
model.load_state_dict(checkpoint, strict=True)
validate(model, device, args, all_iters=all_iters)
exit(0)
# t1,t5 = validate(model, device, args, all_iters=all_iters)
# print("VALDATE: ", t1, " ", t5)
while all_iters < args.total_iters:
all_iters = train(model,
device,
args,
val_interval=args.val_interval,
bn_process=False,
all_iters=all_iters)
validate(model, device, args, all_iters=all_iters)
all_iters = train(model,
device,
args,
val_interval=int(1280000 / args.batch_size),
bn_process=True,
all_iters=all_iters)
top1, top5 = validate(model, device, args, all_iters=all_iters)
save_checkpoint({
'state_dict': model.state_dict(),
},
args.total_iters,
tag='bnps-')
cands[architecture] = [top1, top5]
pickle.dump(
cands,
open("from_scratch_split_{}.pkl".format(args.split_num), 'wb'))
