<filename>app/utils/utils.py
from torch.utils.data import DataLoader, Subset
import torch
from torch import nn, optim
import torchvision
import io

from trust.utils.custom_dataset import load_dataset_custom
from trust.utils.models.resnet import ResNet18
#from trust.utils.models.mobilenet import mymobilenet_v3_small
from trust.utils.utils import SubsetWithTargets
from trust.strategies.smi import SMI
from trust.strategies.random_sampling import RandomSampling


EPOCHS_PER_ROUND = 5
BUDGET = 400
TRAIN_BATCH_SIZE = 20
VAL_BATCH_SIZE = 10
TEST_BATCH_SIZE = 100
DEVICE_ID = 0
DEVICE = "cuda:"+str(DEVICE_ID) if torch.cuda.is_available() else "cpu"
EMBEDDING_TYPE = "features" # TODO
KEEP_EMBEDDING = True

                
def get_model(name, n_classes, device, embedding_type):
    def init_weights(m):
        if isinstance(m, nn.Conv2d):
            torch.nn.init.xavier_uniform_(m.weight)
        elif isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight)
            m.bias.data.fill_(0.01)
    
    if name == "ResNet18":
        if embedding_type == "gradients":
            model = ResNet18(n_classes)
        else:
            model = torchvision.models.resnet18()
    elif name == 'MobileNet':
        if embedding_type == "gradients":
            assert False, "embedding_type 'gradients' is not yet supported for MobileNet" # TODO
        else:
            model = torchvision.models.mobilenet_v3_small(num_classes=n_classes)
    # TODO: Add more models
    model.apply(init_weights)
    model = model.to(device)
    return model

def get_optimizer(model, learning_rate=0.01, m=0.9, wd=5e-4):
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=m, weight_decay=wd)
    return optimizer


def train(i_round, request, subset_history, dataset):
    """
    Run one round of targeted selection.
    Returns dictionary of logged variables for the round.
    """
    train_set = dataset["train_set"]
    val_set = dataset["val_set"]
    lake_set = dataset["lake_set"]
    n_classes = dataset["n_classes"]
    query_idxs = torch.tensor([request["query_idx"]])
    query_set = Subset(train_set, query_idxs)
    # TODO: allow multiple examples in query set?
    # TODO: allow query examples that aren't in train set (from lake, or uploaded)
    if i_round == 0:
        labeled_idxs = torch.tensor([])
        unlabeled_idxs = torch.arange(len(lake_set))
    else:
        # Retrieve saved indices of selected examples, and label them
        labeled_idxs = torch.flatten(torch.tensor(subset_history))
        labeled_subset = SubsetWithTargets(lake_set, labeled_idxs, torch.Tensor(lake_set.targets.float())[labeled_idxs])
        train_set = torch.utils.data.ConcatDataset([train_set, labeled_subset])
        unlabeled_idxs = set(range(len(lake_set))) - set(labeled_idxs.tolist())
        unlabeled_idxs = torch.tensor(sorted(unlabeled_idxs))
        lake_set = Subset(lake_set, unlabeled_idxs)
    trainloader = DataLoader(train_set, batch_size=TRAIN_BATCH_SIZE, shuffle=True, pin_memory=True)
    valloader = DataLoader(val_set, batch_size=VAL_BATCH_SIZE, shuffle=False, pin_memory=True)
    print(f"Number of labeled examples: {len(train_set)}. Number of unlabeled examples: {len(lake_set)}")

    result_dict = {}
    model = get_model(request["model_name"], n_classes, DEVICE, EMBEDDING_TYPE)
    optimizer = get_optimizer(model)
    criterion = nn.CrossEntropyLoss()
    strategy_args = {
        "batch_size":TRAIN_BATCH_SIZE,
        "device":DEVICE,
        "embedding_type":EMBEDDING_TYPE,
        "keep_embedding":KEEP_EMBEDDING,
        "smi_function":request["strategy_name"].lower()
    }
    # Get the initial model by training
    # In the future, we may choose to save the initial model to disk.
    print("Beginning training")
    for i_epoch in range(EPOCHS_PER_ROUND):
        model.train()
        for inputs, targets in trainloader:
            inputs, targets = inputs.to(DEVICE), targets.to(DEVICE, non_blocking=True)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
        # TODO: Add some kind of progress bar or accuracy display for responsiveness
        # TODO: Report initial train/val accuracy for plotting
    print("Training done")
    # Compute values for logging (train/val accuracy)
    with torch.no_grad():
        model.eval()
        for loader, name in [(trainloader, "train"), (valloader, "val")]:
            examples_total = 0
            examples_correct = 0
            for inputs, targets in loader:
                inputs, targets = inputs.to(DEVICE), targets.to(DEVICE, non_blocking=True)
                outputs = model(inputs)
                _, predicted = outputs.max(1)
                examples_total += targets.size(0)
                examples_correct += predicted.eq(targets).sum().item()
            result_dict[name + "_accuracy"] = examples_correct/examples_total

    # Create updated strategy with new data, queries, and model
    # Select examples to be labelled
    if request["strategy_name"].endswith("MI"):
        strategy = SMI(train_set, lake_set, val_set, model, n_classes, strategy_args)
    elif request["strategy_name"] == "Random":
        strategy = RandomSampling(train_set, lake_set, model, n_classes, strategy_args)
    else:
        assert False, "strategy must be one of: 'SIM', 'random'"
    # Update queries
    strategy.update_queries(query_set)
    print("Beginning selection")
    subset = strategy.select(BUDGET)
    print("Selection done")
    # Recover indices w.r.t. original lake set
    true_subset = unlabeled_idxs[subset]
    # TODO: deal with case when len(subset) < budget
    #print(f"true_subset: {true_subset}")
    result_dict["subset"] = true_subset
    result_dict["i_round"] = i_round
    return result_dict
<filename>app/utils/utils.py
from torch.utils.data import DataLoader, Subset
import torch
from torch import nn, optim
import torchvision
import io

from trust.utils.custom_dataset import load_dataset_custom
from trust.utils.models.resnet import ResNet18
#from trust.utils.models.mobilenet import mymobilenet_v3_small
from trust.utils.utils import SubsetWithTargets
from trust.strategies.smi import SMI
from trust.strategies.random_sampling import RandomSampling


EPOCHS_PER_ROUND = 5
BUDGET = 400
TRAIN_BATCH_SIZE = 20
VAL_BATCH_SIZE = 10
TEST_BATCH_SIZE = 100
DEVICE_ID = 0
DEVICE = "cuda:"+str(DEVICE_ID) if torch.cuda.is_available() else "cpu"
EMBEDDING_TYPE = "features" # TODO
KEEP_EMBEDDING = True

                
def get_model(name, n_classes, device, embedding_type):
    def init_weights(m):
        if isinstance(m, nn.Conv2d):
            torch.nn.init.xavier_uniform_(m.weight)
        elif isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight)
            m.bias.data.fill_(0.01)
    
    if name == "ResNet18":
        if embedding_type == "gradients":
            model = ResNet18(n_classes)
        else:
            model = torchvision.models.resnet18()
    elif name == 'MobileNet':
        if embedding_type == "gradients":
            assert False, "embedding_type 'gradients' is not yet supported for MobileNet" # TODO
        else:
            model = torchvision.models.mobilenet_v3_small(num_classes=n_classes)
    # TODO: Add more models
    model.apply(init_weights)
    model = model.to(device)
    return model

def get_optimizer(model, learning_rate=0.01, m=0.9, wd=5e-4):
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=m, weight_decay=wd)
    return optimizer


def train(i_round, request, subset_history, dataset):
    """
    Run one round of targeted selection.
    Returns dictionary of logged variables for the round.
    """
    train_set = dataset["train_set"]
    val_set = dataset["val_set"]
    lake_set = dataset["lake_set"]
    n_classes = dataset["n_classes"]
    query_idxs = torch.tensor([request["query_idx"]])
    query_set = Subset(train_set, query_idxs)
    # TODO: allow multiple examples in query set?
    # TODO: allow query examples that aren't in train set (from lake, or uploaded)
    if i_round == 0:
        labeled_idxs = torch.tensor([])
        unlabeled_idxs = torch.arange(len(lake_set))
    else:
        # Retrieve saved indices of selected examples, and label them
        labeled_idxs = torch.flatten(torch.tensor(subset_history))
        labeled_subset = SubsetWithTargets(lake_set, labeled_idxs, torch.Tensor(lake_set.targets.float())[labeled_idxs])
        train_set = torch.utils.data.ConcatDataset([train_set, labeled_subset])
        unlabeled_idxs = set(range(len(lake_set))) - set(labeled_idxs.tolist())
        unlabeled_idxs = torch.tensor(sorted(unlabeled_idxs))
        lake_set = Subset(lake_set, unlabeled_idxs)
    trainloader = DataLoader(train_set, batch_size=TRAIN_BATCH_SIZE, shuffle=True, pin_memory=True)
    valloader = DataLoader(val_set, batch_size=VAL_BATCH_SIZE, shuffle=False, pin_memory=True)
    print(f"Number of labeled examples: {len(train_set)}. Number of unlabeled examples: {len(lake_set)}")

    result_dict = {}
    model = get_model(request["model_name"], n_classes, DEVICE, EMBEDDING_TYPE)
    optimizer = get_optimizer(model)
    criterion = nn.CrossEntropyLoss()
    strategy_args = {
        "batch_size":TRAIN_BATCH_SIZE,
        "device":DEVICE,
        "embedding_type":EMBEDDING_TYPE,
        "keep_embedding":KEEP_EMBEDDING,
        "smi_function":request["strategy_name"].lower()
    }
    # Get the initial model by training
    # In the future, we may choose to save the initial model to disk.
    print("Beginning training")
    for i_epoch in range(EPOCHS_PER_ROUND):
        model.train()
        for inputs, targets in trainloader:
            inputs, targets = inputs.to(DEVICE), targets.to(DEVICE, non_blocking=True)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
        # TODO: Add some kind of progress bar or accuracy display for responsiveness
        # TODO: Report initial train/val accuracy for plotting
    print("Training done")
    # Compute values for logging (train/val accuracy)
    with torch.no_grad():
        model.eval()
        for loader, name in [(trainloader, "train"), (valloader, "val")]:
            examples_total = 0
            examples_correct = 0
            for inputs, targets in loader:
                inputs, targets = inputs.to(DEVICE), targets.to(DEVICE, non_blocking=True)
                outputs = model(inputs)
                _, predicted = outputs.max(1)
                examples_total += targets.size(0)
                examples_correct += predicted.eq(targets).sum().item()
            result_dict[name + "_accuracy"] = examples_correct/examples_total

    # Create updated strategy with new data, queries, and model
    # Select examples to be labelled
    if request["strategy_name"].endswith("MI"):
        strategy = SMI(train_set, lake_set, val_set, model, n_classes, strategy_args)
    elif request["strategy_name"] == "Random":
        strategy = RandomSampling(train_set, lake_set, model, n_classes, strategy_args)
    else:
        assert False, "strategy must be one of: 'SIM', 'random'"
    # Update queries
    strategy.update_queries(query_set)
    print("Beginning selection")
    subset = strategy.select(BUDGET)
    print("Selection done")
    # Recover indices w.r.t. original lake set
    true_subset = unlabeled_idxs[subset]
    # TODO: deal with case when len(subset) < budget
    #print(f"true_subset: {true_subset}")
    result_dict["subset"] = true_subset
    result_dict["i_round"] = i_round
    return result_dict
#from trust.utils.models.mobilenet import mymobilenet_v3_small # TODO # TODO # TODO: Add more models Run one round of targeted selection.     Returns dictionary of logged variables for the round. # TODO: allow multiple examples in query set? # TODO: allow query examples that aren't in train set (from lake, or uploaded) # Retrieve saved indices of selected examples, and label them # Get the initial model by training # In the future, we may choose to save the initial model to disk. # TODO: Add some kind of progress bar or accuracy display for responsiveness # TODO: Report initial train/val accuracy for plotting # Compute values for logging (train/val accuracy) # Create updated strategy with new data, queries, and model # Select examples to be labelled # Update queries # Recover indices w.r.t. original lake set # TODO: deal with case when len(subset) < budget #print(f"true_subset: {true_subset}")