mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-01-25 14:55:10 +00:00
9ee197d0e9
added branch context; added vocab parallel layers; moved split_batch from load_batch to tensor parallel embedding layers; updated gpt model; updated unit test cases; fixed few collective communicator bugs
25 lines
662 B
Python
25 lines
662 B
Python
import torch
|
|
from colossalai.nn.layer.parallel_2d import reduce_by_batch_2d, split_tensor_2d
|
|
from torch import nn
|
|
|
|
from ._utils import calc_acc
|
|
|
|
|
|
class Accuracy2D(nn.Module):
|
|
"""Accuracy for 2D parallelism
|
|
"""
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, logits, targets):
|
|
"""Calculate the accuracy of predicted labels.
|
|
|
|
:param logits: Predicted labels
|
|
:param targets: True labels from data
|
|
"""
|
|
with torch.no_grad():
|
|
targets = split_tensor_2d(targets)
|
|
correct = calc_acc(logits, targets)
|
|
correct = reduce_by_batch_2d(correct)
|
|
return correct
|