Spaces:
Runtime error
Runtime error
File size: 1,429 Bytes
4e78e8c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import torch
import os
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torchvision.models as models
os.system("gdown --id 1qO2OLR7skDibo1LaMKD3CiOl_jaCTZ0h")
class JHARMNet(nn.Module):
def __init__(self, pretrained_model, num_classes):
super().__init__()
self.premodel = pretrained_model
self.premodel.fc = nn.Linear(2048, num_classes)
nn.init.xavier_uniform_(self.premodel.fc.weight)
def forward(self, x):
out = self.premodel(x)
return out
class HiddenLayer(nn.Module):
def __init__(self, pretrained_model):
super().__init__()
self.premodel = pretrained_model
self.new_layer = nn.Sequential(
nn.Linear(1000, 512),
nn.LeakyReLU(),
nn.Linear(512, 512),
nn.LeakyReLU(),
nn.Linear(512, 256),
nn.LeakyReLU(),
nn.Linear(256, 10)
)
def forward(self, x):
out = self.premodel(x)
out_new_layer = self.new_layer(out)
return out_new_layer
resnet = models.resnet50(pretrained=True)
x = torch.randn((2, 3, 32, 32))
print(x.shape)
model_check = HiddenLayer(resnet)
model_check.load_state_dict(torch.load("CIFAR_end_hll.pt"))
model_check.eval()
print(model_check.forward(x).shape)
|