mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-07-07 12:33:37 +08:00
Dr. Kashif Rasul
37 lines
1.0 KiB
Python
37 lines
1.0 KiB
Python
from typing import Optional
|
|
import inspect
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
def get_module_forward_input_names(module: nn.Module):
|
|
params = inspect.signature(module.forward).parameters
|
|
return list(params)
|
|
|
|
|
|
def copy_parameters(net_source: nn.Module, net_dest: nn.Module) -> None:
|
|
net_dest.load_state_dict(net_source.state_dict())
|
|
|
|
|
|
def weighted_average(
|
|
tensor: torch.Tensor, weights: Optional[torch.Tensor] = None, dim=None
|
|
):
|
|
if weights is not None:
|
|
weighted_tensor = tensor * weights
|
|
if dim is not None:
|
|
sum_weights = torch.sum(weights, dim)
|
|
sum_weighted_tensor = torch.sum(weighted_tensor, dim)
|
|
else:
|
|
sum_weights = weights.sum()
|
|
sum_weighted_tensor = weighted_tensor.sum()
|
|
|
|
sum_weights = torch.max(torch.ones_like(sum_weights), sum_weights)
|
|
|
|
return sum_weighted_tensor / sum_weights
|
|
else:
|
|
if dim is not None:
|
|
return torch.mean(tensor, dim=dim)
|
|
else:
|
|
return tensor.mean()
|