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