text:
import sympy
from ndispers._baseclass import Medium, wl, phi, theta
from ndispers.helper import vars2

class Calcite(Medium):
    """
    calcite (Ca C O_3) crystal

    - Point group : -3m
    - Crystal system : Hexagonal
    - Dielectic principal axis, z // c-axis (x, y-axes are arbitrary)
    - Negative uniaxial, with optic axis parallel to z-axis
    - Tranparency range :

    Dispersion formula for refractive index
    ---------------------------------------
    n(wl)
    = sqrt(1 + A1_o * wl**2 / (wl**2 - B1_o**2) + A2_o * wl**2 / (wl**2 - B2_o**2) + A3_o * wl**2 / (wl**2 - B3_o**2) + A4_o * wl**2 / (wl**2 - B4_o**2)) for o-ray
    = sqrt(1 + A1_e * wl**2 / (wl**2 - B1_e**2) + A2_e * wl**2 / (wl**2 - B2_e**2) + A3_e * wl**2 / (wl**2 - B3_e**2)) for e-ray

    Validity range
    ---------------
    0.2 to 2.2 um for o-ray
    0.2 to 3.3 um for e-ray

    Ref
    ----
    Handbook of Optics: Devices, Measurements, and Properties, Volume II, by Michael Bass (ed),
    Chapter 33: PROPERTIES OF CRYSTALS AND GLASSES, William J. Tropf, Michael E. Thomas, and Terry J. Harris

    Usage
    ------
    >>> abbo = ndispers.media.crystals.AlphaBBO()
    >>> abbo.n(0.6, 0, pol='o') # for o-ray, n does not depend on theta.
    >>> abbo.n(0.6, 0.5*pi, pol='e') # along z-axis, it is pure e-ray.
    >>> abbo.n(0.6, 0.23*pi, pol='e')
    >>> abbo.n(0.6, 0*pi, pol='e') # for theta = 0 rad, it corresponds to o-ray.
    >>> abbo.GVD(0.6, 0.23*pi, pol='e')

    """
    __slots__ = ["_Calcite__plane", "_Calcite__theta_rad", "_Calcite__phi_rad",
                 "_A1_o", "_B1_o", "_A2_o", "_B2_o", "_A3_o", "_B3_o", "_A4_o", "_B4_o",
                 "_A1_e", "_B1_e", "_A2_e", "_B2_e", "_A3_e", "_B3_e"]

    def __init__(self):
        super().__init__()
        self._Calcite__plane = 'arb'
        self._Calcite__theta_rad = 'var'
        self._Calcite__phi_rad = 'arb'

        """ Constants of dispersion formula """
        # For ordinary ray
        self._A1_o = 0.8559
        self._B1_o = 0.0588
        self._A2_o = 0.8391
        self._B2_o = 0.141
        self._A3_o = 0.0009
        self._B3_o = 0.197
        self._A4_o = 0.6845
        self._B4_o = 7.005
        # For extraordinary ray
        self._A1_e = 1.0856
        self._B1_e = 0.07897
        self._A2_e = 0.0988
        self._B2_e = 0.142
        self._A3_e = 0.317
        self._B3_e = 11.468

    @property
    def plane(self):
        return self._Calcite__plane

    @property
    def theta_rad(self):
        return self._Calcite__theta_rad

    @property
    def phi_rad(self):
        return self._Calcite__phi_rad

    @property
    def symbols(self):
        return [wl, theta, phi]

    @property
    def constants(self):
        print(vars2(self))

    def n_o_expr(self):
        """ Sympy expression, dispersion formula for o-ray """
        return sympy.sqrt(1 + self._A1_o * wl**2 / (wl**2 - self._B1_o**2) + self._A2_o * wl**2 / (wl**2 - self._B2_o**2) + self._A3_o * wl**2 / (wl**2 - self._B3_o**2) + self._A4_o * wl**2 / (wl**2 - self._B4_o**2))

    def n_e_expr(self):
        """ Sympy expression, dispersion formula for theta=90 deg e-ray """
        return sympy.sqrt(1 + self._A1_e * wl**2 / (wl**2 - self._B1_e**2) + self._A2_e * wl**2 / (wl**2 - self._B2_e**2) + self._A3_e * wl**2 / (wl**2 - self._B3_e**2))

    def n_expr(self, pol):
        """"
        Sympy expression,
        dispersion formula of a general ray with an angle theta to optic axis. If theta = 0, this expression reduces to 'no_expre'.

        n(theta) = n_e / sqrt( sin(theta)**2 + (n_e/n_o)**2 * cos(theta)**2 )
        """
        if pol == 'o':
            return self.n_o_expr()
        elif pol == 'e':
            return self.n_e_expr() / sympy.sqrt( sympy.sin(theta)**2 + (self.n_e_expr()/self.n_o_expr())**2 * sympy.cos(theta)**2 )
        else:
            raise ValueError("pol = '%s' must be 'o' or 'e'" % pol)

    def n(self, wl_um, theta_rad, pol='o'):
        """
        Refractive index as a function of wavelength, theta and phi angles for each eigen polarization of light.

        input
        ------
        wl_um     :  float, wavelength in um
        theta_rad :  float, 0 to pi radians
        pol       :  str, 'o' or 'e', polarization of light

        return
        -------
        Refractive index, float or array_like
        """
        return super().n(wl_um, theta_rad, 0, pol=pol)

    def dn_wl(self, wl_um, theta_rad, pol='o'):
        return super().dn_wl(wl_um, theta_rad, 0, pol=pol)

    def d2n_wl(self, wl_um, theta_rad, pol='o'):
        return super().d2n_wl(wl_um, theta_rad, 0, pol=pol)

    def d3n_wl(self, wl_um, theta_rad, pol='o'):
        return super().d3n_wl(wl_um, theta_rad, 0, pol=pol)

    def GD(self, wl_um, theta_rad, pol='o'):
        """Group Delay [fs/mm]"""
        return super().GD(wl_um, theta_rad, 0, pol=pol)

    def GV(self, wl_um, theta_rad, pol='o'):
        """Group Velocity [um/fs]"""
        return super().GV(wl_um, theta_rad, 0, pol=pol)

    def ng(self, wl_um, theta_rad, pol='o'):
        """Group index, c/Group velocity"""
        return super().ng(wl_um, theta_rad, 0, pol=pol)

    def GVD(self, wl_um, theta_rad, pol='o'):
        """Group Delay Dispersion [fs^2/mm]"""
        return super().GVD(wl_um, theta_rad, 0, pol=pol)

    def TOD(self, wl_um, theta_rad, pol='o'):
        """Third Order Dispersion [fs^3/mm]"""
        return super().TOD(wl_um, theta_rad, 0, pol=pol)

    def woa_theta(self, wl_um, theta_rad, T_degC, pol='e'):
        return super().woa_theta(wl_um, theta_rad, 0, T_degC, pol=pol)

    def woa_phi(self, wl_um, theta_rad, T_degC, pol='e'):
        return super().woa_phi(wl_um, theta_rad, 0, T_degC, pol=pol)

max_stars_repo_name: akihiko-shimura/ndispers

max_issues_repo_name: akihiko-shimura/ndispers

max_forks_repo_name: akihiko-shimura/ndispers

reason: import sympy


text:
# Copyright 2021 Fredrik Hallgren
#
# 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
#
#     https://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.

"""
Class that encapsulates a kernel function and its parameters.

Contains functionality to evaluate a kernel function and calculate
a kernel matrix, and keeps track of the maximum kernel value.

"""

from itertools import product

import numpy as np

from nystrompca.utils import demean_matrix

class Kernel:

    """
    Wrapper for kernel functions.

    Please see 'nystrompca.base.kernel_machine' for argument descriptions

    """
    def __init__(self, kernel:    str   = 'rbf',
                       sigma:     float = 3,
                       degree:    int   = 2,
                       coef0:     float = 1,
                       normalize: bool  = True ):

        self.sigma        = sigma
        self.degree       = degree
        self.coef0        = coef0
        self.normalize    = normalize

        self.name = kernel

        bound: float

        if kernel == 'rbf':
            kernel_fcn = lambda x, y: rbf(x, y, degree, sigma)
            bound = 1

        elif kernel == 'poly':
            kernel_fcn = lambda x, y: poly(x, y, degree, coef0, normalize)
            if normalize:
                bound = 1
            else:
                bound = np.inf

        elif kernel == 'linear':
            kernel_fcn = lambda x, y: poly(x, y, 1, 0, normalize=False)
            bound = np.inf

        elif kernel == 'laplace':
            kernel_fcn = lambda x, y: rbf(x, y, 2, sigma)
            bound = 1

        elif kernel == 'cauchy':
            kernel_fcn = lambda x, y: cauchy(x, y, sigma)
            bound = 1

        elif kernel == 'cosine':
            kernel_fcn = lambda x, y: poly(x, y, 1, 0, normalize=True)
            bound = 1

        self.kernel_fcn = kernel_fcn

        self.bound = bound

    def __call__(self, x, y):
        return self.kernel_fcn(x, y)

    def get_name(self) -> str:
        """
        Getter for the kernel name

        """
        return self.name

    def get_bound(self) -> float:
        """
        Getter for the kernel function bound

        """
        return self.bound

    def calc_sigma(self, X: np.ndarray):
        """
        Calculate sigma based on a common heuristic

        Calculate the bandwidth parameter for the RBF and Cauchy kernels
        as the average distance between all pairs of data points

        Parameters
        ----------
        X : numpy.ndarray, 2d
            Input data

        """
        assert X.ndim == 2, "Two-dimensional input data expected"

        n = X.shape[0]

        pairwise_diff = np.tile(X, (n,1)) - X.repeat(n, axis=0)
        sigma = np.sqrt((pairwise_diff**2).sum(1)).sum() / n**2

        # Reinitialize with new sigma
        self.__init__(self.name, sigma, self.degree, # type: ignore[misc]
                      self.coef0, self.normalize)

    def matrix(self, X1:     np.ndarray,
                     X2:     np.ndarray = None,
                     demean: bool       = True) -> np.ndarray:
        """
        Calculate kernel matrix.

        Calculate the kernel matrices between two datasets.

        """
        n = X1.shape[0]

        # If only one dataset is supplied we only need to calculate half
        # the kernel evaluations explicitly plus the diagonal
        if X2 is None:
            K = np.zeros((n,n))
            for i in range(n):
                K[i, i] = self(X1[i], X1[i]) / 2
                for j in range(i+1, n):
                    K[i,j] = self(X1[i], X1[j])
            K = K + K.T

        # Otherwise we need to calculate all elements explicitly
        else:
            m = X2.shape[0]
            K = np.zeros((n,m))
            for i, j in product(range(n), range(m)):
                K[i,j] = self(X1[i],X2[j])

        if demean:
            K = demean_matrix(K)

        return K

def rbf(x:      np.ndarray,
        y:      np.ndarray,
        degree: int,
        sigma:  float     ) -> float:
    """
    Radial Basis Functions kernel

    """
    return np.exp(-np.sqrt((x - y) @ (x - y))**degree / (sigma**2))

def poly(x:         np.ndarray,
         y:         np.ndarray,
         degree:    int,
         coef0:     float,
         normalize: bool      = True) -> float:
    """
    Polynomial kernel

    """
    value = (x @ y + coef0)**degree
    if normalize:
        value /= ((x @ x + coef0) * (y @ y + coef0))**(degree/2)

    return value

def cauchy(x:     np.ndarray,
           y:     np.ndarray,
           sigma: float     ) -> float:
    """
    Cauchy kernel

    """
    return 1 / (1 + (x - y) @ (x - y) / sigma**2)

ext: py

max_stars_repo_name: fredhallgren/nystrompca

max_issues_repo_name: fredhallgren/nystrompca

max_forks_repo_name: fredhallgren/nystrompca

reason: import numpy


text:
from sklearn.metrics import auc
import numpy as np
from sklearn.metrics import precision_recall_curve
import sklearn.metrics as metrics
import sys

the_list=['<20','20~30','30~40','40~50','50~60','60~70','70~80','=80']
for the_file in the_list:
    y=np.genfromtxt(the_file,delimiter='\t')[:,1]
    pred=np.genfromtxt(the_file,delimiter='\t')[:,2]

    fpr, tpr, thresholds = metrics.roc_curve(y, pred, pos_label=1)
    the_auc=metrics.auc(fpr, tpr)
    print(the_file,the_auc)

max_stars_repo_name: GuanLab/sepsis

max_issues_repo_name: GuanLab/sepsis

max_forks_repo_name: GuanLab/sepsis


text:
import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter
from torch import Tensor
import numpy as np
import torch.nn.init as nn_init
from unet_parts import *
from math import sqrt
# import pdb

def load_my_state_dict(net, state_dict):
    own_state = net.state_dict()
    for name, param in state_dict.items():
        if name not in own_state:
            continue
        if isinstance(param, Parameter):
            # backwards compatibility for serialized parameters
            param = param.data
        elif isinstance(param, Tensor):
            # backwards compatibility for serialized parameters
            param = param
        own_state[name].copy_(param)

class GaussianNoise(nn.Module):
    def __init__(self, sigma):
        super(GaussianNoise, self).__init__()
        self.sigma = sigma

    def forward(self, input):
        if self.training:
            noise = Variable(input.data.new(input.size()).normal_(std=self.sigma))
            return input + noise
        else:
            return input

class Expression(nn.Module):
    def __init__(self, func):
        super(Expression, self).__init__()
        self.func = func

    def forward(self, input):
        return self.func(input)

class WN_Linear(nn.Linear):
    def __init__(self, in_features, out_features, bias=True, train_scale=False, init_stdv=1.0):
        super(WN_Linear, self).__init__(in_features, out_features, bias=bias)
        if train_scale:
            self.weight_scale = Parameter(torch.ones(self.out_features))
        else:
            self.register_buffer('weight_scale', torch.Tensor(out_features))

        self.train_scale = train_scale
        self.init_mode = False
        self.init_stdv = init_stdv

        self._reset_parameters()

    def _reset_parameters(self):
        self.weight.data.normal_(0, std=0.05)
        if self.bias is not None:
            self.bias.data.zero_()
        if self.train_scale:
            self.weight_scale.data.fill_(1.)
        else:
            self.weight_scale.fill_(1.)

    def forward(self, input):
        if self.train_scale:
            weight_scale = self.weight_scale
        else:
            weight_scale = Variable(self.weight_scale)

        # normalize weight matrix and linear projection
        # norm_weight = self.weight * (weight_scale.unsqueeze(1) / torch.sqrt((self.weight ** 2).sum(1) + 1e-6)).expand_as(self.weight)
        norm_weight = self.weight * (weight_scale / torch.sqrt((self.weight ** 2).sum(1) + 1e-6)).unsqueeze(1).expand_as(self.weight)
        activation = F.linear(input, norm_weight)

        if self.init_mode == True:
            mean_act = activation.mean(0).squeeze(0)
            activation = activation - mean_act.expand_as(activation)

            inv_stdv = self.init_stdv / torch.sqrt((activation ** 2).mean(0) + 1e-6).squeeze(0)
            activation = activation * inv_stdv.expand_as(activation)

            if self.train_scale:
                self.weight_scale.data = self.weight_scale.data * inv_stdv.data
            else:
                self.weight_scale = self.weight_scale * inv_stdv.data
            self.bias.data = - mean_act.data * inv_stdv.data

        else:
            if self.bias is not None:
                activation = activation + self.bias.expand_as(activation)

        return activation

class WN_Conv2d(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, train_scale=False, init_stdv=1.0):
        super(WN_Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        if train_scale:
            self.weight_scale = Parameter(torch.Tensor(out_channels))
        else:
            self.register_buffer('weight_scale', torch.Tensor(out_channels))

        self.train_scale = train_scale
        self.init_mode = False
        self.init_stdv = init_stdv

        self._reset_parameters()

    def _reset_parameters(self):
        self.weight.data.normal_(std=0.05)
        if self.bias is not None:
            self.bias.data.zero_()
        if self.train_scale:
            self.weight_scale.data.fill_(1.)
        else:
            self.weight_scale.fill_(1.)

    def forward(self, input):
        if self.train_scale:
            weight_scale = self.weight_scale
        else:
            weight_scale = Variable(self.weight_scale)
        # normalize weight matrix and linear projection [out x in x h x w]
        # for each output dimension, normalize through (in, h, w) = (1, 2, 3) dims
        # norm_weight = self.weight * (weight_scale[:,None,None,None] / torch.sqrt((self.weight ** 2).sum(3).sum(2).sum(1) + 1e-6)).expand_as(self.weight)
        norm_weight = self.weight * (weight_scale / torch.sqrt((self.weight ** 2).sum(3).sum(2).sum(1) + 1e-6))[:,None,None,None].expand_as(self.weight)
        activation = F.conv2d(input, norm_weight, bias=None,
                              stride=self.stride, padding=self.padding,
                              dilation=self.dilation, groups=self.groups)

        if self.init_mode == True:
            mean_act = activation.mean(3).mean(2).mean(0).squeeze()
            activation = activation - mean_act[None,:,None,None].expand_as(activation)

            inv_stdv = self.init_stdv / torch.sqrt((activation ** 2).mean(3).mean(2).mean(0) + 1e-6).squeeze()
            activation = activation * inv_stdv[None,:,None,None].expand_as(activation)

            if self.train_scale:
                self.weight_scale.data = self.weight_scale.data * inv_stdv.data
            else:
                self.weight_scale = self.weight_scale * inv_stdv.data
            self.bias.data = - mean_act.data * inv_stdv.data

        else:
            if self.bias is not None:
                activation = activation + self.bias[None,:,None,None].expand_as(activation)

        return activation

class WN_ConvTranspose2d(nn.ConvTranspose2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, train_scale=False, init_stdv=1.0):
        super(WN_ConvTranspose2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, output_padding, groups, bias)
        if train_scale:
            self.weight_scale = Parameter(torch.Tensor(out_channels))
        else:
            self.register_buffer('weight_scale', torch.Tensor(out_channels))

        self.train_scale = train_scale
        self.init_mode = False
        self.init_stdv = init_stdv

        self._reset_parameters()

    def _reset_parameters(self):
        self.weight.data.normal_(std=0.05)
        if self.bias is not None:
            self.bias.data.zero_()
        if self.train_scale:
            self.weight_scale.data.fill_(1.)
        else:
            self.weight_scale.fill_(1.)

    def forward(self, input, output_size=None):
        if self.train_scale:
            weight_scale = self.weight_scale
        else:
            weight_scale = Variable(self.weight_scale)
        # normalize weight matrix and linear projection [in x out x h x w]
        # for each output dimension, normalize through (in, h, w)  = (0, 2, 3) dims
        # norm_weight = self.weight * (weight_scale[None,:,None,None] / torch.sqrt((self.weight ** 2).sum(3).sum(2).sum(0) + 1e-6)).expand_as(self.weight)
        norm_weight = self.weight * (weight_scale / torch.sqrt((self.weight ** 2).sum(3).sum(2).sum(0) + 1e-6))[None,:,None,None].expand_as(self.weight)
        output_padding = self._output_padding(input, output_size)
        activation = F.conv_transpose2d(input, norm_weight, bias=None,
                                        stride=self.stride, padding=self.padding,
                                        output_padding=output_padding, groups=self.groups)

        if self.init_mode == True:
            mean_act = activation.mean(3).mean(2).mean(0).squeeze()
            activation = activation - mean_act[None,:,None,None].expand_as(activation)

            inv_stdv = self.init_stdv / torch.sqrt((activation ** 2).mean(3).mean(2).mean(0) + 1e-6).squeeze()
            activation = activation * inv_stdv[None,:,None,None].expand_as(activation)

            if self.train_scale:
                self.weight_scale.data = self.weight_scale.data * inv_stdv.data
            else:
                self.weight_scale = self.weight_scale * inv_stdv.data
            self.bias.data = - mean_act.data * inv_stdv.data

        else:
            if self.bias is not None:
                activation = activation + self.bias[None,:,None,None].expand_as(activation)

        return activation

class thisModule(nn.Module):
    def load_my_state_dict(self, state_dict):
        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name not in own_state:
                continue
            if isinstance(param, Parameter):
                # backwards compatibility for serialized parameters
                param = param.data
            elif isinstance(param, Tensor):
                # backwards compatibility for serialized parameters
                param = param
            own_state[name].copy_(param)

class Discriminative(thisModule):
    def __init__(self, config, ema=False):
        super(Discriminative, self).__init__()

        print '===> Init small-conv for {}'.format(config.dataset)

        self.noise_size = config.noise_size
        self.num_label  = config.num_label

        if hasattr(config, 'double_input_size') and config.double_input_size:
            self.side = 64
        else:
            self.side = 32
        if config.dataset == 'svhn':
            n_filter_1, n_filter_2 = 64, 128
        elif config.dataset == 'cifar':
            n_filter_1, n_filter_2 = 96, 192
        elif config.dataset == 'stl10':
            n_filter_1, n_filter_2 = 128, 192
        elif config.dataset == 'coil20':
            n_filter_1, n_filter_2 = 96, 192
        elif config.dataset == 'imagenet10':
            n_filter_1, n_filter_2 = 96, 192
        else:
            raise ValueError('dataset not found: {}'.format(config.dataset))

        drop_ = 0. if hasattr(config, "drop") and config.drop else 0.5
        # Assume X is of size [batch x 3 x 32 x 32]
        self.core_net = nn.Sequential(  # input side size not aware

            nn.Sequential(GaussianNoise(0.05), nn.Dropout2d(0.15)) if config.dataset == 'svhn' \
                else nn.Sequential(GaussianNoise(0.05), nn.Dropout2d(0.2)),

            WN_Conv2d(         3, n_filter_1, 3, 1, 1), nn.LeakyReLU(0.2),
            WN_Conv2d(n_filter_1, n_filter_1, 3, 1, 1), nn.LeakyReLU(0.2),
            WN_Conv2d(n_filter_1, n_filter_1, 3, 2, 1), nn.LeakyReLU(0.2),

            nn.Dropout2d(drop_) if config.dataset == 'svhn' else nn.Dropout(drop_),

            WN_Conv2d(n_filter_1, n_filter_2, 3, 1, 1), nn.LeakyReLU(0.2),
            WN_Conv2d(n_filter_2, n_filter_2, 3, 1, 1), nn.LeakyReLU(0.2),
            WN_Conv2d(n_filter_2, n_filter_2, 3, 2, 1), nn.LeakyReLU(0.2),

            nn.Dropout2d(drop_) if config.dataset == 'svhn' else nn.Dropout(drop_),

            WN_Conv2d(n_filter_2, n_filter_2, 3, 1, 0), nn.LeakyReLU(0.2),
            WN_Conv2d(n_filter_2, n_filter_2, 1, 1, 0), nn.LeakyReLU(0.2),
            WN_Conv2d(n_filter_2, n_filter_2, 1, 1, 0), nn.LeakyReLU(0.2),

            Expression(lambda tensor: tensor.mean(3).mean(2).squeeze()),
        )

        self.out_net = WN_Linear(n_filter_2, self.num_label, train_scale=True, init_stdv=0.1)

        if ema:
            for param in self.parameters():
                param.detach_()

    def forward(self, X, feat=False):
        if X.dim() == 2:
            X = X.view(X.size(0), 3, self.side, self.side)

        if feat:
            return self.core_net(X)
        else:
            return self.out_net(self.core_net(X))

class Discriminative_out(thisModule):
    def __init__(self, config):
        super(Discriminative_out, self).__init__()

        print '===> Init small-conv for {}'.format(config.dataset)

        self.num_label  = config.num_label

        if config.dataset == 'svhn':
            n_filter_1, n_filter_2 = 64, 128
        elif config.dataset == 'cifar':
            n_filter_1, n_filter_2 = 96, 192
        elif config.dataset == 'imagenet10':
            n_filter_1, n_filter_2 = 96, 192
        else:
            raise ValueError('dataset not found: {}'.format(config.dataset))

        # Assume X is of size [batch x 3 x 32 x 32]

        self.out_net2 = nn.Sequential(
            WN_Linear(n_filter_2, n_filter_2*2, train_scale=True, init_stdv=0.1),
            WN_Linear(n_filter_2*2, self.num_label, train_scale=True, init_stdv=0.1)
        )

        if config.dis_triple:
            self.out_net3 = nn.Sequential(
                WN_Linear(n_filter_2, self.num_label, train_scale=True, init_stdv=0.1)
            )

    def forward(self, X):
            return self.out_net2(X)

def generator(image_side, noise_size=100, large=False, gen_mode='z2i'):
    if gen_mode == 'i2i':
        return UNet(3, 3, large=large, upbilinear=True)
    elif gen_mode == 'z2i':
        return Generator(image_side, noise_size=noise_size, large=large)

class Generator(thisModule):
    def __init__(self, image_side, noise_size=100, large=False):
        super(Generator, self).__init__()

        self.noise_size = noise_size
        self.image_side = image_side

        if self.image_side == 64:
            self.core_net = nn.Sequential(
                nn.Linear(self.noise_size, 2 * 2 * 1024, bias=False), nn.BatchNorm1d(2 * 2 * 1024), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), 1024, 2, 2)),
                nn.ConvTranspose2d(1024, 512, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(512), nn.ReLU(),
                nn.ConvTranspose2d(512, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 128, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                WN_ConvTranspose2d(128,   3, 5, 2, 2, 1, train_scale=True, init_stdv=0.1), nn.Tanh(),
            )
        elif self.image_side == 96:
            self.core_net = nn.Sequential(
                nn.Linear(self.noise_size, 3 * 3 * 1024, bias=False), nn.BatchNorm1d(3 * 3 * 1024), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), 1024, 3, 3)),
                nn.ConvTranspose2d(1024, 512, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(512), nn.ReLU(),
                nn.ConvTranspose2d(512, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 128, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                WN_ConvTranspose2d(128,   3, 5, 2, 2, 1, train_scale=True, init_stdv=0.1), nn.Tanh(),
            )
        elif self.image_side == 128:
            self.core_net = nn.Sequential(
                nn.Linear(self.noise_size, 2 * 2 * 1024, bias=False), nn.BatchNorm1d(2 * 2 * 1024), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), 1024, 2, 2)),
                nn.ConvTranspose2d(1024, 512, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(512), nn.ReLU(),
                nn.ConvTranspose2d(512, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 128, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                WN_ConvTranspose2d(128,   3, 5, 2, 2, 1, train_scale=True, init_stdv=0.1), nn.Tanh(),
            )
        elif self.image_side == 32:   # side: 32
            self.core_net = nn.Sequential(
                nn.Linear(self.noise_size, 4 * 4 * 512, bias=False), nn.BatchNorm1d(4 * 4 * 512), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), 512, 4, 4)),
                nn.ConvTranspose2d(512, 256, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.ConvTranspose2d(256, 128, 5, 2, 2, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                WN_ConvTranspose2d(128,   3, 5, 2, 2, 1, train_scale=True, init_stdv=0.1), nn.Tanh(),
            )
        else:
            assert False, "The image side out of case: {}".format(self.image_side)

    def forward(self, noise):
        output = self.core_net(noise)

        return output

class Encoder(thisModule):
    def __init__(self, image_side, noise_size=100, output_params=False):
        super(Encoder, self).__init__()

        self.noise_size = noise_size
        self.image_side = image_side

        if self.image_side == 32:
            self.core_net = nn.Sequential(
                nn.Conv2d(  3, 128, 5, 2, 2, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                nn.Conv2d(128, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 512, 5, 2, 2, bias=False), nn.BatchNorm2d(512), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), -1)), # 512 * 4 * 4)),
            )
            if output_params:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(4 * 4 * 512, self.noise_size*2, train_scale=True, init_stdv=0.1))
                self.core_net.add_module(str(len(self.core_net._modules)), Expression(lambda x: torch.chunk(x, 2, 1)))
            else:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(4 * 4 * 512, self.noise_size, train_scale=True, init_stdv=0.1))

        elif self.image_side == 96:
            self.core_net = nn.Sequential(
                nn.Conv2d(  3, 128, 5, 2, 2, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                nn.Conv2d(128, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 512, 5, 2, 2, bias=False), nn.BatchNorm2d(512), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), -1)), # 512 * 3 * 3)),
            )
            if output_params:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(3 * 3 * 512, self.noise_size*2, train_scale=True, init_stdv=0.1))
                self.core_net.add_module(str(len(self.core_net._modules)), Expression(lambda x: torch.chunk(x, 2, 1)))
            else:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(3 * 3 * 512, self.noise_size, train_scale=True, init_stdv=0.1))

        elif self.image_side == 64:
            self.core_net = nn.Sequential(
                nn.Conv2d(  3, 128, 5, 2, 2, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                nn.Conv2d(128, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 512, 5, 2, 2, bias=False), nn.BatchNorm2d(512), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), -1)), # 512 * 2 * 2)),
            )
            if output_params:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(2 * 2 * 512, self.noise_size*2, train_scale=True, init_stdv=0.1))
                self.core_net.add_module(str(len(self.core_net._modules)), Expression(lambda x: torch.chunk(x, 2, 1)))
            else:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(2 * 2 * 512, self.noise_size, train_scale=True, init_stdv=0.1))
        elif self.image_side == 128:
            self.core_net = nn.Sequential(
                nn.Conv2d(  3, 128, 5, 2, 2, bias=False), nn.BatchNorm2d(128), nn.ReLU(),
                nn.Conv2d(128, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 256, 5, 2, 2, bias=False), nn.BatchNorm2d(256), nn.ReLU(),
                nn.Conv2d(256, 512, 5, 2, 2, bias=False), nn.BatchNorm2d(512), nn.ReLU(),
                Expression(lambda tensor: tensor.view(tensor.size(0), -1)),
            )
            if output_params:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(2 * 2 * 512, self.noise_size*2, train_scale=True, init_stdv=0.1))
                self.core_net.add_module(str(len(self.core_net._modules)), Expression(lambda x: torch.chunk(x, 2, 1)))
            else:
                self.core_net.add_module(str(len(self.core_net._modules)), WN_Linear(2 * 2 * 512, self.noise_size, train_scale=True, init_stdv=0.1))
        else:
            assert False, "The image side out of case: {}".format(self.image_side)

    def forward(self, input):

        output = self.core_net(input)

        return output

    def load_my_state_dict(self, state_dict):

        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name not in own_state:
                continue
            if isinstance(param, Parameter):
                # backwards compatibility for serialized parameters
                param = param.data
            if isinstance(param, Tensor):
                # backwards compatibility for serialized parameters
                param = param
            own_state[name].copy_(param)

# ref: https://github.com/kevinlu1211/pytorch-unet-resnet-50-encoder/blob/master/u_net_resnet_50_encoder.py
class UNetWithResnetEncoder(thisModule):
    DEPTH = 6

    def __init__(self, n_classes=3, res='50'):
        super(UNetWithResnetEncoder, self).__init__()
        res = str(res)
        if res == '50':
            resnet = torchvision.models.resnet.resnet50(pretrained=True)
            chas = [2048, 1024, 512, 256, 128, 64]
        elif res == '34':
            resnet = torchvision.models.resnet.resnet34(pretrained=True)
            chas = [512, 256, 128, 64, 64, 64]
        elif res == '18':
            resnet = torchvision.models.resnet.resnet18(pretrained=True)
            chas = [512, 256, 128, 64, 64, 64]
        else:
            assert False, "The res = {}".format(res)
        down_blocks = []
        up_blocks = []
        self.input_block = nn.Sequential(*list(resnet.children())[:3])
        self.input_pool = list(resnet.children())[3]
        for bottleneck in list(resnet.children()):
            if isinstance(bottleneck, nn.Sequential):
                down_blocks.append(bottleneck)
        self.down_blocks = nn.ModuleList(down_blocks)
        self.bridge = Bridge(chas[0], chas[0])
        up_blocks.append(UpBlockForUNetWithResNet(chas[0], chas[1]))
        up_blocks.append(UpBlockForUNetWithResNet(chas[1], chas[2]))
        up_blocks.append(UpBlockForUNetWithResNet(chas[2], chas[3]))
        up_blocks.append(UpBlockForUNetWithResNet(in_channels=chas[4] + chas[5], out_channels=chas[4],
                                                  up_conv_in_channels=chas[3], up_conv_out_channels=chas[4]))
        up_blocks.append(UpBlockForUNetWithResNet(in_channels=chas[5] + 3, out_channels=chas[5],
                                                  up_conv_in_channels=chas[4], up_conv_out_channels=chas[5]))
        # self.bridge = Bridge(2048, 2048)
        # up_blocks.append(UpBlockForUNetWithResNet50(2048, 1024))
        # up_blocks.append(UpBlockForUNetWithResNet50(1024, 512))
        # up_blocks.append(UpBlockForUNetWithResNet50(512, 256))
        # up_blocks.append(UpBlockForUNetWithResNet50(in_channels=128 + 64, out_channels=128,
        #                                             up_conv_in_channels=256, up_conv_out_channels=128))
        # up_blocks.append(UpBlockForUNetWithResNet50(in_channels=64 + 3, out_channels=64,
        #                                             up_conv_in_channels=128, up_conv_out_channels=64))

        self.up_blocks = nn.ModuleList(up_blocks)

        self.out = nn.Conv2d(chas[5], n_classes, kernel_size=1, stride=1)
        self._initialize_weights(self.up_blocks)
        self._initialize_weights(self.out)

    def _initialize_weights(self, modules=None):
        if modules is None:
            modules = self
        for m in modules.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()

    def forward(self, x, with_output_feature_map=False, encode=False, skip_encode=False):
        pre_pools = dict()
        pre_pools["layer_0"] = x
        x = self.input_block(x)
        pre_pools["layer_1"] = x
        x = self.input_pool(x)

        for i, block in enumerate(self.down_blocks, 2):
            x = block(x)
            if i == (UNetWithResnetEncoder.DEPTH - 1):
                continue
            pre_pools["layer_{}".format(i)] = x

        x = self.bridge(x)
        if encode:
            del pre_pools
            x = x.mean(3, True).mean(2, True)
            return x.view(x.size(0), -1)
        elif skip_encode:
            key = "layer_{}".format(UNetWithResnetEncoder.DEPTH - 1)
            pre_pools[key] = x
            return pre_pools

        for i, block in enumerate(self.up_blocks, 1):
            key = "layer_{}".format(UNetWithResnetEncoder.DEPTH - 1 - i)
            x = block(x, pre_pools[key])
        output_feature_map = x
        x = self.out(x)
        del pre_pools
        if with_output_feature_map:
            return x, output_feature_map
        else:
            return x

    def decode(self, pre_pools, with_output_feature_map=False):
        key = "layer_{}".format(UNetWithResnetEncoder.DEPTH - 1)
        x = pre_pools[key]
        for i, block in enumerate(self.up_blocks, 1):
            key = "layer_{}".format(UNetWithResnetEncoder.DEPTH - 1 - i)
            x = block(x, pre_pools[key])
        output_feature_map = x
        x = self.out(x)
        del pre_pools
        if with_output_feature_map:
            return x, output_feature_map
        else:
            return x

class ResnetDecoder_skip(thisModule):
    def __init__(self, n_classes=3, res='50'):
        # n_classes: output channels
        super(ResnetDecoder_skip, self).__init__()
        res = str(res)
        if res == '50':
            chas = [2048, 1024, 512, 256, 128, 64]
        elif res == '34':
            chas = [512, 256, 128, 64, 64, 64]
        elif res == '18':
            chas = [512, 256, 128, 64, 64, 64]
        else:
            assert False, "The res = {}".format(res)

        up_blocks = []
        up_blocks.append(UpBlockForUNetWithResNet(chas[0], chas[1]))
        up_blocks.append(UpBlockForUNetWithResNet(chas[1], chas[2]))
        up_blocks.append(UpBlockForUNetWithResNet(chas[2], chas[3]))
        up_blocks.append(UpBlockForUNetWithResNet(in_channels=chas[4] + chas[5], out_channels=chas[4],
                                                  up_conv_in_channels=chas[3], up_conv_out_channels=chas[4]))
        up_blocks.append(UpBlockForUNetWithResNet(in_channels=chas[5] + 3, out_channels=chas[5],
                                                  up_conv_in_channels=chas[4], up_conv_out_channels=chas[5]))

        self.up_blocks = nn.ModuleList(up_blocks)

        self.out = nn.Conv2d(chas[5], n_classes, kernel_size=1, stride=1)
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                if m.bias is not None:
                    m.bias.data.zero_()
                if m.bias is not None:
                    nn.init.xavier_uniform(m.weight)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, pre_pools, with_output_feature_map=False):
        key = "layer_{}".format(UNetWithResnetEncoder.DEPTH - 1)
        x = pre_pools[key]
        for i, block in enumerate(self.up_blocks, 1):
            key = "layer_{}".format(UNetWithResnetEncoder.DEPTH - 1 - i)
            x = block(x, pre_pools[key])
        output_feature_map = x
        x = self.out(x)
        del pre_pools
        if with_output_feature_map:
            return x, output_feature_map
        else:
            return x

    def decode(self, pre_pools, with_output_feature_map=False):
        return self.forward(pre_pools, with_output_feature_map=with_output_feature_map)

class Unet_Discriminator(thisModule):
    def __init__(self, config, in_channels=(2048, 192), ema=False, ucnet=False):
        super(Unet_Discriminator, self).__init__()

        print '===> Init small-fc for {}'.format(config.dataset)

        self.num_label = config.num_label
        self.ucnet = ucnet

        # assert len(in_channels) == 2, \
        #     "in_channels needs 2 int in a tuple, now len={}".format(len(in_channels))
        # (n_filter_1, n_filter_2) = in_channels
        out_net = []
        for i in range(len(in_channels)-1):
            out_net.append(WN_Linear(in_channels[i], in_channels[i+1], train_scale=True, init_stdv=0.1))
            out_net.append(nn.LeakyReLU(0.2))
        out_net.append(WN_Linear(in_channels[-1], self.num_label, train_scale=True, init_stdv=0.1))

        self.out_net = nn.Sequential(*out_net)
        if self.ucnet:
            self.uc_net = WN_Linear(in_channels[-1], self.num_label, train_scale=True, init_stdv=0.1)

        if ema:
            for param in self.parameters():
                param.detach_()

    def forward(self, X, uc=False):
        if self.ucnet and uc:
            for i in range(len(self.out_net)-1):
                X = self.out_net[i](X)
            uc = self.uc_net(X)
            X = self.out_net[-1](X)
            return X, uc
        else:
            return self.out_net(X)

class UNet(thisModule):
    def __init__(self, n_channels, n_classes, large=False, upbilinear=True):
        # n_channels: input channels; n_classes: output channels
        super(UNet, self).__init__()    # cifar: 64 => 16
        if large:
            cnum = 32   # min channel num
        else:
            cnum = 64   # min channel num
        self.inc = inconv(n_channels, cnum)
        self.down1 = down(cnum,   cnum*2)
        self.down2 = down(cnum*2, cnum*4)
        self.down3 = down(cnum*4, cnum*8)
        self.down4 = down(cnum*8, cnum*8)
        self.up1 = up(cnum*16, cnum*4, bilinear=upbilinear)
        self.up2 = up(cnum*8,  cnum*2, bilinear=upbilinear)
        self.up3 = up(cnum*4,  cnum, bilinear=upbilinear)
        self.up4 = up(cnum*2,  cnum, bilinear=upbilinear)
        self.outc = WN_Conv2d(cnum, n_classes, 1)
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()

    def forward(self, x):
        x = x * 0.5 + 0.5  # [-1, 1] -> [0, 1]
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        x = self.outc(x)
        x = F.sigmoid(x) * 2. - 1. # [0, 1] -> [-1, 1]
        return x


text:
from preprocess import use_dataset
from model import DeepShallowConfig, DeepShallowModel
from trainer import generate_square_subsequent_mask
from dataset import TargetDataset, custom_collate_inference_fn

import pandas as pd
import numpy as np
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from transformers import PreTrainedTokenizerFast

import torch
from easydict import EasyDict
import yaml

# Read config.yaml file
with open("config.yaml") as infile:
    SAVED_CFG = yaml.load(infile, Loader=yaml.FullLoader)
    CFG = EasyDict(SAVED_CFG["CFG"])

DEVICE = torch.device(
    "cuda:0" if torch.cuda.is_available() and CFG.DEBUG == False else "cpu"
)

korean_tokenizer = PreTrainedTokenizerFast.from_pretrained("snoop2head/Deep-Shallow-Ko")
english_tokenizer = PreTrainedTokenizerFast.from_pretrained(
    "snoop2head/Deep-Shallow-En"
)

# function to generate output sequence using greedy algorithm
# Greedy Search: Greedy search simply selects the word with the highest probability as its next word
# https://huggingface.co/blog/how-to-generate
def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(DEVICE)
    src_mask = src_mask.to(DEVICE)

    memory = model.encode(src, src_mask)
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(DEVICE)
    for i in range(max_len - 1):
        memory = memory.to(DEVICE)
        tgt_mask = (generate_square_subsequent_mask(ys.size(0)).type(torch.bool)).to(
            DEVICE
        )
        out = model.decode(ys, memory, tgt_mask)
        out = out.transpose(0, 1)
        prob = model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.item()

        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        if next_word == korean_tokenizer.eos_token_id:
            break
    return ys

def inference(transformer, src_input_id):
    # find the index where first pad_token_id appears in src_input_id
    num_tokens = src_input_id.shape[0]
    pad_token_id = korean_tokenizer.pad_token_id

    # get index of the last True
    src_mask = (src_input_id != pad_token_id).squeeze()
    num_tokens_without_pad = src_mask.sum().item() - 1

    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
    transformer.eval()
    tgt_tokens = greedy_decode(
        transformer,
        src_input_id,
        src_mask,
        max_len=num_tokens_without_pad,
        start_symbol=korean_tokenizer.cls_token_id,
    ).flatten()
    target_tokenizer = (
        english_tokenizer if CFG.tgt_language == "en" else korean_tokenizer
    )
    result = " ".join(target_tokenizer.convert_ids_to_tokens(tgt_tokens))
    result = result.replace(" ##", "")
    result = result.replace("[UNK]", "")
    result = result.replace("[CLS]", "")
    return result

if __name__ == "__main__":
    _, _, df_test = use_dataset()
    print(df_test.head(2))

    list_inferenced = []
    slice_index = CFG.num_inference_sample - 1

    list_test = df_test[CFG.src_language].tolist()[:slice_index]
    list_answer = df_test[CFG.tgt_language].tolist()[:slice_index]

    test_dataset = TargetDataset(list_test)

    # test_dataloader = DataLoader(
    #     test_dataset,
    #     batch_size=CFG.valid_batch_size,
    #     collate_fn=custom_collate_inference_fn,
    # )

    config = DeepShallowConfig.from_pretrained(CFG.load_model_name)
    transformer = DeepShallowModel.from_pretrained(CFG.load_model_name, config=config)
    transformer = transformer.to(DEVICE)

    for index_num, item in enumerate(tqdm(test_dataset)):
        src_input_id = item.view(-1, 1)
        result = inference(transformer, src_input_id)
        list_inferenced.append(result)
        if index_num == slice_index:
            break

    df_result = pd.DataFrame(
        {
            "input": list_test,
            "prediction": list_inferenced,
            "label": list_answer,
        }
    )
    df_result.to_csv(
        f"./result/inference_{CFG.num_inference_sample}_samples_{CFG.src_language}_to_{CFG.tgt_language}.csv",
        index=False,
    )


text:
from numpy.core.fromnumeric import size
import pywt
import matplotlib.pyplot as plt
import matplotlib.image as mpimage
import glob
import os
from torch.utils.data import DataLoader
import numpy as np
import torch
def DWT(x_batch):
    x_batch_np = x_batch.detach().cpu().numpy()
    LL_list, LH_list, HL_list, HH_list = [], [], [], []
    for _ in range(len(x_batch_np)): # len(x_batch_np) finds out the number of images in the dataset

        imgs = [mpimage.imread(file) for file in x_batch] #Reading all the images in a particular subdirectory, say, 1 in T1fusion
        tensor_imgs = torch.tensor(imgs)
        for img in tensor_imgs:

            coeffs2 = pywt.dwt2(img, 'db3')
            LL, (LH, HL, HH) = coeffs2
            LL_list.append(LL)
            LH_list.append(LH)
            HL_list.append(HL)
            HH_list.append(HH)

    LL_list = np.array(LL_list)
    LH_list = np.array(LH_list)
    HL_list = np.array(HL_list)
    HH_list = np.array(HH_list)

    return [LL_list, LH_list, HL_list, HH_list]

def IDWT(output_LL, output_LH, output_HL, output_HH):
    output_LL_np = output_LL.detach().cpu().numpy()
    output_LH_np = output_LH.detach().cpu().numpy()
    output_HL_np = output_HL.detach().cpu().numpy()
    output_HH_np = output_HH.detach().cpu().numpy()

    coeff2 = output_LL_np, (output_LH_np, output_HL_np, output_HH_np)
    idwt_fig = pywt.idwt2(coeff2, 'db3')
    return idwt_fig

# new_path = os.path.join(parent_dir, img_folder, sub_directory)
# os.makedirs(new_path)
# for j, a in enumerate([LL, LH, HL, HH]):
#     plt.figure(figsize=(1, 1), dpi=256)

#     plt.imshow(a, interpolation='nearest', cmap=plt.cm.gray)
#     plt.axis('off')
#     plt.tight_layout(pad=0)
#     plt.savefig(new_path +'{}_{}_{}'.format(k+1, i+1, j+1)  +'.png')

#     plt.show()

reason: import numpy,from numpy


text:
#!/usr/bin/python

'''
    mowbot_camera_node.py
    Shaun Bowman
    Aug 9 2018

    openCV / cvbridge ROS implementatoin for monocular vision
    mowbot project
    based on RACECAR/J & MIT RACECAR
    '''

import rospy
import cv2
import numpy as np
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image
from mowbot_line_find.mowbot_line_find import MowbotLineFind
from mowbot_birds_eye.mowbot_birds_eye import MowbotBirdsEye

class MowbotCameraNode:
    myVar = 0
    debug_find_line = True
    debug = False
    debug_one_frame = False
    debug_show_grid = False
    debug_show_birdseye = False
    publish_birdseye = True
    _frame_count = 0
    _image = []
    _frame_stop = 3
    #_image_path = "/home/shaun/catkin_ws/src/mowbot_camera/scripts/experimental/"
    _image_path = "/home/nvidia/racecar-ws/src/mowbot_camera/scripts/experimental/"
    _file_name = "one_image.png"

    line_find = MowbotLineFind()
    birds_eye = MowbotBirdsEye()

    def __init__(self):
        self.bridge = CvBridge()
        self.image_sub = rospy.Subscriber("/usb_cam/image_raw",Image,self.callback)
        self.image_pub = rospy.Publisher("imageBirdsEye", Image, queue_size=10)

    def callback(self,data):

        if self._frame_count < self._frame_stop:
            try:
                cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
            except CvBridgeError as e:
                print(e)

            self.birds_eye.get_birdseye(cv_image)
            self.line_find.updade_img_cur(self.birds_eye.imgBirdsEye)

            if self.debug_find_line:
                self.line_find.find_lines()

            if self.debug:
                cv2.imshow("Image", cv_image)
                cv2.waitKey(1)

            if self.debug_one_frame:
                self._frame_count = self._frame_count + 1
                self._image = cv_image

            if self.debug_show_grid:
                self.line_find.draw_debug_grid()

            if self.debug_show_birdseye:
                cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
                cv2.imshow("Image", self.birds_eye.imgBirdsEye)
                cv2.waitKey(1)

            if self.publish_birdseye:
                try:
                    self.image_pub.publish(self.bridge.cv2_to_imgmsg(self.birds_eye.imgBirdsEye, "bgr8"))
                except CvBridgeError as e:
                    print(e)

        else:

            try:
                cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
            except CvBridgeError as e:
                print(e)

            if self._frame_count == self._frame_stop:
                cv2.imwrite(self._image_path + self._file_name, self._image)
                self._frame_count = self._frame_count + 1

            if self.debug_show_grid:
                self.line_find.draw_debug_grid()

            else:
                self._image = cv_image
                cv2.namedWindow('image', cv2.WINDOW_NORMAL)
                cv2.resizeWindow('image', 960, 540)
                cv2.imshow('image', self._image)
                cv2.waitKey(1)
            pass

if __name__ == "__main__":
    rospy.init_node("mowbot_camera")
    node = MowbotCameraNode()
    rospy.spin()

    try:
        rospy.spin()
    except KeyboardInterrupt:
        print("Shutting down")
    cv2.destroyAllWindows()


text:
##########################################################
# @author: Najeeb Khan                                   #
# @github: najeeb97khan                                  #
# @email: najeeb.khan96@gmail.com                        #
# @version: 1.0.0, 20/01/2018                            #
# @license: MIT                                          #
##########################################################

# Importing Libraries
from __future__ import print_function
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
from sklearn.model_selection import train_test_split
from random import shuffle

def process_data(path_to_file, num_train, num_val):

    mnist = input_data.read_data_sets(path_to_file, one_hot=True)
    X_train, X_val, y_train, y_val = train_test_split(mnist.train.images, mnist.train.labels, random_state=7, test_size=0.3)

    # Shuffling the data
    idx_list = range(num_train)
    shuffle(idx_list)
    X_train, y_train = X_train[idx_list, :], y_train[idx_list, :]
    idx_list = range(num_val)
    shuffle(idx_list)
    X_val, y_val = X_val[idx_list, :], y_val[idx_list, :]
    print('Shape of Training Data: {}'.format(X_train.shape))
    print('Shape of Validation Data: {}'.format(X_val.shape))

    # Normalising the Data
    print('='*25+' Data Characteristics '+'='*25)
    print('Mean Before Normalisation: Training: {}, Validation: {}'.format(np.round(np.mean(X_train), 2), np.round(np.mean(X_val), 2)))
    print('Standard Deviation Before Normalisation: Training: {}, Validation: {}'.format(np.round(np.std(X_train), 2), np.round(np.std(X_val), 2)))

    mean = np.mean(X_train)
    std = np.std(X_train)

    X_train -= mean
    X_train /= std

    X_val -= mean
    X_val /= std

    print('Mean After Normalisation: Training: {}, Validation: {}'.format(np.mean(X_train), np.mean(X_val)))
    print('Standard Deviation After Normalisation: Training: {}, Validation: {}'.format(np.std(X_train), np.std(X_val)))

    return X_train, y_train, X_val, y_val


text:
#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#

import sys
import traceback

import torch
import copy
import gc
import time
import shutil
from typing import Dict
from torch.cuda.amp import autocast
from torch.nn import functional as F
import random
from typing import Union, List, Optional
import numpy as np
from itertools import product

from data.transforms.image_torch import RandomMixup, RandomCutmix
from engine.utils import print_summary
from metrics import Statistics, metric_monitor
from common import DEFAULT_ITERATIONS, DEFAULT_EPOCHS, DEFAULT_LOG_FREQ
from utils import logger
from utils.common_utils import create_directories, move_to_device
from utils.ddp_utils import is_master, dist_barrier
from utils.tensor_utils import reduce_tensor_sum, tensor_to_python_float
from utils.checkpoint_utils import copy_weights, save_checkpoint
from loss_landscape import landscape_utils as ll_utils

class Trainer(object):
    """
    This class defines the training and validation code for training models with CVNets
    """

    def __init__(
        self,
        opts,
        model,
        validation_loader,
        training_loader,
        criterion,
        optimizer,
        scheduler,
        gradient_scalar,
        start_epoch: int = 0,
        start_iteration: int = 0,
        best_metric: float = 0.0,
        model_ema=None,
        *args,
        **kwargs
    ) -> None:
        super(Trainer, self).__init__()

        self.opts = opts

        self.model = model
        self.model_ema = model_ema
        self.criteria = criterion
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.gradient_scalar = gradient_scalar

        self.val_loader = validation_loader
        self.train_loader = training_loader

        self.device = getattr(opts, "dev.device", torch.device("cpu"))

        self.start_epoch = start_epoch
        self.best_metric = best_metric
        self.train_iterations = start_iteration

        self.is_master_node = is_master(opts)
        self.max_iterations_reached = False
        self.max_iterations = getattr(
            self.opts, "scheduler.max_iterations", DEFAULT_ITERATIONS
        )
        self.use_distributed = getattr(self.opts, "ddp.use_distributed", False)
        self.log_freq = getattr(self.opts, "common.log_freq", DEFAULT_LOG_FREQ)
        self.accum_freq = getattr(self.opts, "common.accum_freq", 1)
        self.accum_after_epoch = getattr(self.opts, "common.accum_after_epoch", 0)

        self.mixed_precision_training = getattr(opts, "common.mixed_precision", False)

        self.train_metric_names = getattr(opts, "stats.train", ["loss"])
        if isinstance(self.train_metric_names, str):
            self.train_metric_names = [self.train_metric_names]

        assert isinstance(
            self.train_metric_names, list
        ), "Type of metric names should be list. Got: {}".format(
            type(self.train_metric_names)
        )

        if "loss" not in self.train_metric_names:
            self.train_metric_names.append(self.train_metric_names)

        self.val_metric_names = getattr(opts, "stats.val", ["loss"])
        if isinstance(self.val_metric_names, str):
            self.val_metric_names = [self.val_metric_names]

        assert isinstance(
            self.val_metric_names, list
        ), "Type of metric names should be list. Got: {}".format(
            type(self.val_metric_names)
        )

        if "loss" not in self.val_metric_names:
            self.val_metric_names.append(self.val_metric_names)

        self.save_all_checkpoints = getattr(
            self.opts, "stats.save_all_checkpoints", False
        )
        self.ckpt_metric = getattr(self.opts, "stats.checkpoint_metric", "loss")
        if self.ckpt_metric is None:
            # if checkpoint metric is not specified, then use loss
            self.ckpt_metric = "loss"

        assert (
            self.ckpt_metric in self.val_metric_names
        ), "Checkpoint metric should be part of metric names. Metric names: {}, Checkpoint metric: {}".format(
            self.val_metric_names, self.ckpt_metric
        )
        self.ckpt_metric = self.ckpt_metric.lower()

        self.tb_log_writer = None
        self.bolt_log_writer = None
        if self.is_master_node:
            self.setup_log_writer()

            print_summary(
                opts=self.opts,
                model=self.model,
                criteria=self.criteria,
                optimizer=self.optimizer,
                scheduler=self.scheduler,
            )

        self.adjust_norm_mom = None
        if getattr(opts, "model.normalization.adjust_bn_momentum.enable", False):
            from cvnets.layers import AdjustBatchNormMomentum

            self.adjust_norm_mom = AdjustBatchNormMomentum(opts=opts)
            if self.is_master_node:
                logger.log(
                    "Batch normalization momentum will be annealed during training."
                )
                print(self.adjust_norm_mom)

        # sample-efficient training
        self.cache_dict = None
        self.sample_efficient_training = getattr(
            opts, "dataset.sample_efficient_training.enable", False
        )
        self.sample_confidence = getattr(
            opts, "dataset.sample_efficient_training.sample_confidence", 0.5
        )
        self.find_easy_samples_every_k_epoch = getattr(
            opts,
            "dataset.sample_efficient_training.find_easy_samples_every_k_epochs",
            5,
        )
        self.min_sample_frequency = getattr(
            opts, "dataset.sample_efficient_training.min_sample_frequency", 5
        )
        if self.sample_efficient_training:
            self.train_loader_set = copy.deepcopy(self.train_loader)
            self.sample_ids_orig = self.train_loader_set.get_sample_indices()
            n_samples = len(self.sample_ids_orig)
            self.running_sum_tensor = torch.zeros(
                (n_samples,), device=self.device, dtype=torch.int
            )
            self.running_sum_tensor.requires_grad = False
            if self.is_master_node:
                logger.log("Configuring for sample efficient training")

        # recent versions of PyTorch support setting grads to None, for better performance
        # To be explored in Future
        # self.optimizer.zero_grad(set_to_none=True)
        self.set_grad_to_none = False

    def setup_log_writer(self):
        tensorboard_logging = getattr(self.opts, "common.tensorboard_logging", False)
        if tensorboard_logging:
            try:
                from torch.utils.tensorboard import SummaryWriter
            except ImportError as e:
                logger.log(
                    "Unable to import SummaryWriter from torch.utils.tensorboard. Disabling tensorboard logging"
                )
                SummaryWriter = None

            if SummaryWriter is not None:
                exp_dir = getattr(self.opts, "common.exp_loc", "results/run_1")
                exp_dir = "{}/tb_logs".format(exp_dir)
                create_directories(dir_path=exp_dir, is_master_node=self.is_master_node)
                self.tb_log_writer = SummaryWriter(
                    log_dir=exp_dir, comment="Training and Validation logs"
                )
            else:
                self.tb_log_writer = None

        bolt_logging = getattr(self.opts, "common.bolt_logging", False)
        if bolt_logging:
            try:
                from utils.bolt_logger import BoltLogger
            except ModuleNotFoundError:
                BoltLogger = None

            if BoltLogger is None:
                logger.log("Unable to import bolt. Disabling bolt logging")
                self.bolt_log_writer = None
            else:
                self.bolt_log_writer = BoltLogger()

    def compute_grad_norm(self):
        parameters = [p for p in self.model.parameters() if p.grad is not None]
        if len(parameters) == 0:
            return None

        norm_type = 2.0  # L2 norm

        inv_scale = 1.0 / self.gradient_scalar.get_scale()
        total_norm = torch.norm(
            torch.stack(
                [
                    torch.norm(p.grad.detach() * inv_scale, norm_type).to(self.device)
                    for p in parameters
                ]
            ),
            norm_type,
        )
        if total_norm.isnan() or total_norm.isinf():
            return None
        return total_norm

    def _get_batch_size(self, x):
        if isinstance(x, torch.Tensor):
            return x.shape[0]
        elif isinstance(x, Dict):
            return x["image"].shape[0]

    def apply_mixup_transforms(self, data):
        # Apply mixup transforms on classification tasks
        opts = self.opts
        mixup_transforms = []
        if getattr(opts, "image_augmentation.mixup.enable", False):
            n_classes = getattr(opts, "model.classification.n_classes", None)
            if n_classes is None:
                logger.error("Please specify number of classes. Got None.")
            mixup_transforms.append(RandomMixup(opts=opts, num_classes=n_classes))

        if getattr(opts, "image_augmentation.cutmix.enable", False):
            n_classes = getattr(opts, "model.classification.n_classes", None)
            if n_classes is None:
                logger.error("Please specify number of classes. Got None.")
            mixup_transforms.append(RandomCutmix(opts=opts, num_classes=n_classes))

        if len(mixup_transforms) > 0:
            _mixup_transform = random.choice(mixup_transforms)
            data = _mixup_transform(data)
        return data

    def _zero_grad(self):
        if self.set_grad_to_none:
            self.optimizer.zero_grad(set_to_none=True)
        else:
            self.optimizer.zero_grad()

    def train_epoch(self, epoch):
        time.sleep(2)  # To prevent possible deadlock during epoch transition

        if self.is_master_node:
            logger.double_dash_line()
            logger.debug(
                "Training epoch {} with {} samples".format(
                    epoch, self.train_loader.samples_in_dataset()
                )
            )

        train_stats = Statistics(
            metric_names=self.train_metric_names, is_master_node=self.is_master_node
        )

        self.model.train()
        accum_freq = self.accum_freq if epoch >= self.accum_after_epoch else 1
        max_norm = getattr(self.opts, "common.grad_clip", None)

        # set the gradient to zero or None
        self._zero_grad()

        epoch_start_time = time.time()
        batch_load_start = time.time()
        grad_norm = 0.0
        for batch_id, batch in enumerate(self.train_loader):
            if self.train_iterations > self.max_iterations:
                self.max_iterations_reached = True
                return -1, -1

            # move to device
            batch = move_to_device(opts=self.opts, x=batch, device=self.device)
            # apply mix-up transforms if any
            batch = self.apply_mixup_transforms(data=batch)

            batch_load_toc = time.time() - batch_load_start

            input_img, target_label = batch["image"], batch["label"]

            batch_size = self._get_batch_size(input_img)

            # update the learning rate
            self.optimizer = self.scheduler.update_lr(
                optimizer=self.optimizer, epoch=epoch, curr_iter=self.train_iterations
            )

            # adjust bn momentum
            if self.adjust_norm_mom is not None:
                self.adjust_norm_mom.adjust_momentum(
                    model=self.model, epoch=epoch, iteration=self.train_iterations
                )

            with autocast(enabled=self.mixed_precision_training):
                # prediction
                pred_label = self.model(input_img)
                # compute loss
                loss = self.criteria(
                    input_sample=input_img, prediction=pred_label, target=target_label
                )

                if isinstance(loss, torch.Tensor) and torch.isnan(loss):
                    logger.error("Nan encountered in the loss.")

            # perform the backward pass with gradient accumulation [Optional]
            self.gradient_scalar.scale(loss).backward()

            if (batch_id + 1) % accum_freq == 0:
                if max_norm is not None:
                    # For gradient clipping, unscale the gradients and then clip them
                    self.gradient_scalar.unscale_(self.optimizer)
                    torch.nn.utils.clip_grad_norm_(
                        self.model.parameters(), max_norm=max_norm
                    )

                if "grad_norm" in self.train_metric_names:
                    # compute grad_norm for logging purposes.
                    # We can't use the output of clip_grad_norm_ because it returns the total norm before clipping
                    grad_norm = self.compute_grad_norm()

                # optimizer step
                self.gradient_scalar.step(optimizer=self.optimizer)
                # update the scale for next batch
                self.gradient_scalar.update()
                # set the gradient to zero or None
                self._zero_grad()

                self.train_iterations += 1

                if self.model_ema is not None:
                    self.model_ema.update_parameters(self.model)

            metrics = metric_monitor(
                self.opts,
                pred_label=pred_label,
                target_label=target_label,
                loss=loss,
                grad_norm=grad_norm,
                use_distributed=self.use_distributed,
                metric_names=self.train_metric_names,
            )

            train_stats.update(
                metric_vals=metrics, batch_time=batch_load_toc, n=batch_size
            )

            if batch_id % self.log_freq == 0 and self.is_master_node:
                lr = self.scheduler.retrieve_lr(self.optimizer)
                train_stats.iter_summary(
                    epoch=epoch,
                    n_processed_samples=self.train_iterations,
                    total_samples=self.max_iterations,
                    learning_rate=lr,
                    elapsed_time=epoch_start_time,
                )

            batch_load_start = time.time()

        avg_loss = train_stats.avg_statistics(metric_name="loss")
        train_stats.epoch_summary(epoch=epoch, stage="training")
        avg_ckpt_metric = train_stats.avg_statistics(metric_name=self.ckpt_metric)

        gc.collect()

        return avg_loss, avg_ckpt_metric

    def val_epoch(self, epoch, model, extra_str=""):
        if self.val_loader is None:
            return 0.0, 0.0

        time.sleep(2)  # To prevent possible deadlock during epoch transition
        validation_stats = Statistics(
            metric_names=self.val_metric_names, is_master_node=self.is_master_node
        )

        if "coco_map" in self.val_metric_names:
            from metrics.coco_map import COCOEvaluator

            coco_evaluator = COCOEvaluator(
                opts=self.opts, iou_types=["bbox"], use_distributed=self.use_distributed
            )
        else:
            coco_evaluator = None

        model.eval()
        if model.training and self.is_master_node:
            logger.warning("Model is in training mode. Switching to evaluation mode")
            model.eval()

        with torch.no_grad():
            epoch_start_time = time.time()
            total_samples = len(self.val_loader)
            processed_samples = 0
            lr = self.scheduler.retrieve_lr(self.optimizer)
            for batch_id, batch in enumerate(self.val_loader):
                batch = move_to_device(opts=self.opts, x=batch, device=self.device)

                input_img, target_label = batch["image"], batch["label"]

                batch_size = self._get_batch_size(input_img)

                with autocast(enabled=self.mixed_precision_training):
                    # prediction
                    pred_label = model(input_img)
                    # compute loss
                    loss = self.criteria(
                        input_sample=input_img,
                        prediction=pred_label,
                        target=target_label,
                    )

                processed_samples += batch_size

                metrics = metric_monitor(
                    self.opts,
                    pred_label=pred_label,
                    target_label=target_label,
                    loss=loss,
                    use_distributed=self.use_distributed,
                    metric_names=self.val_metric_names,
                    is_evaluation=True,
                )

                validation_stats.update(
                    metric_vals=metrics, batch_time=0.0, n=batch_size
                )

                if coco_evaluator is not None:
                    coco_evaluator.prepare_predictions(
                        predictions=pred_label, targets=target_label
                    )

                if batch_id % self.log_freq == 0 and self.is_master_node:
                    validation_stats.iter_summary(
                        epoch=epoch,
                        n_processed_samples=processed_samples,
                        total_samples=total_samples,
                        elapsed_time=epoch_start_time,
                        learning_rate=lr,
                    )

        validation_stats.epoch_summary(epoch=epoch, stage="validation" + extra_str)
        avg_loss = validation_stats.avg_statistics(metric_name="loss")
        avg_ckpt_metric = validation_stats.avg_statistics(metric_name=self.ckpt_metric)

        if coco_evaluator is not None:
            # synchronize across different processes and aggregate the results
            coco_evaluator.gather_coco_results()
            coco_map = coco_evaluator.summarize_coco_results()

            if self.ckpt_metric == "coco_map" and "bbox" in coco_map:
                avg_ckpt_metric = round(coco_map["bbox"], 5)

        if avg_ckpt_metric is None:
            avg_ckpt_metric = avg_loss

        gc.collect()

        return avg_loss, avg_ckpt_metric

    def find_easy_samples(self, epoch, model, *args, **kwargs):
        """
        This function identifies easy samples in the training set and removes them from training.

        .. note::
            Currently, this is implemented separately to avoid breaking the training and validation pipeline. In future,
            this will be combined with main training loop to reduce overhead.
        """

        time.sleep(2)  # To prevent possible deadlock during epoch transition

        model.eval()
        if model.training and self.is_master_node:
            logger.warning("Model is in training mode. Switching to evaluation mode")
            model.eval()

        if self.is_master_node:
            logger.log("Trying to find easy samples in epoch {}".format(epoch))

        with torch.no_grad():
            easy_sample_ids_tensor = torch.zeros_like(self.running_sum_tensor)

            for batch_id, batch in enumerate(self.train_loader_set):
                batch = move_to_device(opts=self.opts, x=batch, device=self.device)

                input_img, target_label = batch["image"], batch["label"]

                sample_ids = None
                if "sample_id" in batch:
                    sample_ids = batch["sample_id"]
                else:
                    self.sample_efficient_training = False
                    if self.is_master_node:
                        logger.log(
                            "Sample Ids are required in a batch for sample efficient training. "
                            "sample_id key not found in batch. Disabling sample efficient training."
                        )
                    break

                if sample_ids is None:
                    logger.log("Sample Ids can't be none")
                    break

                with autocast(enabled=self.mixed_precision_training):
                    # prediction
                    pred_label = model(input_img)
                    pred_label = F.softmax(pred_label, dim=-1)

                pred_conf, pred_indices = torch.max(pred_label, dim=-1)

                easy_samples = torch.logical_and(
                    pred_indices.eq(
                        target_label
                    ),  # condition 1: Predicted label == Target label
                    pred_conf
                    >= self.sample_confidence,  # condition 2: prediction confidence >= desired confidence
                )

                if easy_samples.numel() > 0:
                    easy_sample_ids = sample_ids[easy_samples]
                    # find easy samples as per condition 1 and 2 and set their values to 1
                    easy_sample_ids_tensor[easy_sample_ids] = 1

            # synchronize tensors
            if self.use_distributed:
                # sync across all GPUs.
                easy_sample_ids_tensor = reduce_tensor_sum(easy_sample_ids_tensor)

            # some samples which are classified easy earlier may have been classified hard now.
            easy_sample_ids_tensor[easy_sample_ids_tensor == 0] = -1

            if self.is_master_node:
                logger.debug(
                    "Number of easy samples found during epoch {} are {}".format(
                        epoch,
                        easy_sample_ids_tensor[easy_sample_ids_tensor > 0].sum().item(),
                    )
                )

            self.running_sum_tensor = torch.clip(
                self.running_sum_tensor + easy_sample_ids_tensor,
                min=0,
                max=self.min_sample_frequency,
            )

            if self.running_sum_tensor.sum() > 0:
                skip_sample_ids = (
                    self.running_sum_tensor >= self.min_sample_frequency
                ).nonzero(as_tuple=True)[0]

                if skip_sample_ids.numel() > 0:
                    skip_samples = skip_sample_ids.cpu().numpy().tolist()

                    new_sample_ids = [
                        s_id
                        for s_id in self.sample_ids_orig
                        if s_id not in skip_sample_ids
                    ]

                    # update the train loader indices
                    self.train_loader.update_indices(new_sample_ids)

                    if self.is_master_node:
                        logger.debug(
                            "Number of samples to skip after epoch {} are {}".format(
                                epoch, len(skip_samples)
                            )
                        )

    @staticmethod
    def log_metrics(
        lrs: Union[List, float],
        log_writer,
        train_loss: float,
        val_loss: float,
        epoch: int,
        best_metric: float,
        val_ema_loss: Optional[float] = None,
        ckpt_metric_name: Optional[str] = None,
        train_ckpt_metric: Optional[float] = None,
        val_ckpt_metric: Optional[float] = None,
        val_ema_ckpt_metric: Optional[float] = None,
    ) -> None:
        if not isinstance(lrs, list):
            lrs = [lrs]
        for g_id, lr_val in enumerate(lrs):
            log_writer.add_scalar("LR/Group-{}".format(g_id), round(lr_val, 6), epoch)

        log_writer.add_scalar("Train/Loss", round(train_loss, 2), epoch)
        log_writer.add_scalar("Val/Loss", round(val_loss, 2), epoch)
        log_writer.add_scalar("Common/Best Metric", round(best_metric, 2), epoch)
        if val_ema_loss is not None:
            log_writer.add_scalar("Val_EMA/Loss", round(val_ema_loss, 2), epoch)

        # If val checkpoint metric is different from loss, add that too
        if ckpt_metric_name is not None and ckpt_metric_name != "loss":
            if train_ckpt_metric is not None:
                log_writer.add_scalar(
                    "Train/{}".format(ckpt_metric_name.title()),
                    round(train_ckpt_metric, 2),
                    epoch,
                )
            if val_ckpt_metric is not None:
                log_writer.add_scalar(
                    "Val/{}".format(ckpt_metric_name.title()),
                    round(val_ckpt_metric, 2),
                    epoch,
                )
            if val_ema_ckpt_metric is not None:
                log_writer.add_scalar(
                    "Val_EMA/{}".format(ckpt_metric_name.title()),
                    round(val_ema_ckpt_metric, 2),
                    epoch,
                )

    def run(self, train_sampler=None):
        if train_sampler is None and self.is_master_node:
            logger.error("Train sampler cannot be None")

        copy_at_epoch = getattr(self.opts, "ema.copy_at_epoch", -1)
        train_start_time = time.time()
        save_dir = getattr(self.opts, "common.exp_loc", "results")

        cfg_file = getattr(self.opts, "common.config_file", None)
        if cfg_file is not None and self.is_master_node:
            dst_cfg_file = "{}/config.yaml".format(save_dir)
            shutil.copy(src=cfg_file, dst=dst_cfg_file)
            logger.info(
                "Configuration file is stored here: {}".format(
                    logger.color_text(dst_cfg_file)
                )
            )

        keep_k_best_ckpts = getattr(self.opts, "common.k_best_checkpoints", 5)
        ema_best_metric = self.best_metric
        is_ema_best = False

        try:
            max_epochs = getattr(self.opts, "scheduler.max_epochs", DEFAULT_EPOCHS)
            max_checkpoint_metric = getattr(
                self.opts, "stats.checkpoint_metric_max", False
            )
            for epoch in range(self.start_epoch, max_epochs):
                # Note that we are using our owm implementations of data samplers
                # and we have defined this function for both distributed and non-distributed cases
                train_sampler.set_epoch(epoch)
                train_sampler.update_scales(
                    epoch=epoch, is_master_node=self.is_master_node
                )

                train_loss, train_ckpt_metric = self.train_epoch(epoch)

                val_loss, val_ckpt_metric = self.val_epoch(
                    epoch=epoch, model=self.model
                )

                if epoch == copy_at_epoch and self.model_ema is not None:
                    if self.is_master_node:
                        logger.log("Copying EMA weights")
                    # copy model_src weights to model_tgt
                    self.model = copy_weights(
                        model_tgt=self.model, model_src=self.model_ema
                    )
                    if self.is_master_node:
                        logger.log("EMA weights copied")
                        logger.log("Running validation after Copying EMA model weights")
                    self.val_epoch(epoch=epoch, model=self.model)

                if max_checkpoint_metric:
                    is_best = val_ckpt_metric >= self.best_metric
                    self.best_metric = max(val_ckpt_metric, self.best_metric)
                else:
                    is_best = val_ckpt_metric <= self.best_metric
                    self.best_metric = min(val_ckpt_metric, self.best_metric)

                val_ema_loss = None
                val_ema_ckpt_metric = None
                if self.model_ema is not None:
                    val_ema_loss, val_ema_ckpt_metric = self.val_epoch(
                        epoch=epoch, model=self.model_ema.ema_model, extra_str=" (EMA)"
                    )
                    if max_checkpoint_metric:
                        is_ema_best = val_ema_ckpt_metric >= ema_best_metric
                        ema_best_metric = max(val_ema_ckpt_metric, ema_best_metric)
                    else:
                        is_ema_best = val_ema_ckpt_metric <= ema_best_metric
                        ema_best_metric = min(val_ema_ckpt_metric, ema_best_metric)

                # sample efficient training
                if (
                    self.sample_efficient_training
                    and (epoch + 1) % self.find_easy_samples_every_k_epoch == 0
                ):
                    self.find_easy_samples(
                        epoch=epoch,
                        model=self.model
                        if self.model_ema is not None
                        else self.model_ema.ema_model,
                    )

                gc.collect()

                if self.is_master_node:
                    save_checkpoint(
                        iterations=self.train_iterations,
                        epoch=epoch,
                        model=self.model,
                        optimizer=self.optimizer,
                        best_metric=self.best_metric,
                        is_best=is_best,
                        save_dir=save_dir,
                        model_ema=self.model_ema,
                        is_ema_best=is_ema_best,
                        ema_best_metric=ema_best_metric,
                        gradient_scalar=self.gradient_scalar,
                        max_ckpt_metric=max_checkpoint_metric,
                        k_best_checkpoints=keep_k_best_ckpts,
                        save_all_checkpoints=self.save_all_checkpoints,
                    )
                    logger.info(
                        "Checkpoints saved at: {}".format(save_dir), print_line=True
                    )

                if self.is_master_node:
                    lr_list = self.scheduler.retrieve_lr(self.optimizer)

                    if self.tb_log_writer is not None:
                        self.log_metrics(
                            lrs=lr_list,
                            log_writer=self.tb_log_writer,
                            train_loss=train_loss,
                            val_loss=val_loss,
                            epoch=epoch,
                            best_metric=self.best_metric,
                            val_ema_loss=val_ema_loss,
                            ckpt_metric_name=self.ckpt_metric,
                            train_ckpt_metric=train_ckpt_metric,
                            val_ckpt_metric=val_ckpt_metric,
                            val_ema_ckpt_metric=val_ema_ckpt_metric,
                        )
                    if self.bolt_log_writer is not None:
                        self.log_metrics(
                            lrs=lr_list,
                            log_writer=self.bolt_log_writer,
                            train_loss=train_loss,
                            val_loss=val_loss,
                            epoch=epoch,
                            best_metric=self.best_metric,
                            val_ema_loss=val_ema_loss,
                            ckpt_metric_name=self.ckpt_metric,
                            train_ckpt_metric=train_ckpt_metric,
                            val_ckpt_metric=val_ckpt_metric,
                            val_ema_ckpt_metric=val_ema_ckpt_metric,
                        )

                if self.max_iterations_reached:
                    if self.use_distributed:
                        dist_barrier()

                    if self.is_master_node:
                        logger.info("Max. iterations for training reached")
                    break

        except KeyboardInterrupt as e:
            if self.is_master_node:
                logger.log("Keyboard interruption. Exiting from early training")
                raise e
        except Exception as e:
            if "out of memory" in str(e):
                logger.log("OOM exception occured")
                n_gpus = getattr(self.opts, "dev.num_gpus", 1)
                for dev_id in range(n_gpus):
                    mem_summary = torch.cuda.memory_summary(
                        device=torch.device("cuda:{}".format(dev_id)), abbreviated=True
                    )
                    logger.log("Memory summary for device id: {}".format(dev_id))
                    print(mem_summary)
            else:
                logger.log(
                    "Exception occurred that interrupted the training. {}".format(
                        str(e)
                    )
                )
                print(e)
                traceback.print_exc()
                raise e
        finally:
            use_distributed = getattr(self.opts, "ddp.use_distributed", False)
            if use_distributed:
                torch.distributed.destroy_process_group()

            torch.cuda.empty_cache()

            if self.is_master_node and self.tb_log_writer is not None:
                self.tb_log_writer.close()

            if self.is_master_node:
                train_end_time = time.time()
                hours, rem = divmod(train_end_time - train_start_time, 3600)
                minutes, seconds = divmod(rem, 60)
                train_time_str = "{:0>2}:{:0>2}:{:05.2f}".format(
                    int(hours), int(minutes), seconds
                )
                logger.log("Training took {}".format(train_time_str))

            try:
                exit(0)
            except Exception as e:
                pass
            finally:
                pass

    def run_loss_landscape(self):
        # Loss landscape code is adapted from https://github.com/xxxnell/how-do-vits-work
        ll_start_time = time.time()
        try:
            n_points = getattr(self.opts, "loss_landscape.n_points", 32)
            min_x = getattr(self.opts, "loss_landscape.min_x", -1.0)
            max_x = getattr(self.opts, "loss_landscape.max_x", 1.0)
            min_y = getattr(self.opts, "loss_landscape.min_y", -1.0)
            max_y = getattr(self.opts, "loss_landscape.max_y", 1.0)

            if self.is_master_node:
                logger.log(
                    "Loss landscape coord space params: \n\tmin_x={}\n\tmax_x={}\n\tmin_y={}\n\tmax_y={}\n\tn_points={}".format(
                        min_x, max_x, min_y, max_y, n_points
                    )
                )

            ll_metrics = ["loss"]
            ll_stats = Statistics(
                metric_names=ll_metrics, is_master_node=self.is_master_node
            )
            has_module = hasattr(self.model, "module")
            model_name = (
                self.model.module.__class__.__name__
                if has_module
                else self.model.__class__.__name__
            )
            save_dir = getattr(self.opts, "common.exp_loc", "results")

            # copy the model and create bases
            model = copy.deepcopy(self.model)
            weight_state_0 = (
                copy.deepcopy(model.module.state_dict())
                if has_module
                else copy.deepcopy(model.state_dict())
            )
            bases = ll_utils.create_bases(
                model=model, device=self.device, has_module=has_module
            )

            xs = np.linspace(min_x, max_x, n_points)
            ys = np.linspace(min_y, max_y, n_points)

            grid_a, grid_b = np.meshgrid(xs, ys, indexing="xy")
            loss_surface = np.empty_like(grid_a)

            epoch = -1
            for coord_a, coord_b in product(range(n_points), range(n_points)):
                epoch += 1
                coords_list = [grid_a[coord_a, coord_b], grid_b[coord_a, coord_b]]
                weight_state_1 = copy.deepcopy(weight_state_0)
                gs = [{k: r * bs[k] for k in bs} for r, bs in zip(coords_list, bases)]
                gs = {
                    k: torch.sum(torch.stack([g[k] for g in gs]), dim=0)
                    + weight_state_1[k]
                    for k in gs[0]
                }

                # load the weights
                model.module.load_state_dict(
                    gs
                ) if has_module else model.load_state_dict(gs)

                model = model.to(device=self.device)
                model.eval()

                total_samples = len(self.val_loader)
                with torch.no_grad():
                    epoch_start_time = time.time()
                    processed_samples = 0
                    for batch_id, batch in enumerate(self.val_loader):
                        batch = move_to_device(
                            opts=self.opts, x=batch, device=self.device
                        )
                        input_img, target_label = batch["image"], batch["label"]

                        batch_size = self._get_batch_size(x=input_img)
                        processed_samples += batch_size

                        # make the prediction and compute loss
                        pred_label = model(input_img)
                        loss = self.criteria(
                            input_sample=input_img,
                            prediction=pred_label,
                            target=target_label,
                        )

                        if isinstance(loss, torch.Tensor) and torch.isnan(loss):
                            logger.error("Nan encountered in the loss.")

                        metrics = metric_monitor(
                            self.opts,
                            pred_label=pred_label,
                            target_label=target_label,
                            loss=loss,
                            use_distributed=self.use_distributed,
                            metric_names=ll_metrics,
                            is_evaluation=True,
                        )

                        ll_stats.update(
                            metric_vals=metrics, batch_time=0.0, n=batch_size
                        )

                        if batch_id % self.log_freq == 0 and self.is_master_node:
                            ll_stats.iter_summary(
                                epoch=epoch,
                                n_processed_samples=processed_samples,
                                total_samples=total_samples,
                                elapsed_time=epoch_start_time,
                                learning_rate=0.0,
                            )

                    avg_loss = ll_stats.avg_statistics(metric_name="loss")
                    loss_surface[coord_a, coord_b] = avg_loss
                    if self.is_master_node:
                        print(
                            "x: {:.2f}, y: {:.2f}, loss: {:.2f}".format(
                                coords_list[0], coords_list[1], avg_loss
                            )
                        )

                    if self.is_master_node:
                        lr_list = [0.0]

                        if self.tb_log_writer is not None:
                            self.log_metrics(
                                lrs=lr_list,
                                log_writer=self.tb_log_writer,
                                train_loss=0.0,
                                val_loss=avg_loss,
                                epoch=epoch,
                                best_metric=self.best_metric,
                                val_ema_loss=None,
                                ckpt_metric_name=None,
                                train_ckpt_metric=None,
                                val_ckpt_metric=None,
                                val_ema_ckpt_metric=None,
                            )
                        if self.bolt_log_writer is not None:
                            self.log_metrics(
                                lrs=lr_list,
                                log_writer=self.bolt_log_writer,
                                train_loss=0.0,
                                val_loss=avg_loss,
                                epoch=epoch,
                                best_metric=self.best_metric,
                                val_ema_loss=None,
                                ckpt_metric_name=None,
                                train_ckpt_metric=None,
                                val_ckpt_metric=None,
                                val_ema_ckpt_metric=None,
                            )

                    gc.collect()
                    # take a small nap
                    time.sleep(1)

            if self.is_master_node:
                ll_utils.plot_save_graphs(
                    save_dir=save_dir,
                    model_name=model_name,
                    grid_a=grid_a,
                    grid_b=grid_b,
                    loss_surface=loss_surface,
                    resolution=n_points,
                )
        except KeyboardInterrupt as e:
            if self.is_master_node:
                logger.log("Keyboard interruption. Exiting from early training")
                raise e
        except Exception as e:
            if "out of memory" in str(e):
                logger.log("OOM exception occured")
                n_gpus = getattr(self.opts, "dev.num_gpus", 1)
                for dev_id in range(n_gpus):
                    mem_summary = torch.cuda.memory_summary(
                        device=torch.device("cuda:{}".format(dev_id)), abbreviated=True
                    )
                    logger.log("Memory summary for device id: {}".format(dev_id))
                    print(mem_summary)
            else:
                logger.log(
                    "Exception occurred that interrupted the training. {}".format(
                        str(e)
                    )
                )
                print(e)
                raise e
        finally:
            if self.use_distributed:
                torch.distributed.destroy_process_group()

            torch.cuda.empty_cache()

            if self.is_master_node:
                ll_end_time = time.time()
                hours, rem = divmod(ll_end_time - ll_start_time, 3600)
                minutes, seconds = divmod(rem, 60)
                train_time_str = "{:0>2}:{:0>2}:{:05.2f}".format(
                    int(hours), int(minutes), seconds
                )
                logger.log("Loss landspace evaluation took {}".format(train_time_str))

            try:
                exit(0)
            except Exception as e:
                pass
            finally:
                pass

max_stars_repo_name: apple/ml-cvnets

max_issues_repo_name: apple/ml-cvnets

max_forks_repo_name: apple/ml-cvnets