samplers.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
from scipy.stats import truncnorm, uniform
from tqdm import tqdm

from algorithms import *
from models import *
from utils import *

MIN = - 2.58
MAX = 2.58
NUM_ITERATIONS = 100
MOMENTUM = .9
OFFLINE_NUM_TASKS = 1000
EPS = 1e-3


class RegularSampler:

    def __init__(self, batch_size, num_iterations, tasks, args):

        self.total_num_updates = batch_size * num_iterations
        self.num_updates = 0

    def weight(self, task, value):

        return 1.

    def update(self, value):

        self.num_updates += 1


class OfflineSampler(RegularSampler):

    def _compute_value(self, task):

        value, _ = evaluate_task(task, self.ways, self.support_shots,
                                 self.query_shots, self.model, self.adapt,
                                 None, False)
        return value.item()

    def __init__(self, batch_size, num_iterations, tasks, args):

        super(OfflineSampler, self).__init__(batch_size, num_iterations, tasks,
                                             args)

        model = MODELS[args['model']](args['ways'])
        algo_model, adapt = ALGORITHMS[args['algorithm']]
        model = algo_model(model)
        model.load_state_dict(torch.load(args['params_path']))
        model.to('cuda')
        model.eval()

        self.ways = args['ways']
        self.support_shots = args['support_shots']
        self.query_shots = args['query_shots']
        self.model = model
        self.adapt = adapt

        values = []
        for idx in tqdm(range(OFFLINE_NUM_TASKS)):
            task = tasks[idx]
            value = self._compute_value(task)
            values.append(value)

        self.mean = np.mean(values)
        self.var = np.var(values, ddof=1)

    def weight(self, task, value):

        value = self._compute_value(task)
        return self._weight(value)


class OnlineSampler(RegularSampler):

    def __init__(self, batch_size, num_iterations, tasks, args):

        super(OnlineSampler, self).__init__(batch_size, num_iterations, tasks,
                                            args)
        self.mean = None
        self.var = None
        self.batch_size = batch_size

    def _update(self, value):

        self.num_updates += 1
        if self.mean is None and self.var is None:
            mean = value
            self.mean = mean
        elif self.var is None:
            mean = value
            var = (value - self.mean) ** 2.
            self.mean = (MOMENTUM * self.mean) + ((1. - MOMENTUM) * mean)
            self.var = var
        else:
            mean = value
            var = (value - self.mean) ** 2.
            self.mean = (MOMENTUM * self.mean) + ((1. - MOMENTUM) * mean)
            self.var = (MOMENTUM * self.var) + ((1. - MOMENTUM) * var)

    def weight(self, task, value):

        if self.num_updates < self.batch_size * NUM_ITERATIONS:
            return 1.
        return self._weight(value)

    def update(self, value):

        return self._update(value)


class Importance:

    def _convert_value(self, value):

        return (value - self.mean) / (self.var ** .5)

    def _stable_weight(self, target, value):

        observed = truncnorm.pdf(value, MIN, MAX)

        if target < EPS and observed < EPS:
            return 1.
        if observed < EPS:
            return target / (observed + EPS)

        return target / observed


class EasyImportance(Importance):

    def _weight(self, value):

        value = self._convert_value(value)

        target = uniform.pdf(value, MIN, - MIN)

        return self._stable_weight(target, value)


class HardImportance(Importance):

    def _weight(self, value):

        value = self._convert_value(value)

        target = uniform.pdf(value, 0., MAX)

        return self._stable_weight(target, value)


class CurriculumImportance(Importance):

    def _weight(self, value):

        value = self._convert_value(value)

        lam = self.num_updates / self.total_num_updates
        mean = ((1. - lam) * MIN) + (lam * MAX)
        target = truncnorm.pdf(value, MIN - mean, MAX - mean, mean, 1.)

        return self._stable_weight(target, value)


class UniformImportance(Importance):

    def _weight(self, value):

        value = self._convert_value(value)

        target = uniform.pdf(value, MIN, MAX - MIN)

        return self._stable_weight(target, value)


class OfflineEasyImportanceSampler(OfflineSampler, EasyImportance):

    def dummy(self):

        return


class OfflineHardImportanceSampler(OfflineSampler, HardImportance):

    def dummy(self):

        return


class OfflineCurriculumImportanceSampler(OfflineSampler, CurriculumImportance):

    def dummy(self):

        return


class OfflineUniformImportanceSampler(OfflineSampler, UniformImportance):

    def dummy(self):

        return


class OnlineUniformImportanceSampler(OnlineSampler, UniformImportance):

    def dummy(self):

        return


SAMPLERS = {
    'regular': RegularSampler,
    'offline_easy': OfflineEasyImportanceSampler,
    'offline_hard': OfflineHardImportanceSampler,
    'offline_curriculum': OfflineCurriculumImportanceSampler,
    'offline_uniform': OfflineUniformImportanceSampler,
    'online_uniform': OnlineUniformImportanceSampler
}