File size: 2,986 Bytes
2ff0f4b |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import torch
import torch.nn.functional as F
import numpy as np
import torch.nn as nn
import random
def weight_init(module):
for n, m in module.named_children():
print('initialize: '+n)
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.ones_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Sequential):
for f, g in m.named_children():
print('initialize: ' + f)
if isinstance(g, nn.Conv2d):
nn.init.kaiming_normal_(g.weight, mode='fan_in', nonlinearity='relu')
if g.bias is not None:
nn.init.zeros_(g.bias)
elif isinstance(g, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.ones_(g.weight)
if g.bias is not None:
nn.init.zeros_(g.bias)
elif isinstance(g, nn.Linear):
nn.init.kaiming_normal_(g.weight, mode='fan_in', nonlinearity='relu')
if g.bias is not None:
nn.init.zeros_(g.bias)
elif isinstance(m, nn.AdaptiveAvgPool2d) or isinstance(m, nn.AdaptiveMaxPool2d) or isinstance(m, nn.ModuleList) or isinstance(m, nn.BCELoss):
a=1
else:
pass
def init_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)
def BCEDiceLoss(inputs, targets):
# print(inputs.shape, targets.shape)
bce = F.binary_cross_entropy(inputs, targets)
inter = (inputs * targets).sum()
eps = 1e-5
dice = (2 * inter + eps) / (inputs.sum() + targets.sum() + eps)
# print(bce.item(), inter.item(), inputs.sum().item(), dice.item())
return bce + 1 - dice
def BCE(inputs, targets):
# print(inputs.shape, targets.shape)
bce = F.binary_cross_entropy(inputs, targets)
return bce
def adjust_learning_rate(args, optimizer, epoch, iter, max_batches, lr_factor=1):
if args.lr_mode == 'step':
lr = args.lr * (0.1 ** (epoch // args.step_loss))
elif args.lr_mode == 'poly':
cur_iter = iter
max_iter = max_batches * args.max_epochs
lr = args.lr * (1 - cur_iter * 1.0 / max_iter) ** 0.9
else:
raise ValueError('Unknown lr mode {}'.format(args.lr_mode))
if epoch == 0 and iter < 200:
lr = args.lr * 0.9 * (iter + 1) / 200 + 0.1 * args.lr # warm_up
lr *= lr_factor
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
|