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)