Spaces:
Runtime error
Runtime error
| from torch import nn | |
| class BaseNetwork(nn.Module): | |
| def __init__(self): | |
| super(BaseNetwork, self).__init__() | |
| def forward(self, x, y): | |
| pass | |
| def print_network(self): | |
| if isinstance(self, list): | |
| self = self[0] | |
| num_params = 0 | |
| for param in self.parameters(): | |
| num_params += param.numel() | |
| print('Network [%s] was created. Total number of parameters: %.1f million. ' | |
| 'To see the architecture, do print(network).' | |
| % (type(self).__name__, num_params / 1000000)) | |
| def set_requires_grad(self, requires_grad=False): | |
| """Set requies_grad=Fasle for all the networks to avoid unnecessary computations | |
| Parameters: | |
| requires_grad (bool) -- whether the networks require gradients or not | |
| """ | |
| for param in self.parameters(): | |
| param.requires_grad = requires_grad | |
| def init_weights(self, init_type='xavier', gain=0.02): | |
| def init_func(m): | |
| classname = m.__class__.__name__ | |
| if classname.find('BatchNorm2d') != -1: | |
| if hasattr(m, 'weight') and m.weight is not None: | |
| nn.init.normal_(m.weight.data, 1.0, gain) | |
| if hasattr(m, 'bias') and m.bias is not None: | |
| nn.init.constant_(m.bias.data, 0.0) | |
| elif hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): | |
| if init_type == 'normal': | |
| nn.init.normal_(m.weight.data, 0.0, gain) | |
| elif init_type == 'xavier': | |
| nn.init.xavier_normal_(m.weight.data, gain=gain) | |
| elif init_type == 'xavier_uniform': | |
| nn.init.xavier_uniform_(m.weight.data, gain=1.0) | |
| elif init_type == 'kaiming': | |
| nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') | |
| elif init_type == 'orthogonal': | |
| nn.init.orthogonal_(m.weight.data, gain=gain) | |
| elif init_type == 'none': # uses pytorch's default init method | |
| m.reset_parameters() | |
| else: | |
| raise NotImplementedError('initialization method [%s] is not implemented' % init_type) | |
| if hasattr(m, 'bias') and m.bias is not None: | |
| nn.init.constant_(m.bias.data, 0.0) | |
| self.apply(init_func) | |
| # propagate to children | |
| for m in self.children(): | |
| if hasattr(m, 'init_weights'): | |
| m.init_weights(init_type, gain) | |