wehub-resource-sync
44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
from .common import preferred_dtype
|
|
|
|
|
|
class MultiOutputModel(torch.nn.Module):
|
|
|
|
def __init__(self, hidden_dim, weight_value):
|
|
super(MultiOutputModel, self).__init__()
|
|
self.linear = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
|
|
self.linear.weight.data.fill_(weight_value)
|
|
self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
|
|
|
|
def forward(self, inputs, targets):
|
|
losses = []
|
|
for x, y in zip(inputs, targets):
|
|
hidden_dim = self.linear(x)
|
|
loss = self.cross_entropy_loss(hidden_dim, y)
|
|
losses.append(loss)
|
|
return tuple(losses)
|
|
|
|
|
|
def multi_output_dataloader(model, total_samples, hidden_dim, device, inputs, targets):
|
|
assert len(inputs) == len(targets)
|
|
batch_size = model.train_micro_batch_size_per_gpu()
|
|
|
|
train_data = [
|
|
torch.full(size=(total_samples, hidden_dim),
|
|
fill_value=x,
|
|
device=device,
|
|
dtype=preferred_dtype(),
|
|
requires_grad=True) for x in inputs
|
|
]
|
|
|
|
train_label = [torch.empty(total_samples, device=device, dtype=torch.long).fill_(y) for y in targets]
|
|
|
|
train_dataset = torch.utils.data.TensorDataset(*train_data, *train_label)
|
|
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
|
|
return train_loader
|