def _init_model(self):
genotype = eval('genotypes.%s' % self.args.arch)
model = Network(self.args.init_channels, self.num_classes,
self.args.layers, self.args.auxiliary, genotype,
self.args.parse_method)
flops, params = profile(model,
inputs=(torch.randn(1, 3, 32, 32), ),
verbose=False)
self.logger.info('flops = %fM', flops / 1e6)
self.logger.info('param size = %fM', params / 1e6)
# Try move model to multi gpus
if torch.cuda.device_count() > 1 and self.args.multi_gpus:
self.logger.info('use: %d gpus', torch.cuda.device_count())
model = nn.DataParallel(model)
else:
self.logger.info('gpu device = %d' % self.device_id)
torch.cuda.set_device(self.device_id)
self.model = model.to(self.device)
criterion = nn.CrossEntropyLoss()
self.criterion = criterion.to(self.device)
self.optimizer = torch.optim.SGD(model.parameters(),
self.args.learning_rate,
momentum=self.args.momentum,
weight_decay=self.args.weight_decay)
self.best_acc_top1 = 0
# optionally resume from a checkpoint
if self.args.resume:
if os.path.isfile(self.args.resume):
print("=> loading checkpoint {}".format(self.args.resume))
checkpoint = torch.load(self.args.resume,
map_location=self.device)
self.dur_time = checkpoint['dur_time']
self.args.start_epoch = checkpoint['epoch']
self.best_acc_top1 = checkpoint['best_acc_top1']
self.args.drop_path_prob = checkpoint['drop_path_prob']
self.model.load_state_dict(checkpoint['state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
self.args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(
self.args.resume))
self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
self.optimizer,
float(self.args.epochs),
eta_min=0,
last_epoch=-1
if self.args.start_epoch == 0 else self.args.start_epoch)
# reload the scheduler if possible
if self.args.resume and os.path.isfile(self.args.resume):
checkpoint = torch.load(self.args.resume)
self.scheduler.load_state_dict(checkpoint['scheduler'])
def main():
np.random.seed(args.seed)
if torch.cuda.is_available():
device = torch.device('cuda:{}'.format(args.gpu))
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
cudnn.deterministic = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
else:
device = torch.device('cpu')
logging.info('No gpu device available')
torch.manual_seed(args.seed)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers,
args.auxiliary, genotype)
model = model.to(device)
utils.load(model, args.model_path, args.gpu)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
_, test_transform = utils._data_transforms_cifar10(args.cutout,
args.cutout_length)
test_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=test_transform)
test_queue = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=2)
model.drop_path_prob = args.drop_path_prob
test_acc, test_obj = infer(test_queue, model, criterion, args.gpu)
logging.info('test_acc %f', test_acc)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
# torch.cuda.set_device(args.gpu)
device = torch.device("cuda")
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled=True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
# read data
train_transform, valid_transform = utils._data_transforms_cifar10(args)
if args.dataset == 'cifar10':
args.data = '/home/work/dataset/cifar'
train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform)
valid_data = dset.CIFAR10(root=args.data, train=False, download=True, transform=valid_transform)
classes = 10
if args.dataset == 'cifar100':
args.data = '/home/work/dataset/cifar100'
train_data = dset.CIFAR100(root=args.data, train=True, download=True, transform=train_transform)
valid_data = dset.CIFAR100(root=args.data, train=False, download=True, transform=valid_transform)
classes = 100
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=2)
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=2)
# model
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, classes, args.layers, args.auxiliary, genotype)
model = model.cuda()
model.drop_path_prob = args.drop_path_prob
flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32).cuda(),), verbose=False)
logging.info('flops = %fM', flops / 1e6)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
optimizer = torch.optim.SGD(
model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, float(args.epochs))
best_val_acc = 0.
if args.resume:
# state = torch.load('/home/work/lixudong/code_work/sgas/cnn/full_train_s3_1-20200608/weights.pt')
# state = torch.load('/home/work/lixudong/code_work/sgas/cnn/full_train_s2_factor1-20200609/weights.pt', map_location='cpu')
# state = torch.load('/home/work/lixudong/code_work/sgas/cnn/full_train_s3_factor1-20200609/weights.pt', map_location='cpu')
# state = torch.load('/home/work/lixudong/code_work/sgas/cnn/full_train_s3_0-20200608/weights.pt', map_location='cpu')
# state = torch.load('/home/work/lixudong/code_work/sgas/cnn/full_train_s2_0-20200608/weights.pt', map_location='cpu')
state = torch.load('/home/work/lixudong/code_work/sgas/cnn/full_train_s3_2-20200608/weights.pt', map_location='cpu')
model.load_state_dict(state)
model = model.to(device)
for i in range(args.start_epoch):
scheduler.step()
best_val_acc = 97.19#97.34#97.32#94.92#94.6#97.2
for epoch in range(args.start_epoch, args.epochs):
scheduler.step()
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion, optimizer)
logging.info('train_acc %f', train_acc)
with torch.no_grad():
valid_acc, valid_obj = infer(valid_queue, model, criterion)
if valid_acc > best_val_acc:
best_val_acc = valid_acc
utils.save(model, os.path.join(args.save, 'best_weights.pt'))
# logging.info('valid_acc %f\tbest_val_acc %f', valid_acc, best_val_acc)
logging.info('val_acc: {:.6}, best_val_acc: \033[31m{:.6}\033[0m'.format(valid_acc, best_val_acc))
state = {
'epoch': epoch,
'model_state': model.state_dict(),
'optimizer': optimizer.state_dict(),
'best_val_acc': best_val_acc
}
torch.save(state, os.path.join(args.save, 'weights.pt.tar'))
def main():
np.random.seed(args.seed)
random.seed(args.seed)
if torch.cuda.is_available():
device = torch.device('cuda:{}'.format(args.gpu))
cudnn.benchmark = False
torch.manual_seed(args.seed)
cudnn.enabled = True
cudnn.deterministic = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
else:
device = torch.device('cpu')
logging.info('No gpu device available')
torch.manual_seed(args.seed)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers,
args.auxiliary, genotype)
model = model.to(device)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
total_params = sum(x.data.nelement() for x in model.parameters())
logging.info('Model total parameters: {}'.format(total_params))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
train_transform, valid_transform = utils._data_transforms_cifar10(
args.cutout, args.cutout_length)
train_data = dset.CIFAR10(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=valid_transform)
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=0)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size // 2,
shuffle=False,
pin_memory=True,
num_workers=0)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs))
for epoch in range(args.epochs):
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion, optimizer,
args.gpu)
logging.info('train_acc %f', train_acc)
valid_acc, valid_obj = infer(valid_queue, model, criterion, args.gpu)
logging.info('valid_acc %f', valid_acc)
scheduler.step()
utils.save(model, os.path.join(args.save, 'weights.pt'))
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers,
args.auxiliary, genotype)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs")
#dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
model.to(device)
#model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
train_transform, valid_transform = utils._data_transforms_cifar10(args)
if args.set == 'cifar100':
train_data = dset.CIFAR100(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = dset.CIFAR100(root=args.data,
train=False,
download=True,
transform=valid_transform)
else:
train_data = dset.CIFAR10(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=valid_transform)
#train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform)
#valid_data = dset.CIFAR10(root=args.data, train=False, download=True, transform=valid_transform)
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=2)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=20,
shuffle=False,
pin_memory=True,
num_workers=2)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs))
best_acc = 0.0
for epoch in range(args.epochs):
scheduler.step()
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
#model.drop_path_prob = args.drop_path_prob * (epoch+1) / args.epochs
train_acc, train_obj = train(train_queue, model, criterion, optimizer,
args, epoch)
logging.info('train_acc %f', train_acc)
valid_acc, valid_obj = infer(valid_queue, model, criterion)
if valid_acc > best_acc:
best_acc = valid_acc
logging.info('valid_acc %f, best_acc %f', valid_acc, best_acc)
utils.save(model, os.path.join(args.save, 'weights.pt'))
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled=True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers, args.auxiliary, genotype)
model = model.to('cuda')
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.to('cuda')
# not apply weight decay to BN layers
if args.bn_no_decay:
logging.info('BN layers are excluded from weight decay')
bn_params, other_params = utils.split_bn_params(model)
logging.debug('bn: %s', [p.dtype for p in bn_params])
logging.debug('other: %s', [p.dtype for p in other_params])
param_group = [{'params': bn_params, 'weight_decay': 0},
{'params': other_params}]
else:
param_group = model.parameters()
optimizer = torch.optim.SGD(
param_group,
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay
)
logging.info('optimizer: %s', optimizer)
train_transform, valid_transform = utils.data_transforms_cifar10(args)
train_data = datasets.CIFAR10(root=args.data, train=True, download=True, transform=train_transform)
valid_data = datasets.CIFAR10(root=args.data, train=False, download=True, transform=valid_transform)
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=args.num_workers)
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=args.num_workers)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, float(args.epochs))
init_epoch = 0
best_acc = 0
if args.recover:
states = torch.load(args.recover)
model.load_state_dict(states['stete'])
init_epoch = states['epoch'] + 1
best_acc = states['best_acc']
logging.info('checkpoint loaded')
scheduler.step(init_epoch)
logging.info('scheduler is set to epoch %d. learning rate is %s', init_epoch, scheduler.get_lr())
for epoch in range(init_epoch, args.epochs):
logging.info('epoch %d lr %s', epoch, scheduler.get_lr())
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion, optimizer)
logging.info('train_acc %.4f', train_acc)
with torch.no_grad():
valid_acc, valid_obj = infer(valid_queue, model, criterion)
logging.info('valid_acc %f', valid_acc)
logging.info('epoch %03d overall train_acc=%.4f valid_acc=%.4f', epoch, train_acc, valid_acc)
scheduler.step()
# gpu info
utils.gpu_usage(args.debug)
if valid_acc > best_acc:
best_acc = valid_acc
logging.info('best acc: %.4f', best_acc)
utils.save_checkpoint(state={'stete': model.state_dict(),
'epoch': epoch,
'best_acc': best_acc,},
is_best=False, save=args.save)
def main(self, arch, epochs=600, gpu=0, load_weights=False, train_portion=0.7, seed=0, save='EXP'):
# Set up save file and logging
self.save = 'eval-{}-{}'.format(save, time.strftime("%Y%m%d-%H%M%S"))
utils.create_exp_dir(self.save, scripts_to_save=glob.glob('*.py'))
log_format = '%(asctime)s %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
format=log_format, datefmt='%m/%d %I:%M:%S %p')
fh = logging.FileHandler(os.path.join(self.save, 'log.txt'))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
self.arch = arch
self.epochs = epochs
self.load_weights = load_weights
self.gpu = gpu
self.train_portion = train_portion
if self.train_portion == 1:
self.validation_set = False
self.seed = seed
print('Train class params')
print('arch: {}, epochs: {}, gpu: {}, load_weights: {}, train_portion: {}'
.format(arch, epochs, gpu, load_weights, train_portion))
# cpu-gpu switch
if not torch.cuda.is_available():
logging.info('no gpu device available')
torch.manual_seed(self.seed)
device = torch.device('cpu')
else:
torch.cuda.manual_seed_all(self.seed)
random.seed(self.seed)
torch.manual_seed(self.seed)
device = torch.device(self.gpu)
cudnn.benchmark = False
cudnn.enabled=True
cudnn.deterministic=True
logging.info('gpu device = %d' % self.gpu)
Genotype = namedtuple('Genotype', 'normal normal_concat reduce reduce_concat')
genotype = eval(self.convert_to_genotype(arch))
model = Network(self.init_channels, self.CIFAR_CLASSES, self.layers, self.auxiliary, genotype)
model = model.to(device)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
total_params = sum(x.data.nelement() for x in model.parameters())
logging.info('Model total parameters: {}'.format(total_params))
criterion = nn.CrossEntropyLoss()
criterion = criterion.to(device)
optimizer = torch.optim.SGD(
model.parameters(),
self.learning_rate,
momentum=self.momentum,
weight_decay=self.weight_decay
)
train_transform, test_transform = utils._data_transforms_cifar10(self.cutout, self.cutout_length)
train_data = dset.CIFAR10(root=self.data, train=True, download=True, transform=train_transform)
test_data = dset.CIFAR10(root=self.data, train=False, download=True, transform=test_transform)
num_train = len(train_data)
indices = list(range(num_train))
if self.validation_set:
split = int(np.floor(self.train_portion * num_train))
else:
split = num_train
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=self.batch_size,
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]),
pin_memory=True, num_workers=0)
if self.validation_set:
valid_queue = torch.utils.data.DataLoader(
train_data, batch_size=self.batch_size // 2,
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[split:num_train]),
pin_memory=True, num_workers=0)
test_queue = torch.utils.data.DataLoader(
test_data, batch_size=self.batch_size // 2, shuffle=False, pin_memory=True, num_workers=0)
if self.load_weights:
logging.info('loading saved weights')
ml = 'cuda:{}'.format(self.gpu) if torch.cuda.is_available() else 'cpu'
model.load_state_dict(torch.load('weights.pt', map_location = ml))
logging.info('loaded saved weights')
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, float(self.epochs))
valid_accs = []
test_accs = []
for epoch in range(self.epochs):
scheduler.step()
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
model.drop_path_prob = self.drop_path_prob * epoch / self.epochs
train_acc, train_obj = self.train(train_queue, model, criterion, optimizer)
logging.info('train_acc %f', train_acc)
if self.validation_set:
valid_acc, valid_obj = self.infer(valid_queue, model, criterion)
logging.info('valid_acc %f', valid_acc)
else:
valid_acc, valid_obj = 0, 0
test_acc, test_obj = self.infer(test_queue, model, criterion, test_data=True)
logging.info('test_acc %f', test_acc)
utils.save(model, os.path.join(self.save, 'weights-new.pt'))
if epoch in list(range(max(0, epochs - 5), epochs)):
valid_accs.append((epoch, valid_acc))
test_accs.append((epoch, test_acc))
return valid_accs, test_accs
def main():
#if not torch.cuda.is_available():
# logging.info('no gpu device available')
# sys.exit(1)
np.random.seed(args.seed)
'''
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled=True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
'''
device = torch.device("cpu")
torch.manual_seed(args.seed)
logging.info('use cpu')
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers,
args.auxiliary, genotype)
param_list = list(model.parameters())
# print(model)
# model = model.cuda()
model.to(device)
utils.load(model, args.model_path)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
#criterion = criterion.cuda()
criterion.to(device)
_, test_transform = utils._data_transforms_cifar10(args)
test_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=test_transform)
test_queue = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=2)
model.drop_path_prob = args.drop_path_prob
METRICS = Metrics(list(model.parameters()), p=1)
PERFORMANCE_STATISTICS = {}
ARCH_STATISTICS = {}
metrics_path = './metrics_stat_test.xlsx'
weights_path = './weights_stat_test.xlsx'
epoch = 0
io_metrics = METRICS.evaluate(epoch)
PERFORMANCE_STATISTICS[f'in_S_epoch_{epoch}'] = io_metrics.input_channel_S
PERFORMANCE_STATISTICS[
f'out_S_epoch_{epoch}'] = io_metrics.output_channel_S
PERFORMANCE_STATISTICS[
f'mode12_S_epoch_{epoch}'] = io_metrics.mode_12_channel_S
PERFORMANCE_STATISTICS[f'fc_S_epoch_{epoch}'] = io_metrics.fc_S
PERFORMANCE_STATISTICS[
f'in_rank_epoch_{epoch}'] = io_metrics.input_channel_rank
PERFORMANCE_STATISTICS[
f'out_rank_epoch_{epoch}'] = io_metrics.output_channel_rank
PERFORMANCE_STATISTICS[
f'mode12_rank_epoch_{epoch}'] = io_metrics.mode_12_channel_rank
PERFORMANCE_STATISTICS[f'fc_rank_epoch_{epoch}'] = io_metrics.fc_rank
PERFORMANCE_STATISTICS[
f'in_condition_epoch_{epoch}'] = io_metrics.input_channel_condition
PERFORMANCE_STATISTICS[
f'out_condition_epoch_{epoch}'] = io_metrics.output_channel_condition
PERFORMANCE_STATISTICS[
f'mode12_condition_epoch_{epoch}'] = io_metrics.mode_12_channel_condition
# write metrics data to xls file
metrics_df = pd.DataFrame(data=PERFORMANCE_STATISTICS)
metrics_df.to_excel(metrics_path)
test_acc, test_obj = infer(test_queue, model, criterion)
logging.info('test_acc %f', test_acc)
def main():
tStart = time.time()
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers,
args.auxiliary, genotype)
devices = eval(args.use_gpus)
if not isinstance(devices, tuple):
devices = [devices]
torch.cuda.set_device(devices[0])
setattr(args, 'batch_size', args.batch_size * len(devices))
setattr(args, 'learning_rate', args.learning_rate * len(devices))
logging.info('use gpu: %s' % args.use_gpus)
logging.info("args = %s", args)
model = nn.DataParallel(model, device_ids=devices)
model.to('cuda')
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
train_transform, valid_transform = utils._data_transforms_cifar10(args)
if args.set == 'cifar100':
train_data = dset.CIFAR100(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = dset.CIFAR100(root=args.data,
train=False,
download=True,
transform=valid_transform)
else:
train_data = dset.CIFAR10(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=valid_transform)
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=4)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs))
best_valid_acc = 0.0
best_acc = 0.0
for epoch in range(args.epochs):
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion, optimizer)
logging.info('train_acc %f', train_acc)
scheduler.step()
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
valid_acc, valid_obj = infer(valid_queue, model, criterion)
if valid_acc > best_acc:
best_acc = valid_acc
logging.info('valid_acc %f, best_acc %f', valid_acc, best_acc)
best_valid_acc = max(best_valid_acc, valid_acc)
utils.save(model.module, os.path.join(args.save, 'weights.pt'))
logging.info('best valid_acc %f', best_valid_acc)
tEnd = time.time()
print("Excute time：%.2f second" % (tEnd - tStart))
