def __init__(self,
latent_size: int = 2048,
input_size: int = 140,
output_size: int = 140,
kernel_size: int = 3,
stride: int = 1,
batch_size: int = 256,
weight_KLD: float = 1.0):
super().__init__()
self.latent_size = latent_size
self.encoder = nn.Sequential(
Encoder_4_sampling_bn_1px_deep_convonly_skip(input_size,
kernel_size,
stride,
n_latent=latent_size),
Flatten())
# hidden => mui: 2048 is the number of latent variables fixed in the encoder/decoder models
self.fc1 = nn.Linear(self.latent_size, self.latent_size)
# hidden => logvar is a linear transformation of 2048 units to another 2048 units
self.fc2 = nn.Linear(self.latent_size, self.latent_size)
self.decoder = nn.Sequential(
Unflatten(latent_size, 1, 1),
Decoder_4_sampling_bn_1px_deep_convonly_skip(output_size,
kernel_size,
stride,
n_latent=latent_size))
self.cur_mu = torch.zeros([batch_size, self.latent_size],
dtype=torch.float)
self.cur_logvar = torch.zeros([batch_size, self.latent_size],
dtype=torch.float)
# The weight factor for the KL divergence part of loss. Currently set to 1
self.weight_KLD = nn.Parameter(torch.Tensor([weight_KLD]),
requires_grad=False)
data_loaders = subsetWeightedSampler.get_data_loaders(
dataset,
imbalance_factor=imbalance_factor,
batch_size=batch_size,
num_workers=num_workers)
input_size = 140
output_size = input_size
valid_size = 2
kernel_size = 3
stride = 1
n_fmaps = 16 # fixed in model class
n_latent = 2048
model = AE_Encoder_Classifier(
Encoder_4_sampling_bn_1px_deep_convonly_skip(input_size,
kernel_size,
stride,
n_latent=n_latent),
Classifier3Layered(n_latent=n_latent))
checkpoint = torch.load(state_dict_path,
map_location=lambda storage, loc: storage)
state_dict = checkpoint['model_state_dict']
model.load_encoder_state_dict(state_dict)
model.freeze_encoder_weights(expr=r'^.*\.encoding_conv.*$')
model.reset_state()
for name, param in model.named_parameters():
print(name, param.requires_grad)
criterion = torch.nn.NLLLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.00000075)
def predict_bbox_from_json(bbox_idx, verbose=True):
if verbose:
print('(' + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") +
') Starting Parallel Prediction ... bbox: {}'.format(bbox_idx))
run_root = os.path.dirname(os.path.abspath(__file__))
cache_HDD_root = os.path.join(run_root, '.cache/')
datasources_json_path = os.path.join(run_root,
'datasources_predict_parallel.json')
state_dict_path = os.path.join(
run_root,
'../../training/ae_classify_v09_3layer_unfreeze_latent_debris_clean_transform_add_clean2_wiggle/.log/run_w_pr/epoch_700/model_state_dict'
)
device = 'cpu'
output_wkw_root = '/tmpscratch/webknossos/Connectomics_Department/2018-11-13_scMS109_1to7199_v01_l4_06_24_fixed_mag8_artifact_pred'
output_label = 'probs_sparse'
batch_size = 128
input_shape = (140, 140, 1)
output_shape = (1, 1, 1)
num_workers = 12
kernel_size = 3
stride = 1
n_fmaps = 16
n_latent = 2048
input_size = 140
output_size = input_size
model = AE_Encoder_Classifier(
Encoder_4_sampling_bn_1px_deep_convonly_skip(input_size,
kernel_size,
stride,
n_latent=n_latent),
Classifier3Layered(n_latent=n_latent))
datasources = WkwData.datasources_bbox_from_json(
datasources_json_path,
bbox_ext=[1024, 1024, 1024],
bbox_idx=bbox_idx,
datasource_idx=0)
dataset = WkwData(input_shape=input_shape,
target_shape=output_shape,
data_sources=datasources,
stride=(35, 35, 1),
cache_HDD=False,
cache_RAM=False,
cache_HDD_root=cache_HDD_root)
prediction_loader = torch.utils.data.DataLoader(dataset=dataset,
batch_size=batch_size,
num_workers=num_workers)
checkpoint = torch.load(state_dict_path,
map_location=lambda storage, loc: storage)
state_dict = checkpoint['model_state_dict']
model.load_state_dict(state_dict)
output_prob_fn = lambda x: np.exp(x[:, 1, 0, 0])
# output_dtype = np.uint8
output_dtype = np.float32
# output_dtype_fn = lambda x: (logit(x) + 16) * 256 / 32
output_dtype_fn = lambda x: x
# output_dtype_fni = lambda x: expit(x / 256 * 32 - 16)
output_dtype_fni = lambda x: x
predictor = Predictor(model=model,
dataloader=prediction_loader,
output_prob_fn=output_prob_fn,
output_dtype_fn=output_dtype_fn,
output_dtype=output_dtype,
output_label=output_label,
output_wkw_root=output_wkw_root,
output_wkw_compress=True,
device=device,
interpolate=None)
predictor.predict(verbose=verbose)
