mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-29 04:05:35 +00:00
* Add handson to ColossalAI. * Change names of handsons and edit sequence parallel example. * Edit wrong folder name * resolve conflict * delete readme
13 lines
371 B
Python
13 lines
371 B
Python
import torch
|
|
import math
|
|
|
|
def init_normal(tensor, sigma):
|
|
"""Init method based on N(0, sigma)."""
|
|
torch.nn.init.normal_(tensor, mean=0.0, std=sigma)
|
|
|
|
|
|
def output_init_normal(tensor, sigma, num_layers):
|
|
"""Init method based on N(0, sigma/sqrt(2*num_layers)."""
|
|
std = sigma / math.sqrt(2.0 * num_layers)
|
|
torch.nn.init.normal_(tensor, mean=0.0, std=std)
|