FilesExpand file tree

main

/

resnet_flax.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

111 lines (90 loc) · 3.33 KB

main

/

resnet_flax.py

Copy path

Top

File metadata and controls

• Code
• Blame

111 lines (90 loc) · 3.33 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

import flax.linen as nn

import jax

import jax.numpy as jnp

class FlaxUpsample2D(nn.Module):

out_channels: int

dtype: jnp.dtype = jnp.float32

def setup(self):

self.conv = nn.Conv(

self.out_channels,

kernel_size=(3, 3),

strides=(1, 1),

padding=((1, 1), (1, 1)),

dtype=self.dtype,

)

def __call__(self, hidden_states):

batch, height, width, channels = hidden_states.shape

hidden_states = jax.image.resize(

hidden_states,

shape=(batch, height * 2, width * 2, channels),

method="nearest",

)

hidden_states = self.conv(hidden_states)

return hidden_states

class FlaxDownsample2D(nn.Module):

out_channels: int

dtype: jnp.dtype = jnp.float32

def setup(self):

self.conv = nn.Conv(

self.out_channels,

kernel_size=(3, 3),

strides=(2, 2),

padding=((1, 1), (1, 1)), # padding="VALID",

dtype=self.dtype,

)

def __call__(self, hidden_states):

# pad = ((0, 0), (0, 1), (0, 1), (0, 0)) # pad height and width dim

# hidden_states = jnp.pad(hidden_states, pad_width=pad)

hidden_states = self.conv(hidden_states)

return hidden_states

class FlaxResnetBlock2D(nn.Module):

in_channels: int

out_channels: int = None

dropout_prob: float = 0.0

use_nin_shortcut: bool = None

dtype: jnp.dtype = jnp.float32

def setup(self):

out_channels = self.in_channels if self.out_channels is None else self.out_channels

self.norm1 = nn.GroupNorm(num_groups=32, epsilon=1e-5)

self.conv1 = nn.Conv(

out_channels,

kernel_size=(3, 3),

strides=(1, 1),

padding=((1, 1), (1, 1)),

dtype=self.dtype,

)

self.time_emb_proj = nn.Dense(out_channels, dtype=self.dtype)

self.norm2 = nn.GroupNorm(num_groups=32, epsilon=1e-5)

self.dropout = nn.Dropout(self.dropout_prob)

self.conv2 = nn.Conv(

out_channels,

kernel_size=(3, 3),

strides=(1, 1),

padding=((1, 1), (1, 1)),

dtype=self.dtype,

)

use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

self.conv_shortcut = None

if use_nin_shortcut:

self.conv_shortcut = nn.Conv(

out_channels,

kernel_size=(1, 1),

strides=(1, 1),

padding="VALID",

dtype=self.dtype,

)

def __call__(self, hidden_states, temb, deterministic=True):

residual = hidden_states

hidden_states = self.norm1(hidden_states)

hidden_states = nn.swish(hidden_states)

hidden_states = self.conv1(hidden_states)

temb = self.time_emb_proj(nn.swish(temb))

temb = jnp.expand_dims(jnp.expand_dims(temb, 1), 1)

hidden_states = hidden_states + temb

hidden_states = self.norm2(hidden_states)

hidden_states = nn.swish(hidden_states)

hidden_states = self.dropout(hidden_states, deterministic)

hidden_states = self.conv2(hidden_states)

if self.conv_shortcut is not None:

residual = self.conv_shortcut(residual)

return hidden_states + residual