File size: 1,438 Bytes
457db56 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import matplotlib.pyplot as plt
def truncated_normal_(tensor, mean=0, std=1):
size = tensor.shape
tmp = tensor.new_empty(size + (4,)).normal_()
valid = (tmp < 2) & (tmp > -2)
ind = valid.max(-1, keepdim=True)[1]
tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
tensor.data.mul_(std).add_(mean)
def init_weights(m):
if type(m) == nn.Conv2d or type(m) == nn.ConvTranspose2d:
nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
#nn.init.normal_(m.weight, std=0.001)
#nn.init.normal_(m.bias, std=0.001)
truncated_normal_(m.bias, mean=0, std=0.001)
def init_weights_orthogonal_normal(m):
if type(m) == nn.Conv2d or type(m) == nn.ConvTranspose2d:
nn.init.orthogonal_(m.weight)
truncated_normal_(m.bias, mean=0, std=0.001)
#nn.init.normal_(m.bias, std=0.001)
def l2_regularisation(m):
l2_reg = None
for W in m.parameters():
if l2_reg is None:
l2_reg = W.norm(2)
else:
l2_reg = l2_reg + W.norm(2)
return l2_reg
def save_mask_prediction_example(mask, pred, iter):
plt.imshow(pred[0,:,:],cmap='Greys')
plt.savefig('images/'+str(iter)+"_prediction.png")
plt.imshow(mask[0,:,:],cmap='Greys')
plt.savefig('images/'+str(iter)+"_mask.png") |