utils.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import random
import torch

torch.backends.cudnn.benchmark = True


def setup(seed=42):

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


loss_fn = torch.nn.functional.cross_entropy


def acc_fn(yh, y):

    return (yh.argmax(dim=1) == y).float().mean()


def euclidean_distance(a, b):

    n = a.shape[0]
    m = b.shape[0]
    dist = ((a.unsqueeze(1).expand(n, m, -1) -
             b.unsqueeze(0).expand(n, m, -1)) ** 2).sum(dim=2)

    return dist


class Normalize(torch.nn.Module):

    def __init__(self):

        super(Normalize, self).__init__()
        self.scale = torch.nn.Parameter(torch.ones(1))

    def forward(self, x):

        return self.scale * torch.nn.functional.normalize(x)


def task_to_support_query(task, ways, support_shots, query_shots):

    x, y = task
    ind = torch.sort(y)[1]
    x = x[ind]
    y = y[ind]

    shots = support_shots + query_shots
    support_ind = np.zeros(y.size(0), dtype=bool)
    for support_shot in range(support_shots):
        support_ind[(np.arange(ways) * shots) + support_shot] = True
    query_ind = torch.from_numpy(~support_ind)
    support_ind = torch.from_numpy(support_ind)

    support_x = x[support_ind]
    support_y = y[support_ind]
    query_x = x[query_ind]
    query_y = y[query_ind]

    return support_x, support_y, query_x, query_y


def evaluate_task(task, ways, support_shots, query_shots, model, adapt,
                  Sampler, train):

    support_x, support_y, query_x, query_y = task_to_support_query(
        task, ways, support_shots, query_shots)

    support_x = support_x.to('cuda', non_blocking=True)
    support_y = support_y.to('cuda', non_blocking=True)
    query_x = query_x.to('cuda', non_blocking=True)
    query_y = query_y.to('cuda', non_blocking=True)

    query_yh = adapt(support_x, support_y, query_x, model, train)

    loss = loss_fn(query_yh, query_y)
    acc = acc_fn(query_yh, query_y)

    if not train:
        return loss, acc

    weight = Sampler.weight(task, loss.item())
    Sampler.update(loss.item())

    return loss, acc, weight