diff --git a/code/model_Unet.py b/code/model_Unet.py index b1df0cef0e0a554885414f20382891caa3e483b6..a4185db49ee2c6f520fafd3e3643855f8f7af41e 100644 --- a/code/model_Unet.py +++ b/code/model_Unet.py @@ -31,13 +31,13 @@ class TestModule(pl.LightningModule): # Define the model self.model = smp.Unet( - encoder_name="resnet18", # Also consider using smaller or larger encoders + encoder_name="resnet34", # Also consider using smaller or larger encoders encoder_weights= "imagenet", # Do the pretrained weights help? Try with or without in_channels=1, # We use 1 chanel transmittance as input classes=1, # classes == output channels. We use one output channel for cyto data activation="sigmoid" ) - self.loss_f = RMILoss(with_logits=True) #torch.nn.L1Loss() #RMILoss(with_logits=True) #torch.nn.MSELoss() + self.loss_f = torch.nn.L1Loss() #torch.nn.L1Loss() #RMILoss(with_logits=True) #torch.nn.MSELoss() def forward(self, x): x = self.model(x)