def test(model, device, test_loader, criterion, L1_loss_enable=False):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
#test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
test_loss += criterion(output, target).item() # sum up batch loss
pred = output.argmax(
dim=1,
keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
#test_loss /= len(test_loader.dataset) # For F.nll_loss
test_loss /= len(test_loader) # criterion = nn.CrossEntropyLoss()
if (L1_loss_enable == True):
regloss = regularization.L1_Loss_calc(model, 0.0005)
regloss /= len(
test_loader.dataset
) # by batch size which is here total test dataset size
test_loss += regloss
print(
'\nTest set: Average loss: {:.6f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
acc = 100. * correct / len(test_loader.dataset)
return np.round(acc, 2), test_loss
def train(model,
device,
train_loader,
optimizer,
scheduler,
criterion,
L1_loss_enable=False):
model.train()
pbar = tqdm(train_loader)
train_loss = 0
correct = 0
processed = 0
for batch_idx, (data, target) in enumerate(pbar):
# get samples
data, target = data.to(device), target.to(device)
# Init
optimizer.zero_grad()
y_pred = model(data)
loss = criterion(y_pred, target)
if (L1_loss_enable == True):
regloss = regularization.L1_Loss_calc(model, 0.0005)
regloss /= len(data) # by batch size
loss += regloss
train_loss += loss.item()
# Backpropagation
loss.backward()
optimizer.step()
scheduler.step()
# Update pbar-tqdm
pred = y_pred.argmax(
dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
processed += len(data)
pbar.set_description(
desc=
f'Loss={train_loss / (batch_idx + 1):0.6f} Batch_id={batch_idx} Accuracy={100 * correct / processed:0.2f}'
)
train_loss /= len(train_loader)
acc = 100. * correct / len(train_loader.dataset) # processed #
return np.round(acc, 2), np.round(train_loss, 6)
