| |
| |
| |
|
|
| import os |
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| __all__ = ["AccuracyPredictor"] |
|
|
|
|
| class AccuracyPredictor(nn.Module): |
| def __init__( |
| self, |
| arch_encoder, |
| hidden_size=400, |
| n_layers=3, |
| checkpoint_path=None, |
| device="cuda:0", |
| base_acc_val = None |
| ): |
| super(AccuracyPredictor, self).__init__() |
| self.arch_encoder = arch_encoder |
| self.hidden_size = hidden_size |
| self.n_layers = n_layers |
| self.device = device |
| self.base_acc_val = base_acc_val |
| |
| layers = [] |
| for i in range(self.n_layers): |
| layers.append( |
| nn.Sequential( |
| nn.Linear( |
| self.arch_encoder.n_dim if i == 0 else self.hidden_size, |
| self.hidden_size, |
| ), |
| nn.ReLU(inplace=True), |
| ) |
| ) |
| layers.append(nn.Linear(self.hidden_size, 1, bias=False)) |
| self.layers = nn.Sequential(*layers) |
| if self.base_acc_val!=None : |
| self.base_acc = nn.Parameter( |
| torch.tensor(self.base_acc_val, device=self.device), requires_grad=False |
| ) |
| else: |
| self.base_acc = nn.Parameter( |
| torch.zeros(1, device=self.device), requires_grad=False |
| ) |
|
|
| if checkpoint_path is not None and os.path.exists(checkpoint_path): |
| checkpoint = torch.load(checkpoint_path, map_location="cpu") |
| if "state_dict" in checkpoint: |
| checkpoint = checkpoint["state_dict"] |
| self.load_state_dict(checkpoint) |
| print("Loaded checkpoint from %s" % checkpoint_path) |
|
|
| self.layers = self.layers.to(self.device) |
|
|
| def forward(self, x): |
| y = self.layers(x).squeeze() |
| return y + self.base_acc |
|
|
| def predict_acc(self, arch_dict_list): |
| X = [self.arch_encoder.arch2feature(arch_dict) for arch_dict in arch_dict_list] |
| X = torch.tensor(np.array(X)).float().to(self.device) |
| return self.forward(X) |
|
|
