使用来自 __init__ 的参数定义 python 类方法

Question

What I want to achieve

I want to define a class with two modes A and B, so that the forward method of the class changes accordingly.

class MyClass():
    def __init__(self, constant):
        self.constant=constant
            
    def forward(self, x1,x2,function):
        if function=='A':
            return x1+self.constant
        
        elif function=='B':
            return x1*x2+self.constant
        
        else:
            print('please provide the correct function')
model1 = MyClass(2)
model1.forward(2, None, 'A')
output>>>4
model2 = MyClass(2)
model2.forward(2, 2, 'B')
output>>>6

It works, but it is not optimal, since every time when calling the forward method, it will check which function to use. However, the forward function is already set and will never be changed once the class is define, therefore, checking which function to use inside forward is super redundant in my case. (For those who notice this, I am writing my neural network model using PyTorch, two models share 90% of the network architecture, the only 10% differences is the way they do feedforward).

My desired version

I want to set the forward method when the class is defined, so that I can achieve this

model1 = MyClass(2, 'A')
model1.forward(2)
output>>>4

model2 = MyClass(2, 'B')
model2.forward(2, 2)
output>>>6

So I rewrote my class to be:

class MyClass():
    def __init__(self, constant, function):
        self.constant=constant # There would be a lot of shared parameters for the two methods
        self.function=function # This controls the feedforward method of this class
        
    if self.function=='A':
        def forward(self, x1):
            return x1+self.constant
        
    elif self.function=='B':
        def forward(self, x1, x2):
            return x1*x2+self.constant
        
    else:
        print('please provide the correct function')

However, it gives me the following error.

NameError: name 'self' is not defined

How do I write the class for that it defines different forward method based on the args from __init__ ?

Answer 1

You have been trying to redefine the class with your code, such that each new object would change the definition of forward for all objects, before and after. Fortunately, you didn't figure out how to do that.

Instead, make the chosen function an attribute of the object. Code the two functions you want, and then assign the desired variant as you create each instance.

class MyClass():
    def __init__(self, constant, function):
        self.constant=constant
        if function == 'A':
            self.forward = self.forwardA
        elif function=='B':
            self.forward = self.forwardB
        else:
            print('please provide the correct function')
            
    def forwardA(self, x1):
        return x1+self.constant
        
    def forwardB(self, x1, x2):
        return x1*x2+self.constant

# Main
model1 = MyClass(2, 'A')
print(model1.forward(2))

model2 = MyClass(2, 'B')
print(model2.forward(2, 2))

Output:

4
6

Answer 2

You can also try factoring out a base class. It will probably play nicer with mypy and will make it easier to not mix up whatever class you're using.

class MyClassBase():                                      
    def __init__(self, constant):                         
         self.constant=constant                           
                                                          
    def forward(self, *args, **kwargs):              
        raise NotImplementedError('use a derived class')  
                                                          
class MyClassA(MyClassBase):                              
    def __init__(self, constant):                         
        super().__init__(constant)                        
                                                          
    def forward(self,x1):                                 
        return x1 + self.constant                         
                                                          
class MyClassB(MyClassBase):                              
    def __init__(self, constant):                         
        super().__init__(constant)                        
                                                          
    def forward(self, x1, x2):                            
        return x1*x2 + self.constant                      
                                                          
a = MyClassA(2)                                           
b = MyClassB(2)                                           
                                                          
print(a.forward(2))                                       
print(b.forward(2,2))                                     

Answer 3

I would say if you're going to use the constant variable names across returns in a class method, we can define it in this way to tweak the logic. (We can do this in a different way, play only with Params) Hope this looks good.

class MyClass():
    def __init__(self, constant, function):
        self.constant=constant
        self.function=function
    
    def forward(self, x1 = None, x2 = None):
        if self.function=='A':
            return x1+self.constant
        elif self.function=='B':
            return x1*x2+self.constant
        else:
            print('please provide the correct function')

model1 = MyClass(2, 'A')
model1.forward(2)

model2 = MyClass(2, 'B')
model2.forward(2, 2)

Answer 4

Definition

I would also go for inheritance in this case; given what you want:

model2 = MyClass(2, 'B')

It would be much easier (and way more readable for others) to just do:

model2 = MyClassB(2)

Given that, similar to what @Nathan Chappell provided in his answer but shorter (no need for redefining __init__ for example):

import torch


class Base(torch.nn.Module):
    def __init__(self, constant):
        super().__init__()
        self.constant = constant


class MyClassA(Base):
    def forward(self, x1):
        return x1 + self.constant


class MyClassB(Base):
    def forward(self, x1, x2):
        return x1 * x2 + self.constant

Calling

You should use torch.nn.Module 's __call__ method instead of forward as it works correctly with hooks (see this answer ), hence it should be:

model1 = MyClassA(5)
model1(torch.randn(10, 5))

instead of:

model1 = MyClassA(5)
model1.forward(torch.randn(10, 5))