Function bert_example

examples/pytorch/bert/bert_example.py:90–316

(args)

def bert_example(args):
    torch.manual_seed(0)
    random.seed(0)
    np.random.seed(0)

    if dist.is_mpi_available():
        try:
            dist.init_process_group(backend='mpi')
            rank = dist.get_rank()
            world_size = dist.get_world_size()
        except:
            rank = dist.get_rank()
            world_size = dist.get_world_size()
    else:
        rank = 0

    batch_size = args['batch_size']
    seq_len = args['seq_len']
    if args['weight_path'] is not None:
        if 'large' in args['weight_path']:
            layer_num = 24
            head_num = 16
            head_size = 64
        elif 'base' in args['weight_path']:
            layer_num = 12
            head_num = 12
            head_size = 64
        else:
            layer_num = args['layer_num']
            head_num = args['head_num']
            head_size = args['head_size']
    else:
        layer_num = args['layer_num']
        head_num = args['head_num']
        head_size = args['head_size']
    hidden_dim = head_num * head_size

    if args['int8_mode'] == 1:
        per_channel = True
    elif args['int8_mode'] == 2 or args['int8_mode'] == 3:
        per_channel = False
    elif args['int8_mode'] != 0:
        raise ValueError("wrong int8_mode argument")

    if rank == 0:
        print("\n=============== Argument ===============")
        for key in args:
            print("{}: {}".format(key, args[key]))
        print("========================================\n")

    inp = torch.empty(batch_size, seq_len, hidden_dim).cuda()
    torch.nn.init.normal_(inp, -0.02, 0.02)
    if args['avg_seq_len'] > 0:
        mem_seq_lens = torch.ones((batch_size,)) * args['avg_seq_len']
        mem_seq_lens = mem_seq_lens.to(torch.int32).cuda()
    elif args['avg_seq_len'] == -1:
        mem_seq_lens = torch.randint(1, seq_len + 1, (batch_size,), dtype=torch.int32).cuda()
    else: