def train():
"""
A simple Neural Network
:return:
"""
logging.info("Start Training!")
corpus_type = 'train'
summary = Summary(args.batch_size, args.max_len)
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(summary.parameters(), lr=args.lr)
summary.train()
start_epoch = 0
if args.model_name:
checkpoint = torch.load(args.load_model)
summary = checkpoint['model']
start_epoch = checkpoint['epochs']
start_epoch += 1 if start_epoch != 0 else start_epoch
for epoch in range(start_epoch, args.epoch):
epoch_loss = 0
batch_num = 0
for i, batch in enumerate(bachify_data(corpus_type)):
batch_df, batch_label, _, _ = batch
batch_df = torch.tensor(batch_df)
batch_label = torch.tensor(batch_label)
binary_output = summary(batch_df)
# calculate loss
loss = criterion(binary_output, batch_label)
loss.backward()
optimizer.step()
epoch_loss += loss
batch_num += 1
logging.info("Epoch {}: Total loss is {}, Avg loss is {}".format(epoch, epoch_loss, epoch_loss/batch_num))
# store model
model_name = "{}_epoch_model.tar".format(epoch)
directory = os.path.join(args.save_path, model_name)
if not os.path.exists(directory):
os.makedirs(directory)
torch.save({
'model': summary.state_dict(),
'loss': epoch_loss / batch_num,
"epochs": epoch
}, directory)
logging.info("Finish Training!")
