SV12/ERA_Session13

1from tqdm import tqdm2import torch3import torch.nn.functional as F4 5 6def train(7    model,8    device,9    train_loader,10    optimizer,11    criterion,12    scheduler,13    L1=False,14    l1_lambda=0.01,15):16    model.train()17    pbar = tqdm(train_loader)18 19    train_losses = []20    train_acc = []21    lrs = []22 23    correct = 024    processed = 025    train_loss = 026 27    for batch_idx, (data, target) in enumerate(pbar):28        data, target = data.to(device), target.to(device)29        optimizer.zero_grad()30        y_pred = model(data)31 32        # Calculate loss33        loss = criterion(y_pred, target)34        if L1:35            l1_loss = 036            for p in model.parameters():37                l1_loss = l1_loss + p.abs().sum()38            loss = loss + l1_lambda * l1_loss39        else:40            loss = loss41 42        train_loss += loss.item()43        train_losses.append(loss.item())44 45        # Backpropagation46        loss.backward()47        optimizer.step()48        scheduler.step()49 50        # Update pbar-tqdm51        pred = y_pred.argmax(52            dim=1, keepdim=True53        )  # get the index of the max log-probability54        correct += pred.eq(target.view_as(pred)).sum().item()55        processed += len(data)56 57        pbar.set_description(58            desc=f"Loss={loss.item():0.2f} Accuracy={100*correct/processed:0.2f}"59        )60        train_acc.append(100 * correct / processed)61        lrs.append(scheduler.get_last_lr())62 63    return train_losses, train_acc, lrs64 65 66def test(model, device, criterion, test_loader):67    model.eval()68    test_loss = 069    correct = 070    with torch.no_grad():71        for data, target in test_loader:72            data, target = data.to(device), target.to(device)73            output = model(data)74            test_loss += F.cross_entropy(output, target, reduction="sum").item()75            pred = output.argmax(dim=1, keepdim=True)76            correct += pred.eq(target.view_as(pred)).sum().item()77 78    test_loss /= len(test_loader.dataset)79 80    print(81        "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n".format(82            test_loss,83            correct,84            len(test_loader.dataset),85            100.0 * correct / len(test_loader.dataset),86        )87    )88    test_acc = 100.0 * correct / len(test_loader.dataset)89 90    return test_loss, test_acc91