# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01audio_models.ipynb (unless otherwise specified).

__all__ = ['ResNeSt_', 'ResNeSt50', 'ResNeSt101', 'ResNeSt200', 'ResNeSt269', 'ResNeSt50_fast_1s1x64d',
           'ResNeSt50_fast_1s2x40d', 'ResNeSt50_fast_1s4x24d', 'ResNeSt50_fast_2s1x64d', 'ResNeSt50_fast_2s2x40d',
           'ResNeSt50_fast_4s1x64d', 'ResNeSt50_fast_4s2x40d']

# Cell
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from fastai.vision.all import AdaptiveConcatPool2d, Flatten
from fastai.vision.all import children_and_parameters, create_head, Module
from torchvision import models
from IPython.core.debugger import set_trace
from .core import *
from ..vision.models import *

# Cell
class ResNeSt_(Module):
    arch=None
    def __init__(self, num_classes, head_ps=0.8, pretrained=True, **kwargs):
        try:
            model = torch.hub.load('zhanghang1989/ResNeSt', self.arch, pretrained=pretrained)
        except:
            torch.hub.list('zhanghang1989/ResNeSt', force_reload=True)
            model = torch.hub.load('zhanghang1989/ResNeSt', self.arch, pretrained=pretrained)
        model = nn.Sequential(*[*children_and_parameters(model)[:-2],
                *children_and_parameters(create_head(2048*2, num_classes, ps=head_ps))])
        w = model[0][0].weight.data[:,0].unsqueeze(1)
        model[0][0] = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2),
                                padding=(1, 1), bias=False)
        model[0][0].weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

class ResNeSt50(ResNeSt_):
    arch='resnest50'

class ResNeSt101(ResNeSt_):
    arch='resnest101'

class ResNeSt200(ResNeSt_):
    arch='resnest200'

class ResNeSt269(ResNeSt_):
    arch='resnest269'

class ResNeSt50_fast_1s1x64d(ResNeSt_):
    arch='resnest50_fast_1s1x64d'

class ResNeSt50_fast_1s2x40d(ResNeSt_):
    arch='resnest50_fast_1s2x40d'

class ResNeSt50_fast_1s4x24d(ResNeSt_):
    arch='resnest50_fast_1s4x24d'

class ResNeSt50_fast_2s1x64d(ResNeSt_):
    arch='resnest50_fast_2s1x64d'

class ResNeSt50_fast_2s2x40d(ResNeSt_):
    arch='resnest50_fast_2s2x40d'

class ResNeSt50_fast_4s1x64d(ResNeSt_):
    arch='resnest50_fast_4s1x64d'

class ResNeSt50_fast_4s2x40d(ResNeSt_):
    arch='resnest50_fast_4s2x40d'