File size: 2,062 Bytes
803ef9e |
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 51 52 53 54 55 56 57 58 59 60 61 62 |
import torch.nn as nn
from model import get_model, get_head
from eval.sgd import eval_sgd
from eval.knn import eval_knn
from eval.get_data import get_data
import torch
import tqdm
class BaseMethod(nn.Module):
"""
Base class for self-supervised loss implementation.
It includes encoder and head for training, evaluation function.
"""
def __init__(self, cfg):
super().__init__()
self.model, self.out_size = get_model(cfg.arch, cfg.dataset)
self.head = get_head(self.out_size, cfg)
self.knn = cfg.knn
self.num_pairs = cfg.num_samples * (cfg.num_samples - 1) // 2
self.eval_head = cfg.eval_head
self.emb_size = cfg.emb
def forward(self, samples):
raise NotImplementedError
def get_acc(self, ds_clf, ds_test):
self.eval()
if self.eval_head:
model = lambda x: self.head(self.model(x))
out_size = self.emb_size
else:
model, out_size = self.model, self.out_size
# torch.cuda.empty_cache()
x_train, y_train = get_data(model, ds_clf, out_size, "cuda")
x_test, y_test = get_data(model, ds_test, out_size, "cuda")
acc_knn = eval_knn(x_train, y_train, x_test, y_test, self.knn)
acc_linear = eval_sgd(x_train, y_train, x_test, y_test)
del x_train, y_train, x_test, y_test
self.train()
return acc_knn, acc_linear
def get_acc_knn(self, ds_clf, ds_test):
self.eval()
if self.eval_head:
model = lambda x: self.head(self.model(x))
out_size = self.emb_size
else:
model, out_size = self.model, self.out_size
# torch.cuda.empty_cache()
x_train, y_train = get_data(model, ds_clf, out_size, "cuda")
x_test, y_test = get_data(model, ds_test, out_size, "cuda")
acc_knn = eval_knn(x_train, y_train, x_test, y_test, self.knn)
del x_train, y_train, x_test, y_test
self.train()
return acc_knn
def step(self, progress):
pass
|