File size: 2,057 Bytes
6fc43ab |
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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import torch
import torch.nn as nn
import sys
from icecream import ic
# sys.path.append('/home/skowshik/ADRD_repo/adrd_tool/adrd/')
from .net_resnet3d import r3d_18
# from dev.data.dataset_csv import CSVDataset
class ResNetModel(nn.Module):
''' ... '''
def __init__(
self,
tgt_modalities,
mri_feature = 'img_MRI_T1',
):
''' ... '''
super().__init__()
self.mri_feature = mri_feature
self.img_net_ = r3d_18()
# self.modules_emb_src = nn.Sequential(
# nn.BatchNorm1d(9),
# nn.Linear(9, d_model)
# )
# classifiers (binary only)
self.modules_cls = nn.ModuleDict()
for k, info in tgt_modalities.items():
if info['type'] == 'categorical' and info['num_categories'] == 2:
# categorical
self.modules_cls[k] = nn.Linear(64, 1)
else:
# unrecognized
raise ValueError
def forward(self, x):
''' ... '''
tgt_iter = self.modules_cls.keys()
img_x_batch = x[self.mri_feature]
img_out = self.img_net_(img_x_batch)
# ic(img_out.shape)
# run linear classifiers
out = [self.modules_cls[k](img_out).squeeze(1) for i, k in enumerate(tgt_iter)]
out = torch.stack(out, dim=1)
# ic(out.shape)
# out to dict
out = {k: out[:, i] for i, k in enumerate(tgt_iter)}
return out
if __name__ == '__main__':
''' for testing purpose only '''
# import torch
# import numpy as np
# seed = 0
# print('Loading training dataset ... ')
# dat_trn = CSVDataset(mode=0, split=[1, 700], seed=seed)
# print(len(dat_trn))
# tgt_modalities = dat_trn.label_modalities
# net = ResNetModel(tgt_modalities).to('cuda')
# x = dat_trn.features
# x = {k: torch.as_tensor(np.array([x[i][k] for i in range(len(x))])).to('cuda') for k in x[0]}
# ic(x)
# # print(net(x).shape)
# print(net(x))
|