Preprocessing Layer 만들기
Class Normalize(nn.Module): def __init__(self, mean, std): super(Normalize, self).__init__() self.mean = torch.tensor(mean, device='cuda') self.std = torch.tensor(std, device='cuda') def forward(self, input): x = input * 255 x = x - self.mean x = x / self.std return x if onnx: model = nn.Sequential( Normalize([mean_r, mean_g, mean_b], [std_r, std_g, std_b]), model ) return model
2023. 1. 7.