def GraSi3N4_init(r_seed, model, config=None):
np.set_printoptions(linewidth=np.inf)
random.seed(r_seed - 1)
np.random.seed(r_seed)
if config is None:
config = DefaultConfig(42)
else:
config = config
if config.model != model:
print("\n!!!GraSi3N4_init model CHANGE!!! {}=>{}\n".format(
config.model, model))
config.model = model
t0 = time.time()
n_dict = N_Dict(config)
n_dict.Load("Si3N4", "./data/Si3N4_310nm-14280nm.txt", scale=1000)
n_dict.Load("Graphene", "./data/Graphene_240nm-30000nm.txt")
n_dict.InitMap2(["Si3N4", "Graphene"], config.lenda_tic)
config.n_dict = n_dict
return config
def main(**kwargs):
args = DefaultConfig()
args.parse(kwargs)
args.model = 'LSTM'
args.device = 0
args.id = 'word4'
if not torch.cuda.is_available():
args.cuda = False
args.device = None
torch.manual_seed(args.seed) # set random seed for cpu
train_iter, val_iter, test_iter, args.vocab_size, vectors = data.load_data(args)
args.print_config()
global best_score
# init model
model = getattr(models, args.model)(args, vectors)
print(model)
# 模型保存位置
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
save_path = os.path.join(args.save_dir, '{}_{}.pth'.format(args.model, args.id))
if args.cuda:
torch.cuda.set_device(args.device)
torch.cuda.manual_seed(args.seed) # set random seed for gpu
model.cuda()
# 目标函数和优化器
criterion = F.cross_entropy
lr1, lr2 = args.lr1, args.lr2
optimizer = model.get_optimizer(lr1, lr2, args.weight_decay)
for i in range(args.max_epochs):
total_loss = 0.0
correct = 0
total = 0
model.train()
for idx, batch in enumerate(train_iter):
# 训练模型参数
# 使用BatchNorm层时，batch size不能为1
if len(batch) == 1:
continue
text, label = batch.text, batch.label
if args.cuda:
text, label = text.cuda(), label.cuda()
optimizer.zero_grad()
pred = model(text)
loss = criterion(pred, label)
loss.backward()
optimizer.step()
# 更新统计指标
total_loss += loss.item()
predicted = pred.max(1)[1]
total += label.size(0)
correct += predicted.eq(label).sum().item()
if idx % 80 == 79:
print('[{}, {}] loss: {:.3f} | Acc: {:.3f}%({}/{})'.format(i + 1, idx + 1, total_loss / 20,
100. * correct / total, correct, total))
total_loss = 0.0
# 计算再验证集上的分数，并相应调整学习率
f1score = val(model, val_iter, args)
if f1score > best_score:
best_score = f1score
checkpoint = {
'state_dict': model.state_dict(),
'config': args
}
torch.save(checkpoint, save_path)
print('Best tmp model f1score: {}'.format(best_score))
if f1score < best_score:
model.load_state_dict(torch.load(save_path, map_location={'cuda:5':'cuda:0'})['state_dict'])
lr1 *= args.lr_decay
lr2 = 2e-4 if lr2 == 0 else lr2 * 0.8
optimizer = model.get_optimizer(lr1, lr2, 0)
print('* load previous best model: {}'.format(best_score))
print('* model lr:{} emb lr:{}'.format(lr1, lr2))
if lr1 < args.min_lr:
print('* training over, best f1 score: {}'.format(best_score))
break
# 保存训练最终的模型
args.best_score = best_score
final_model = {
'state_dict': model.state_dict(),
'config': args
}
best_model_path = os.path.join(args.save_dir, '{}_{}_{}.pth'.format(args.model, args.text_type, best_score))
torch.save(final_model, best_model_path)
print('Best Final Model saved in {}'.format(best_model_path))
# 在测试集上运行模型并生成概率结果和提交结果
if not os.path.exists('result/'):
os.mkdir('result/')
probs, test_pred = test(model, test_iter, args)
result_path = 'result/' + '{}_{}_{}'.format(args.model, args.id, args.best_score)
np.save('{}.npy'.format(result_path), probs)
print('Prob result {}.npy saved!'.format(result_path))
test_pred[['id', 'class']].to_csv('{}.csv'.format(result_path), index=None)
print('Result {}.csv saved!'.format(result_path))
t2 = time.time()
print('time use: {}'.format(t2 - t1))
