Commit 41956de8 authored by Andrei Roibu's avatar Andrei Roibu
Browse files

fixed if-bug in reset parameters

parent 9a385a85
...@@ -209,7 +209,8 @@ class BrainMapperAE3D(nn.Module): ...@@ -209,7 +209,8 @@ class BrainMapperAE3D(nn.Module):
for _, submodule in module.named_children(): for _, submodule in module.named_children():
if isinstance(submodule, (torch.nn.ConvTranspose3d, torch.nn.Conv3d, torch.nn.InstanceNorm3d)) == True: if isinstance(submodule, (torch.nn.ConvTranspose3d, torch.nn.Conv3d, torch.nn.InstanceNorm3d)) == True:
submodule.reset_parameters() submodule.reset_parameters()
if custom_weight_reset_flag == True & isinstance(submodule, (torch.nn.Conv3d, torch.nn.ConvTranspose3d)): if custom_weight_reset_flag == True:
if isinstance(submodule, (torch.nn.Conv3d, torch.nn.ConvTranspose3d)):
gain = np.sqrt(np.divide(2, 1 + np.power(0.25, 2))) gain = np.sqrt(np.divide(2, 1 + np.power(0.25, 2)))
fan, _ = calculate_fan(submodule.weight) fan, _ = calculate_fan(submodule.weight)
std = np.divide(gain, np.sqrt(fan)) std = np.divide(gain, np.sqrt(fan))
...@@ -217,7 +218,8 @@ class BrainMapperAE3D(nn.Module): ...@@ -217,7 +218,8 @@ class BrainMapperAE3D(nn.Module):
for _, subsubmodule in submodule.named_children(): for _, subsubmodule in submodule.named_children():
if isinstance(subsubmodule, (torch.nn.ConvTranspose3d, torch.nn.Conv3d, torch.nn.InstanceNorm3d)) == True: if isinstance(subsubmodule, (torch.nn.ConvTranspose3d, torch.nn.Conv3d, torch.nn.InstanceNorm3d)) == True:
subsubmodule.reset_parameters() subsubmodule.reset_parameters()
if custom_weight_reset_flag == True & isinstance(subsubmodule, (torch.nn.Conv3d, torch.nn.ConvTranspose3d)): if custom_weight_reset_flag == True:
if isinstance(subsubmodule, (torch.nn.Conv3d, torch.nn.ConvTranspose3d)):
gain = np.sqrt(np.divide(2, 1 + np.power(0.25, 2))) gain = np.sqrt(np.divide(2, 1 + np.power(0.25, 2)))
fan, _ = calculate_fan(subsubmodule.weight) fan, _ = calculate_fan(subsubmodule.weight)
std = np.divide(gain, np.sqrt(fan)) std = np.divide(gain, np.sqrt(fan))
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment