DeepMIH/model.py at main · TomTomTommi/DeepMIH

History

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

import warnings

import os

import torch.optim

import torch.nn as nn

import config as c

from hinet import Hinet_stage1

from hinet import Hinet_stage2

class Model_1(nn.Module):

def __init__(self):

super(Model_1, self).__init__()

self.model = Hinet_stage1()

def forward(self, x, rev=False):

if not rev:

out = self.model(x)

else:

out = self.model(x, rev=True)

return out

class Model_2(nn.Module):

def __init__(self):

super(Model_2, self).__init__()

self.model = Hinet_stage2()

def forward(self, x, rev=False):

if not rev:

out = self.model(x)

else:

out = self.model(x, rev=True)

return out

def init_model(mod):

for key, param in mod.named_parameters():

split = key.split('.')

if param.requires_grad:

param.data = c.init_scale * torch.randn(param.data.shape).cuda()

if split[-2] == 'conv5':

param.data.fill_(0.)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

model.py

model.py

Files

model.py

Latest commit

History

model.py

File metadata and controls