def main():
args = parse_args()
update_config(cfg, args)
if args.load_path is None:
raise AttributeError("Please specify load path.")
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
# Set the random seed manually for reproducibility.
np.random.seed(cfg.SEED)
torch.manual_seed(cfg.SEED)
torch.cuda.manual_seed_all(cfg.SEED)
# model and optimizer
if cfg.MODEL.NAME == 'model':
if args.load_path and os.path.exists(args.load_path):
checkpoint = torch.load(args.load_path)
genotype = checkpoint['genotype']
else:
raise AssertionError('Please specify the model to evaluate')
model = Network(cfg.MODEL.INIT_CHANNELS, cfg.MODEL.NUM_CLASSES, cfg.MODEL.LAYERS, genotype)
model.drop_path_prob = 0.0
else:
model = eval('resnet.{}(num_classes={})'.format(cfg.MODEL.NAME, cfg.MODEL.NUM_CLASSES))
model = model.cuda()
# resume && make log dir and logger
if args.load_path and os.path.exists(args.load_path):
checkpoint = torch.load(args.load_path)
# load checkpoint
model.load_state_dict(checkpoint['state_dict'])
args.path_helper = checkpoint['path_helper']
logger = create_logger(os.path.dirname(args.load_path))
logger.info("=> loaded checkpoint '{}'".format(args.load_path))
else:
raise AssertionError('Please specify the model to evaluate')
logger.info(args)
logger.info(cfg)
# dataloader
test_dataset_verification = VoxcelebTestset(
Path(cfg.DATASET.DATA_DIR), cfg.DATASET.PARTIAL_N_FRAMES
)
test_loader_verification = torch.utils.data.DataLoader(
dataset=test_dataset_verification,
batch_size=1,
num_workers=cfg.DATASET.NUM_WORKERS,
pin_memory=True,
shuffle=False,
drop_last=False,
)
validate_verification(cfg, model, test_loader_verification)
def start_test():
args = parse_args()
sys.stdout = Logger(os.path.join(args.log_test_dir, "test_log.txt"))
# load GPU
str_ids = args.gpus.split(',')
gpu_ids = []
for str_id in str_ids:
gid = int(str_id)
if gid >= 0:
gpu_ids.append(gid)
# set gpu ids
if len(gpu_ids) > 0:
torch.cuda.set_device(gpu_ids[0])
cudnn.benchmark = True
with open('%s/opts_test.yaml' % args.log_test_dir, 'w') as fp:
yaml.dump(vars(args), fp, default_flow_style=False)
model = Network(args).cuda()
main(model, args)
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.cuda()
utils.load(model, args.model_path)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
_, 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
test_acc, test_obj = infer(test_queue, model, criterion)
logging.info('test_acc %f', test_acc)
def train(args):
"""
Train a model on the train set defined in labels.json
"""
config_path = args.conf
with open(config_path) as f:
config = yaml.load(f, Loader=yaml.FullLoader)
with open(config['labels_file']) as f:
dataset = json.load(f)
# print(dataset)
train_generator = DataGenerator(config, dataset['train'], shuffle=True,
use_data_augmentation=config['data_aug']['use_data_aug'])
#----------val generator--------
val_generator = DataGenerator(config, dataset['val'], shuffle=True, use_data_augmentation=False)
train_model = Network(config)
trainer = Trainer(config, train_model, train_generator, val_generator)
trainer.train()
def start(args):
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cudnn.enabled = True
cudnn.benchmark = True
logging.info("args = %s", args)
dataset = LoadData(args.data_name)
in_dim = dataset.num_atom_type
num_classes = 1
criterion = nn.L1Loss()
criterion = criterion.cuda()
print(f"=> input dimension: {in_dim}, number classes: {num_classes}")
genotype = ZINC_Net
print('=> loading from genotype: \n', genotype)
# model = Network(genotype, args.layers, in_dim, args.feature_dim, num_classes, criterion, args.data_type, args.readout, args.dropout)
model = Network(args, genotype, num_classes, in_dim, criterion)
model = model.cuda()
logging.info("=> param size = %f", count_parameters_in_MB(model) * 1e6)
train_data, val_data, test_data = dataset.train, dataset.val, dataset.test
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
pin_memory=True,
num_workers=args.workers,
collate_fn=dataset.collate,
shuffle=True)
valid_queue = torch.utils.data.DataLoader(val_data,
batch_size=args.batch_size,
pin_memory=True,
num_workers=args.workers,
collate_fn=dataset.collate,
shuffle=False)
test_queue = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
pin_memory=True,
num_workers=args.workers,
collate_fn=dataset.collate,
shuffle=False)
if args.optimizer == 'SGD':
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), eta_min=args.learning_rate_min)
elif args.optimizer == 'ADAM':
optimizer = torch.optim.Adam(model.parameters(),
lr=0.001,
weight_decay=0.0)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='min',
factor=0.5,
patience=10,
verbose=True)
for epoch in range(args.epochs):
logging.info('[EPOCH]\t%d', epoch)
if args.optimizer == 'SGD':
scheduler.step()
lr = scheduler.get_lr()[0]
logging.info('[LR]\t%f', lr)
# training
train_mae, train_obj = train(train_queue, model, criterion, optimizer)
# validation
valid_mae, valid_obj = infer(valid_queue,
model,
criterion,
stage='validating')
# testing
test_mae, test_obj = infer(test_queue,
model,
criterion,
stage='testing ')
desc = '[train] mae: {:.3f}, loss: {:.3f}\t[validate] mae:{:.3f}, loss: {:.3f}\t[test] mae: {:.3f}, loss: {:.3f}'.format(
train_mae, train_obj, valid_mae, valid_obj, test_mae, test_obj)
logging.info(desc)
if args.optimizer == 'ADAM':
scheduler.step(valid_obj)
if optimizer.param_groups[0]['lr'] < 1e-5:
print("\n!! LR SMALLER OR EQUAL TO MIN LR THRESHOLD.")
break
torch.cuda.set_device(gpu_ids[0])
cudnn.benchmark = True
with open('%s/opts_test.yaml' % args.log_test_dir, 'w') as fp:
yaml.dump(vars(args), fp, default_flow_style=False)
model = Network(args).cuda()
main(model, args)
if __name__ == '__main__':
args = parse_args()
sys.stdout = Logger(os.path.join(args.log_test_dir, "test_log.txt"))
# load GPU
str_ids = args.gpus.split(',')
gpu_ids = []
for str_id in str_ids:
gid = int(str_id)
if gid >= 0:
gpu_ids.append(gid)
# set gpu ids
if len(gpu_ids) > 0:
torch.cuda.set_device(gpu_ids[0])
cudnn.benchmark = True
with open('%s/opts_test.yaml' % args.log_test_dir, 'w') as fp:
yaml.dump(vars(args), fp, default_flow_style=False)
model = Network(args).cuda()
main(model, args)
# -*- coding: utf-8 -*-
"""
Created on Fri Sep 25 13:48:22 2020
@author: dilpreet
"""
from models.model import Network
from preprocessing.preprocessing import get_data
if __name__ == "__main__":
dataset, class_dist = get_data()
prob_dist = class_dist / class_dist.sum()
model1 = Network(dataset, class_dist, prob_dist, mode='cnn1d')
model1.fit()
model2 = Network(dataset, class_dist, prob_dist, mode='cnn2d')
model2.fit()
model3 = Network(dataset, class_dist, prob_dist, mode='rnn')
model3.fit()
def main():
args = parse_args()
update_config(cfg, args)
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
# Set the random seed manually for reproducibility.
np.random.seed(cfg.SEED)
torch.manual_seed(cfg.SEED)
torch.cuda.manual_seed_all(cfg.SEED)
# Loss
criterion = CrossEntropyLoss(cfg.MODEL.NUM_CLASSES).cuda()
# load arch
genotype = eval(args.text_arch)
model = Network(cfg.MODEL.INIT_CHANNELS, cfg.MODEL.NUM_CLASSES,
cfg.MODEL.LAYERS, genotype)
model = model.cuda()
optimizer = optim.Adam(
model.parameters(),
lr=cfg.TRAIN.LR,
weight_decay=cfg.TRAIN.WD,
)
# resume && make log dir and logger
if args.load_path and os.path.exists(args.load_path):
checkpoint_file = os.path.join(args.load_path, 'Model',
'checkpoint_best.pth')
assert os.path.exists(checkpoint_file)
checkpoint = torch.load(checkpoint_file)
# load checkpoint
begin_epoch = checkpoint['epoch']
last_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
best_acc1 = checkpoint['best_acc1']
optimizer.load_state_dict(checkpoint['optimizer'])
args.path_helper = checkpoint['path_helper']
logger = create_logger(args.path_helper['log_path'])
logger.info("=> loaded checkloggpoint '{}'".format(checkpoint_file))
else:
exp_name = args.cfg.split('/')[-1].split('.')[0]
args.path_helper = set_path('logs_scratch', exp_name)
logger = create_logger(args.path_helper['log_path'])
begin_epoch = cfg.TRAIN.BEGIN_EPOCH
best_acc1 = 0.0
last_epoch = -1
logger.info(args)
logger.info(cfg)
logger.info(f"selected architecture: {genotype}")
logger.info("Number of parameters: {}".format(count_parameters(model)))
# copy model file
this_dir = os.path.dirname(__file__)
shutil.copy2(os.path.join(this_dir, './models', cfg.MODEL.NAME + '.py'),
args.path_helper['ckpt_path'])
# dataloader
train_dataset = DeepSpeakerDataset(Path(cfg.DATASET.DATA_DIR),
cfg.DATASET.PARTIAL_N_FRAMES, 'train')
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE,
num_workers=cfg.DATASET.NUM_WORKERS,
pin_memory=True,
shuffle=True,
drop_last=True,
)
test_dataset = DeepSpeakerDataset(Path(cfg.DATASET.DATA_DIR),
cfg.DATASET.PARTIAL_N_FRAMES,
'test',
is_test=True)
test_loader = torch.utils.data.DataLoader(
dataset=test_dataset,
batch_size=1,
num_workers=cfg.DATASET.NUM_WORKERS,
pin_memory=True,
shuffle=True,
drop_last=True,
)
# training setting
writer_dict = {
'writer': SummaryWriter(args.path_helper['log_path']),
'train_global_steps': begin_epoch * len(train_loader),
'valid_global_steps': begin_epoch // cfg.VAL_FREQ,
}
# training loop
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
cfg.TRAIN.END_EPOCH,
cfg.TRAIN.LR_MIN,
last_epoch=last_epoch)
for epoch in tqdm(range(begin_epoch, cfg.TRAIN.END_EPOCH),
desc='train progress'):
model.train()
model.drop_path_prob = cfg.MODEL.DROP_PATH_PROB * epoch / cfg.TRAIN.END_EPOCH
train_from_scratch(cfg, model, optimizer, train_loader, criterion,
epoch, writer_dict)
if epoch % cfg.VAL_FREQ == 0 or epoch == cfg.TRAIN.END_EPOCH - 1:
acc = validate_identification(cfg, model, test_loader, criterion)
# remember best acc@1 and save checkpoint
is_best = acc > best_acc1
best_acc1 = max(acc, best_acc1)
# save
logger.info('=> saving checkpoint to {}'.format(
args.path_helper['ckpt_path']))
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
'path_helper': args.path_helper,
'genotype': genotype,
}, is_best, args.path_helper['ckpt_path'],
'checkpoint_{}.pth'.format(epoch))
lr_scheduler.step(epoch)
dataloaders = {
x: data_config(args.dir,
args.batch_size,
x,
args.max_length,
args.embedding_type,
transform=data_transforms[x])
for x in ['train', 'val']
}
# loss function
if args.CMPM:
print("import CMPM")
compute_loss = Loss(args).cuda()
model = Network(args).cuda()
# compute the model size:
print('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
# load checkpoint:
if args.resume is not None:
start_epoch, model = load_checkpoint(model, args.resume)
else:
print("Do not load checkpoint,Epoch start from 0")
start_epoch = 0
# opitimizer:
opitimizer = optimizer_function(args, model)
exp_lr_scheduler = lr_scheduler(opitimizer, args)
format=log_format, datefmt='%m/%d %I:%M:%S %p')
if not os.path.isdir(result_dir):
os.makedirs(result_dir)
fh = logging.FileHandler(os.path.join(result_dir, 'train.log'))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
train_datasets = [None] * len(HEADS)
train_loaders = [None] * len(HEADS)
for i in range(len(HEADS)):
train_datasets[i] = loadedDataset(train_dir, HEADS[i], args.frame_num, args.time_freq, args.patch_size)
train_loaders[i] = torch.utils.data.DataLoader(
train_datasets[i], batch_size=args.batch_size, shuffle=True, pin_memory=True)
model = Network(in_channels=3, out_channels=3).cuda()
rmse_loss = RMSE_Loss()
criterion = nn.MSELoss()
if os.path.exists(os.path.join(save_dir, 'checkpoint.pth.tar')):
# load existing model
print('==> loading existing model')
model_info = torch.load(os.path.join(save_dir, 'checkpoint.pth.tar'))
model.load_state_dict(model_info['state_dict'])
optimizer = torch.optim.Adam(model.parameters())
optimizer.load_state_dict(model_info['optimizer'])
cur_epoch = model_info['epoch']
else:
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
# create model
from models.model import Network
parser = argparse.ArgumentParser(description = 'Train')
# training parameters
parser.add_argument('--origin_size', default=1999, type=int, help='origin image size')
parser.add_argument('--patch_size', default=512, type=int, help='image patch size')
parser.add_argument('--clip_size', default=184, type=int, help='clipped image size')
args = parser.parse_args()
test_dir = './img_data/test/'
save_dir = './save_model/'
result_dir = './result/final/'
HEADS = ['U', 'V', 'W', 'X', 'Y', 'Z']
model = Network(in_channels=3, out_channels=3).cuda()
model.eval()
if os.path.exists(os.path.join(save_dir, 'checkpoint.pth.tar')):
# load existing model
print('==> loading existing model')
g_info = torch.load(os.path.join(save_dir, 'checkpoint.pth.tar'))
model.load_state_dict(g_info['state_dict'])
else:
print('==> No trained model detected!')
exit(1)
if not os.path.isdir(result_dir):
os.makedirs(result_dir)
for head in HEADS:
def start(args):
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
# seed
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cudnn.benchmark = True
cudnn.enabled = True
logging.info("args = %s", args)
dataset = LoadData(args.data_name)
if args.data_name == 'SBM_PATTERN':
in_dim = 3
num_classes = 2
else:
in_dim = 7
num_classes = 6
print(f"=> input dimension: {in_dim}, number classes: {num_classes}")
criterion = MyCriterion(num_classes)
criterion = criterion.cuda()
if args.data_name == 'SBM_PATTERN':
genotype = PATTERN_Net
elif args.data_name == 'SBM_CLUSTER':
genotype = CLUSTER_Net
else:
print("Unknown dataset.")
exit()
print('=> loading from genotype: \n', genotype)
model = Network(args, genotype, num_classes, in_dim, criterion)
model = model.cuda()
logging.info("param size = %fMB", count_parameters_in_MB(model))
train_data, val_data, test_data = dataset.train, dataset.val, dataset.test
train_queue = torch.utils.data.DataLoader(
train_data, batch_size = args.batch_size,
pin_memory = True,
num_workers=args.workers,
collate_fn = dataset.collate,
shuffle = True)
valid_queue = torch.utils.data.DataLoader(
val_data, batch_size = args.batch_size,
pin_memory = True,
num_workers=args.workers,
collate_fn = dataset.collate,
shuffle = False)
test_queue = torch.utils.data.DataLoader(
test_data, batch_size=args.batch_size,
pin_memory=True,
num_workers=args.workers,
collate_fn=dataset.collate,
shuffle = False)
if args.optimizer == 'SGD':
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), eta_min=args.learning_rate_min)
elif args.optimizer == 'ADAM':
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
factor=0.5,
patience=5,
verbose=True)
for epoch in range(args.epochs):
logging.info('[EPOCH]\t%d', epoch)
if args.optimizer == 'SGD':
scheduler.step()
lr = scheduler.get_lr()[0]
logging.info('[LR]\t%f', lr)
macro_acc, micro_acc, train_obj = train(train_queue, model, criterion, optimizer)
# validation
macro_acc, micro_acc, valid_obj = infer(valid_queue, model, criterion, stage = 'validating')
# testing
macro_acc, micro_acc, test_obj = infer(test_queue, model, criterion, stage = ' testing ')
if args.optimizer == 'ADAM':
scheduler.step(valid_obj)
if optimizer.param_groups[0]['lr'] < 1e-5:
print("\n!! LR SMALLER OR EQUAL TO MIN LR THRESHOLD.")
break
def start(args):
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cudnn.enabled = True
cudnn.benchmark = True
logging.info("args = %s", args)
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
dataset = LoadData(args.data_name)
in_dim = dataset.train[0][0].ndata['feat'][0].size(0)
num_classes = len(np.unique(np.array(dataset.train[:][1])))
print(f"=> input dimension: {in_dim}, number classes: {num_classes}")
if args.data_name == 'MNIST':
genotype = MNIST_Net
elif args.data_name == 'CIFAR10':
genotype = CIFAR10_Net
else:
print("Unknown dataset.")
exit()
print('=> loading from genotype: \n', genotype)
model = Network(args, genotype, num_classes, in_dim, criterion)
model = model.cuda()
logging.info("param size = %fMB", count_parameters_in_MB(model))
train_data, val_data, test_data = dataset.train, dataset.val, dataset.test
train_queue = torch.utils.data.DataLoader(
train_data, batch_size = args.batch_size,
pin_memory = True,
num_workers=args.workers,
collate_fn = dataset.collate,
shuffle = True) #新增shuffle
valid_queue = torch.utils.data.DataLoader(
val_data, batch_size = args.batch_size,
pin_memory = True,
num_workers=args.workers,
collate_fn = dataset.collate,
shuffle = False)
test_queue = torch.utils.data.DataLoader(
test_data, batch_size=args.batch_size,
pin_memory=True,
num_workers=args.workers,
collate_fn=dataset.collate,
shuffle = False)
if args.optimizer == 'SGD':
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), eta_min=args.learning_rate_min)
elif args.optimizer == 'ADAM':
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=3e-6)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
factor=0.5,
patience=5,
verbose=True)
for epoch in range(args.epochs):
logging.info('[EPOCH]\t%d', epoch)
if args.optimizer == 'SGD':
scheduler.step()
lr = scheduler.get_lr()[0]
logging.info('[LR]\t%f', lr)
# training
train_acc, train_obj = train(train_queue, model, criterion, optimizer)
# validation
valid_acc, valid_obj = infer(valid_queue, model, criterion, stage = 'validating')
# testing
test_acc, test_obj = infer(test_queue, model, criterion, stage = 'testing ')
desc = '[train] acc: {:.3f}, loss: {:.3f}\t[validate] acc:{:.3f}, loss: {:.3f}\t[test] acc: {:.3f}, loss: {:.3f}'.format(
train_acc, train_obj, valid_acc, valid_obj, test_acc, test_obj
)
logging.info(desc)
if args.optimizer == 'ADAM':
scheduler.step(valid_obj)
if optimizer.param_groups[0]['lr'] < 1e-5:
print("\n!! LR SMALLER OR EQUAL TO MIN LR THRESHOLD.")
break
def main():
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)
logging.info("args = %s", args)
num_gpus = torch.cuda.device_count()
genotype = eval("core.genotypes.%s" % args.arch)
print('---------Genotype---------')
logging.info(genotype)
print('--------------------------')
model = Network(args.init_channels, args.input_channels, num_classes, args.layers, args.auxiliary, genotype)
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.cuda()
else:
model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(num_classes, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(
model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay
)
data_augmentations = args.data_aug
if data_augmentations is None:
data_augmentations = transforms.ToTensor()
elif isinstance(type(data_augmentations), list):
data_augmentations = transforms.Compose(data_augmentations)
elif not isinstance(data_augmentations, transforms.Compose):
raise NotImplementedError
# Dataset
train_data = K49(args.data_dir, True, data_augmentations)
test_data = K49(args.data_dir, False, data_augmentations)
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(
test_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=2)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, float(args.epochs))
best_acc_top1 = 0
for epoch in range(args.epochs):
if args.lr_scheduler == 'cosine':
scheduler.step()
current_lr = scheduler.get_lr()[0]
elif args.lr_scheduler == 'linear':
current_lr = adjust_lr(optimizer, epoch)
else:
print('Wrong lr type, exit')
sys.exit(1)
logging.info('Epoch: %d lr %e', epoch, current_lr)
if epoch < 5 and args.batch_size > 256:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr * (epoch + 1) / 5.0
logging.info('Warming-up Epoch: %d, LR: %e', epoch, current_lr * (epoch + 1) / 5.0)
if num_gpus > 1:
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
else:
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
epoch_start = time.time()
train_acc, train_obj = train(train_queue, model, criterion_smooth, optimizer)
logging.info('Train_acc: %f', train_acc)
valid_acc_top1, valid_obj = infer(valid_queue, model, criterion)
logging.info('Valid_acc_top1: %f', valid_acc_top1)
epoch_duration = time.time() - epoch_start
logging.info('Epoch time: %ds.', epoch_duration)
is_best = False
if valid_acc_top1 > best_acc_top1:
best_acc_top1 = valid_acc_top1
is_best = True
utils.save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer' : optimizer.state_dict(),
}, is_best, log_path)
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("core.genotypes.%s" % args.arch)
#if args.set == "KMNIST":
# model = NetworkKMNIST(args.init_channels, args.input_channels, num_classes, args.layers, args.auxiliary, genotype)
#elif args.set == "K49":
model = Network(args.init_channels, args.input_channels, num_classes,
args.layers, args.auxiliary, genotype)
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)
# Data augmentations
train_transform, valid_transform = utils.data_transforms_Kuzushiji(args)
# Dataset
if args.set == "KMNIST":
train_data = KMNIST(args.data_dir, True, train_transform)
test_data = KMNIST(args.data_dir, False, valid_transform)
elif args.set == "K49":
train_data = K49(args.data_dir, True, train_transform)
test_data = K49(args.data_dir, False, valid_transform)
else:
raise ValueError("Unknown Dataset %s" % args.dataset)
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(test_data,
batch_size=args.batch_size,
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/%d lr %e', epoch, args.epochs,
scheduler.get_lr()[0])
genotype = eval("core.genotypes.%s" % args.arch)
print('---------Genotype---------')
logging.info(genotype)
print('--------------------------')
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)
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(log_path, 'weights.pt'))