def train(model, head, iterator, optimizer, criterion):
epoch_loss = 0
epoch_acc = 0
model.train()
for batch in tqdm(iterator, desc='Training...'):
text = batch['tokens']
text_lengths = batch['tokens_lengths']
label = batch['label']
optimizer.zero_grad()
output, hidden = model(text, text_lengths)
prediction = head(output, hidden)
loss = criterion(prediction, label)
acc = categorical_accuracy(prediction, label)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_acc += acc.item()
return epoch_loss / len(iterator), epoch_acc / len(iterator)
def evaluate(model, head, iterator, criterion):
epoch_loss = 0
epoch_acc = 0
model.eval()
with torch.no_grad():
for batch in tqdm(iterator, desc='Evaluating...'):
text = batch['tokens']
text_lengths = batch['tokens_lengths']
label = batch['label']
output, hidden = model(text, text_lengths)
prediction = head(output, hidden)
loss = criterion(prediction, label)
acc = categorical_accuracy(prediction, label)
epoch_loss += loss.item()
epoch_acc += acc.item()
return epoch_loss / len(iterator), epoch_acc / len(iterator)
def eval_fn(data_loader, model, device):
model.eval()
eval_loss, eval_acc = 0.0, 0.0
fin_targets = []
fin_outputs = []
with torch.no_grad():
for bi, d in tqdm(enumerate(data_loader), total=len(data_loader)):
ids = d["ids"]
token_type_ids = d["token_type_ids"]
mask = d["mask"]
targets = d["targets"]
ids = ids.to(device, dtype=torch.long)
token_type_ids = token_type_ids.to(device, dtype=torch.long)
mask = mask.to(device, dtype=torch.long)
targets = targets.to(device, dtype=torch.long)
outputs = model(
ids=ids,
mask=mask,
token_type_ids=token_type_ids
)
loss = loss_fn(outputs, targets)
eval_loss += loss.item()
pred_labels = torch.argmax(outputs, axis=1)
# (pred_labels == targets).sum().item()
eval_acc += categorical_accuracy(outputs, targets).item()
fin_targets.extend(targets.cpu().detach().numpy().tolist())
fin_outputs.extend(torch.argmax(
outputs, dim=1).cpu().detach().numpy().tolist())
eval_loss /= len(data_loader)
eval_acc /= len(data_loader)
return fin_outputs, fin_targets, eval_loss, eval_acc
def train_fn(data_loader, model, optimizer, device, scheduler):
model.train()
train_loss, train_acc = 0.0, 0.0
for bi, d in tqdm(enumerate(data_loader), total=len(data_loader)):
ids = d["ids"]
token_type_ids = d["token_type_ids"]
mask = d["mask"]
targets = d["targets"]
ids = ids.to(device, dtype=torch.long)
token_type_ids = token_type_ids.to(device, dtype=torch.long)
mask = mask.to(device, dtype=torch.long)
targets = targets.to(device, dtype=torch.long)
optimizer.zero_grad()
outputs = model(
ids=ids,
mask=mask,
token_type_ids=token_type_ids
)
loss = loss_fn(outputs, targets)
loss.backward()
optimizer.step()
scheduler.step()
train_loss += loss.item()
pred_labels = torch.argmax(outputs, dim=1)
# (pred_labels == targets).sum().item()
train_acc += categorical_accuracy(outputs, targets).item()
train_loss /= len(data_loader)
train_acc /= len(data_loader)
return train_loss, train_acc
def test(epoch, model, dataloader):
model.eval()
for i, (X, y) in enumerate(dataloader):
target = utils.one_hot(y, model.final_caps.n_unit)
X, target = Variable(X), Variable(target)
if args.use_gpu:
X, target = X.cuda(), target.cuda()
y_hat = model(X)
loss = model.loss(y_hat, target)
preds = model.capsule_prediction(y_hat)
acc = utils.categorical_accuracy(y.float(), preds.cpu().data)
logger.log(epoch,
i,
len(dataloader.dataset),
start + '_TEST',
loss=loss.data[0],
acc=acc)
def train(epoch, model, dataloader, optim, decoder, decoder_optim):
model.train()
decoder_criterion = nn.MSELoss()
for ix, (X, y) in enumerate(dataloader):
target = utils.one_hot(y, model.final_caps.n_unit)
X, target = Variable(X), Variable(target)
if args.use_gpu:
X, target = X.cuda(), target.cuda()
y_hat = model(X)
loss = model.loss(y_hat, target)
loss.backward()
optim.step()
optim.zero_grad()
# train the decoder
imgs = decoder(y_hat.detach())
decoder_loss = decoder_criterion(imgs, X)
decoder_loss.backward()
decoder_optim.step()
decoder_optim.zero_grad()
if ix % args.log_interval == 0:
preds = model.capsule_prediction(y_hat)
acc = utils.categorical_accuracy(y.float(), preds.cpu().data)
logger.log(epoch,
ix,
len(dataloader.dataset),
start + '_TRAIN',
loss=loss.data[0],
acc=acc,
decoder_loss=decoder_loss.data[0])
logger.images('generated_fmnist', imgs.data.cpu())
return loss.data[0]