os.environ['CUDA_VISIBLE_DEVICES'] = opt['gpu_id']
t0 = time.time()
device0 = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
device1 = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
rootDir = '/data/Jinwei/Bayesian_QSM/' + opt['weight_dir']
# network
unet3d = Unet(
input_channels=1,
output_channels=1,
num_filters=[2**i for i in range(5, 10)], # or range(3, 8)
use_deconv=1,
flag_rsa=0)
unet3d.to(device0)
weights_dict = torch.load(rootDir +
'/linear_factor=1_validation=6_test=7_unet3d.pt')
unet3d.load_state_dict(weights_dict)
resnet = ResBlock(
input_dim=2,
filter_dim=32,
output_dim=1,
)
resnet.to(device1)
weights_dict = torch.load(rootDir +
'/linear_factor=1_validation=6_test=7_resnet.pt')
resnet.load_state_dict(weights_dict)
# optimizer
batch_size = 1
B0_dir = (0, 0, 1)
voxel_size = dataLoader_train.voxel_size
volume_size = dataLoader_train.volume_size
trainLoader = data.DataLoader(dataLoader_train,
batch_size=batch_size,
shuffle=True)
# network of HOBIT
unet3d = Unet(input_channels=1,
output_channels=1,
num_filters=[2**i for i in range(5, 10)],
use_deconv=1,
flag_rsa=0)
unet3d.to(device0)
weights_dict = torch.load(rootDir + '/weight_2nets/unet3d_fine.pt')
# weights_dict = torch.load(rootDir+'/weight_2nets/linear_factor=1_validation=6_test=7_unet3d.pt')
unet3d.load_state_dict(weights_dict)
# QSMnet
unet3d_ = Unet(input_channels=1,
output_channels=1,
num_filters=[2**i for i in range(5, 10)],
use_deconv=1,
flag_rsa=0)
unet3d_.to(device0)
weights_dict = torch.load(rootDir +
'/weight_2nets/rsa=0_validation=6_test=7_.pt'
) # used to compute ich metric
# weights_dict = torch.load(rootDir+'/weight_cv/rsa=0_validation=6_test=7.pt')
initializers={
'w': 'he_normal',
'b': 'normal'
})
criterion = nn.BCELoss(size_average=False)
else:
print("Choose a model.\nAborting....")
sys.exit()
if not os.path.exists('logs'):
os.mkdir('logs')
logFile = 'logs/' + fName + '.txt'
makeLogFile(logFile)
net.to(device)
optimizer = torch.optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-5)
nTrain = len(train_loader)
nValid = len(valid_loader)
nTest = len(test_loader)
minLoss = 1e8
convIter = 0
convCheck = 20
for epoch in range(args.epochs):
trLoss = []
vlLoss = []
vlGed = [0]
klEpoch = [0]
# # network
unet3d = Unet(input_channels=1,
output_channels=2,
num_filters=[2**i for i in range(3, 8)],
bilateral=1,
use_deconv=1,
use_deconv2=1,
renorm=1,
flag_r_train=flag_r_train)
# unet3d = Unet(
# input_channels=1,
# output_channels=2,
# num_filters=[2**i for i in range(3, 8)],
# flag_rsa=2
# )
unet3d.to(device)
if flag_init == 0:
weights_dict = torch.load(
rootDir +
'/weight/weights_sigma={0}_smv={1}_mv8'.format(sigma, 1) + '.pt')
else:
weights_dict = torch.load(
rootDir + folder_weights_VI +
'/weights_lambda_tv={0}.pt'.format(Lambda_tv))
weights_dict['r'] = (torch.ones(1) * r).to(device)
unet3d.load_state_dict(weights_dict)
# optimizer
optimizer = optim.Adam(unet3d.parameters(), lr=lr, betas=(0.5, 0.999))
epoch = 0
dataLoader_test = UTFI_data_loader(
dataFolder='/data/Jinwei/Unsupervised_total_field_inversion/data/test',
flag_train=0)
testLoader = data.DataLoader(dataLoader_test, batch_size=1, shuffle=False)
# network
model = Unet(input_channels=1,
output_channels=2,
num_filters=[2**i for i in range(5, 10)],
bilateral=1,
use_deconv=1,
use_deconv2=1,
renorm=0,
flag_r_train=0,
flag_UTFI=1)
model.to(device)
model.load_state_dict(
torch.load(
rootDir +
'/weights/weight_pdf={0}_tv={1}.pt'.format(lambda_pdf, lambda_tv)))
model.eval()
chi_bs, chi_ls = [], []
with torch.no_grad():
for idx, (ifreqs, masks, data_weights, wGs) in enumerate(testLoader):
ifreqs = ifreqs.to(device, dtype=torch.float)
masks = masks.to(device, dtype=torch.float)
outputs = model(ifreqs)
chi_b, chi_l = outputs[:, 0:1, ...] * (
1 - masks), outputs[:, 1:2, ...] * masks
chi_bs.append(chi_b[0, ...].cpu().detach())
