Support "Automatic Mixture Precision ( AMP )"

[ItoMasaki]

What I did

• Add Automatic Mixture Precision ( AMP ) function into Model class

What I did not

• None

What you can do

• Memory on cuda can be reduced if use_amp argument is True in Model.

What you can not do

• None

Operation check

• Check cuda memory it used ( nvidia-smi )

• Sample program

  
  class Normal(dists.Normal):
  def __init__(self):
      super().__init__(var=["y"], cond_var=["x"])
  
      self.nn = nn.Linear(1, 10000)
  
      self.output_loc = nn.Linear(10000, 1)
      self.output_scale = nn.Sequential(
          nn.Linear(10000, 1),
          nn.Softplus(),
      )
  
  def forward(self, x):
      h = self.nn(x)
      loc = self.output_loc(h)
      scale = self.output_scale(h)
  
      return {"loc": loc, "scale": scale}

p = Normal() p.to("cuda") model = Model(loss=-LogProb(p).mean(), distributions=[p], use_amp=False)

x = torch.ones(10000, 1).to("cuda") y = torch.zeros(10000, 1).to("cuda") loss = model.train({"x": x, "y": y})

- Result is like below

GPU RAM ( If use_amp argument is False ) 1305MiB / 15360MiB

GPU RAM ( If use_amp argument is True ) 737MiB / 15360MiB

## Reference information

- [AUTOMATIC MIXED PRECISION PACKAGE - TORCH.AMP](https://pytorch.org/docs/stable/amp.html)