| from model import BrainEncodingModel | |
| from config_utils import load_from_yaml | |
| import torch | |
| cfg = load_from_yaml("./config.yaml") | |
| model = BrainEncodingModel(cfg) | |
| sd_path = './ckpt.pth' | |
| sd = torch.load(sd_path) | |
| model.load_state_dict(sd) | |
| model.eval().cuda() | |
| x = torch.randn(1, 3, 224, 224) | |
| def transform_image(x): | |
| means = [0.485, 0.456, 0.406] | |
| stds = [0.229, 0.224, 0.225] | |
| x = (x - torch.tensor(means).view(1, 3, 1, 1)) / torch.tensor(stds).view(1, 3, 1, 1) | |
| return x | |
| x = transform_image(x) | |
| x = x.cuda() | |
| subject = 'subj01' # could be 1 of 8 subjects | |
| with torch.no_grad(): | |
| out = model(x, subject) | |
| print(out.shape) | |
| # torch.Size([1, 327684]) |