def main():
config = get_config()
set_seed(config.train.seed)
message_info(config)
# create log path
val_logFile, writer_logFile, save_path = create_logFile(config)
# dist --init
if config.dist:
dist.init_process_group(backend=config.dist_backend,
init_method=config.dist_init_method)
torch.cuda.set_device(config.dist_local_rank)
# tensorboard
if get_rank() == 0:
writer = get_tensorboard_writer(writer_logFile, purge_step=None)
else:
writer = None
# model
model = get_model(config)
# optimizer
optimizer = get_optimizer(config, model)
if config.apex:
model, optimizer = apex.amp.initialize(model,
optimizer,
opt_level=config.apex_mode)
if config.dist:
if config.dist_sync_bn:
if config.apex:
model = convert_syncbn_model(model)
else:
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
if config.apex:
model = DDP(model, delay_allreduce=True)
else:
model = nn.parallel.DistributedDataParallel(model,
device_ids=[config.dist_local_rank],
output_device=config.dist_local_rank)
# loss
train_loss, val_loss = get_loss(config)
# load_data
data = pd.read_csv(config.train.dataset)
skf = KFold(n_splits=10, shuffle=True, random_state=452)
for fold_idx, (train_idx, val_idx) in enumerate(
skf.split(data['filename'].values, data['filename'].values)):
if fold_idx == config.train.fold:
break
# create dataloader
train_data = data.iloc[train_idx]
val_data = data.iloc[val_idx]
train_loader = create_dataloader(config, train_data, 'train')
val_loader = create_dataloader(config, val_data, 'val')
if get_rank() == 0:
logger.info(f"Splited train set: {train_data.shape}")
logger.info(f"Splited val set: {val_data.shape}")
# scheduler
scheduler = CosineWarmupLr(config, optimizer, len(train_loader))
best_precision, lowest_loss = 0, 100
for epoch in range(config.train.epoches):
# if config.dist:
# train_loader.sampler.set_epoch(epoch)
# train
train(config, epoch, train_loader, model, optimizer, scheduler, train_loss, writer)
# val
if epoch % config.train.val_preiod == 0:
precision, avg_loss = val(config, val_loader, model, val_loss, writer)
if get_rank() == 0:
with open(val_logFile, 'a') as acc_file:
acc_file.write(
f'Fold: {fold_idx:2d}, '
f'Epoch: {epoch:2d}, '
f'Precision: {precision:.8f}, '
f'Loss: {avg_loss:.8f}\n')
is_best = precision > best_precision
is_lowest_loss = avg_loss < lowest_loss
best_precision = max(precision, best_precision)
lowest_loss = min(avg_loss, lowest_loss)
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'best_precision': best_precision,
'lowest_loss': lowest_loss,
}
save_checkpoint(state, epoch, is_best, is_lowest_loss, save_path)
torch.save(model,
os.path.join(save_path, 'whole_model.pth'))
if get_rank() == 0:
writer.close()
torch.cuda.empty_cache()
# save
submission = pd.DataFrame({"ID": ids, "Label": prediction})
submission.to_csv('../' + 'svm_submission.csv', index=False, header=False)
if __name__ == "__main__":
config = get_config('../configs/c15.yml')
# #构建保存特征的文件夹
feature_path = '../features/'
os.makedirs(feature_path, exist_ok=True)
# data -train
train_data = pd.read_csv(config.train.dataset)
train_loader = create_dataloader(config, train_data, 'train')
# data -test
test_data = pd.read_csv(config.test.dataset)
test_loader = create_dataloader(config, test_data, 'test')
#############################################################
#### 保存训练集特征
train_feature_path = feature_path + 'psdufeature_aug.pkl'
train_label_path = feature_path + 'psdulabel.pkl'
save_path = feature_path + 'psdusvm.m'
test_feature_path = feature_path + 'testfeature_old.pkl'
test_id_path = feature_path + 'testid_old.pkl'
# save_feature(train_loader, train_feature_path, train_label_path)
# save_feature(test_loader, test_feature_path, test_id_path)
if args.image:
logger.info('test mode: single image')
result = single_image(args.image, models)
logger.info('result:', result)
print('result:', result, 'label:', config.labels_list[int(result)])
sys.exit(0)
else:
logger.info('single image: None')
if config.test.dataset:
logger.info('test mode: csv')
# skf = KFold(n_splits=0)
test_data = pd.read_csv(config.test.dataset)
logger.info(f"test set: {test_data.shape}")
test_loader = create_dataloader(config, test_data, 'test')
result = test_csv(test_loader, models, config.labels_list, half=False)
if not os.path.exists(config.test.log_file):
with open(config.test.log_file, 'w') as f:
pass
with open(config.test.log_file, 'a') as fc:
fc.write(
'\n%s %s\n' %
(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(
time.time())), f'{config.model}'))
fc.write(result)
print(result)