Apply the model and return all of the intermediate tensors. :param x: an [N x C x ...] Tensor of inputs. :param timesteps: a 1-D batch of timesteps. :param y: an [N] Tensor of labels, if class-conditional. :return: a dict with the following keys:
(self, x, timesteps, y=None)
| 491 | return self.out(h) |
| 492 | |
| 493 | def get_feature_vectors(self, x, timesteps, y=None): |
| 494 | """ |
| 495 | Apply the model and return all of the intermediate tensors. |
| 496 | |
| 497 | :param x: an [N x C x ...] Tensor of inputs. |
| 498 | :param timesteps: a 1-D batch of timesteps. |
| 499 | :param y: an [N] Tensor of labels, if class-conditional. |
| 500 | :return: a dict with the following keys: |
| 501 | - 'down': a list of hidden state tensors from downsampling. |
| 502 | - 'middle': the tensor of the output of the lowest-resolution |
| 503 | block in the model. |
| 504 | - 'up': a list of hidden state tensors from upsampling. |
| 505 | """ |
| 506 | hs = [] |
| 507 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) |
| 508 | if self.num_classes is not None: |
| 509 | assert y.shape == (x.shape[0],) |
| 510 | emb = emb + self.label_emb(y) |
| 511 | result = dict(down=[], up=[]) |
| 512 | h = x.type(self.inner_dtype) |
| 513 | for module in self.input_blocks: |
| 514 | h = module(h, emb) |
| 515 | hs.append(h) |
| 516 | result["down"].append(h.type(x.dtype)) |
| 517 | h = self.middle_block(h, emb) |
| 518 | result["middle"] = h.type(x.dtype) |
| 519 | for module in self.output_blocks: |
| 520 | cat_in = th.cat([h, hs.pop()], dim=1) |
| 521 | h = module(cat_in, emb) |
| 522 | result["up"].append(h.type(x.dtype)) |
| 523 | return result |
| 524 | |
| 525 | |
| 526 | class SuperResModel(UNetModel): |
no test coverage detected